From 92c244947285d69f2015846c4ab69b9eafed8e06 Mon Sep 17 00:00:00 2001
From: Felix Scherz <felixwscherz@gmail.com>
Date: Wed, 28 Aug 2024 13:34:45 +0200
Subject: [PATCH] feat: support for writing to buffers

---
 fastparquet/test/test_api.py |  9 +++++++++
 fastparquet/writer.py        | 14 +++++++++-----
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/fastparquet/test/test_api.py b/fastparquet/test/test_api.py
index 62cf749c..47235206 100644
--- a/fastparquet/test/test_api.py
+++ b/fastparquet/test/test_api.py
@@ -1546,3 +1546,12 @@ def test_read_a_non_pandas_parquet_file(tempdir):
 
     assert parquet_file.count() == 2
     assert parquet_file.head(1).equals(pd.DataFrame({"foo": [0], "bar": ["a"]}))
+
+
+def test_writing_to_buffer_does_not_close():
+    df = pd.DataFrame({"val": [1, 2]})
+    buffer = io.BytesIO()
+    write(buffer, df, file_scheme="simple")
+    assert not buffer.closed
+    parquet_file = ParquetFile(buffer)
+    assert parquet_file.count() == 2
diff --git a/fastparquet/writer.py b/fastparquet/writer.py
index b023eb2e..218aea56 100644
--- a/fastparquet/writer.py
+++ b/fastparquet/writer.py
@@ -957,11 +957,8 @@ def write_simple(fn, data, fmd, row_group_offsets=None, compression=None,
     if isinstance(data, pd.DataFrame):
         data = iter_dataframe(data, row_group_offsets)
     mode = 'rb+' if append else 'wb'
-    if hasattr(fn, "write"):
-        of = fn
-    else:
-        of = open_with(fn, mode)
-    with of as f:
+
+    def write_to_file(f):
         if append:
             f.seek(-8, 2)
             head_size = struct.unpack('<I', f.read(4))[0]
@@ -981,6 +978,13 @@ def write_simple(fn, data, fmd, row_group_offsets=None, compression=None,
         f.write(struct.pack(b"<I", foot_size))
         f.write(MARKER)
 
+    if hasattr(fn, "write"):
+        write_to_file(fn)
+    else:
+        of = open_with(fn, mode)
+        with of as f:
+            write_to_file(f)
+
 
 def write_multi(dn, data, fmd, row_group_offsets=None, compression=None,
                 file_scheme='hive', write_fmd=True, open_with=default_open,