update arrow

polars/polars-arrow/Cargo.toml

@@ -9,8 +9,8 @@ description = "Arrow interfaces for Polars DataFrame library"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-# arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "086c54d00aee7b26ba3c8d0c5c683c531e2c75c1", default-features = false }
-arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "from_iter_variants2", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "b46a636f31d70a20bd54df9c7f9e9363053ade08", default-features = false }
+# arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "from_iter_variants", default-features = false }
 # arrow = { package = "arrow2", version = "0.9", default-features = false, features = ["compute_concatenate"] }
 hashbrown = "0.12"
 num = "^0.4"

polars/polars-core/Cargo.toml

@@ -160,10 +160,10 @@ unsafe_unwrap = "^0.1.0"
 
 [dependencies.arrow]
 package = "arrow2"
-# git = "https://github.com/jorgecarleitao/arrow2"
-git = "https://github.com/ritchie46/arrow2"
-# rev = "54797de5e4860cebc4eb73ad1890457cd1a658eb"
-branch = "from_iter_variants2"
+git = "https://github.com/jorgecarleitao/arrow2"
+# git = "https://github.com/ritchie46/arrow2"
+rev = "b46a636f31d70a20bd54df9c7f9e9363053ade08"
+# branch = "from_iter_variants"
 # version = "0.9"
 default-features = false
 features = [

polars/polars-io/Cargo.toml

@@ -31,8 +31,8 @@ private = []
 [dependencies]
 ahash = "0.7"
 anyhow = "1.0"
-# arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "086c54d00aee7b26ba3c8d0c5c683c531e2c75c1", default-features = false }
-arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "from_iter_variants2", default-features = false }
+arrow = { package = "arrow2", git = "https://github.com/jorgecarleitao/arrow2", rev = "b46a636f31d70a20bd54df9c7f9e9363053ade08", default-features = false }
+# arrow = { package = "arrow2", git = "https://github.com/ritchie46/arrow2", branch = "from_iter_variants", default-features = false }
 # arrow = { package = "arrow2", version = "0.9", default-features = false }
 csv-core = { version = "0.1.10", optional = true }
 dirs = "4.0"

polars/polars-io/src/parquet/write.rs

@@ -1,7 +1,7 @@
 use super::ArrowResult;
 use arrow::datatypes::PhysicalType;
 use arrow::error::ArrowError;
-use arrow::io::parquet::write::{self, *};
+use arrow::io::parquet::write::{self, FileWriter, *};
 use arrow::io::parquet::write::{array_to_pages, DynIter, DynStreamingIterator, Encoding};
 use polars_core::prelude::*;
 use rayon::prelude::*;
@@ -103,13 +103,11 @@ where
             })
             .collect::<Vec<_>>();
 
-        // clone is needed because parquet schema is moved into `write_file`
-        let parquet_schema_iter = parquet_schema.clone();
         let row_group_iter = rb_iter.map(|batch| {
             let columns = batch
                 .columns()
                 .par_iter()
-                .zip(parquet_schema_iter.columns().par_iter())
+                .zip(parquet_schema.columns().par_iter())
                 .zip(encodings.par_iter())
                 .map(|((array, descriptor), encoding)| {
                     let encoded_pages =
@@ -127,17 +125,17 @@ where
                     .into_iter()
                     .map(|column| Ok(DynStreamingIterator::new(Bla::new(column)))),
             );
-            ArrowResult::Ok(row_group)
+            ArrowResult::Ok((row_group, batch.columns()[0].len()))
         });
 
-        write::write_file(
-            &mut self.writer,
-            row_group_iter,
-            &schema,
-            parquet_schema,
-            options,
-            None,
-        )?;
+        let mut writer = FileWriter::try_new(&mut self.writer, schema, options)?;
+        // write the headers
+        writer.start()?;
+        for group in row_group_iter {
+            let (group, len) = group?;
+            writer.write(group, len)?;
+        }
+        let _ = writer.end(None)?;
 
         Ok(())
     }