Merge branch 'main' into compat

fastparquet/test/test_api.py

@@ -1555,3 +1555,12 @@ def test_gh929(tempdir):
     df.to_parquet(f"{tempdir}/test_datetimetz_index.parquet", engine="fastparquet")
     result = pd.read_parquet(f"{tempdir}/test_datetimetz_index.parquet", engine="fastparquet")
     assert result.index.equals(df.index)
+
+
+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

@@ -975,11 +975,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]
@@ -999,6 +996,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,