Allow custom tree builder for parquet::record::RowIter

It will allow to read parquet with custom batch_size. Currently the only possible batch_size for parquet::record::RowIter is 1024
apache · Sep 6, 2023 · b32fd9c · b32fd9c
1 parent 65edbb1
commit b32fd9c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/parquet/src/record/reader.rs b/parquet/src/record/reader.rs
@@ -747,6 +747,12 @@ impl<'a> RowIter<'a> {
         }
     }
 
+    /// Sets tree builder for this row iter.
+    pub fn with_tree_builder(mut self, tree_builder: TreeBuilder) -> Self {
+        self.tree_builder = tree_builder;
+        self
+    }
+
     /// Returns common tree builder, so the same settings are applied to both iterators
     /// from file reader and row group.
     #[inline]