feat: support for writing to buffers

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

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,