Added support for writing Parquet to StringIO - #51

ankane · Mar 18, 2024 · 91f06be · 91f06be
1 parent c154ebe
commit 91f06be
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 - Updated Polars to 0.38.2
 - Added support for writing JSON to string
+- Added support for writing Parquet to StringIO
 
 ## 0.9.0 (2024-03-03)
 

diff --git a/ext/polars/src/dataframe.rs b/ext/polars/src/dataframe.rs
@@ -603,7 +603,13 @@ impl RbDataFrame {
                 .finish(&mut self.df.borrow_mut())
                 .map_err(RbPolarsErr::from)?;
         } else {
-            todo!();
+            let buf = get_file_like(rb_f, true)?;
+            ParquetWriter::new(buf)
+                .with_compression(compression)
+                .with_statistics(statistics)
+                .with_row_group_size(row_group_size)
+                .finish(&mut self.df.borrow_mut())
+                .map_err(RbPolarsErr::from)?;
         }
 
         Ok(())

diff --git a/test/parquet_test.rb b/test/parquet_test.rb
@@ -43,6 +43,16 @@ def test_write_parquet
     assert_nil df.write_parquet(temp_path)
   end
 
+  def test_write_parquet_io
+    require "stringio"
+
+    df = Polars::DataFrame.new({"a" => [1, 2, 3], "b" => ["one", "two", "three"]})
+    io = StringIO.new
+    df.write_parquet(io)
+    io.rewind
+    assert_frame df, Polars.read_parquet(io)
+  end
+
   def test_write_parquet_struct
     df = Polars::DataFrame.new({"a" => [{"f1" => 1}, {"f1" => 2}]})
     assert_nil df.write_parquet(temp_path)