Enable External ArrowColumnWriter Access

parquet/src/arrow/arrow_writer/levels.rs

@@ -49,10 +49,7 @@ use std::ops::Range;
 
 /// Performs a depth-first scan of the children of `array`, constructing [`LevelInfo`]
 /// for each leaf column encountered
-pub(crate) fn calculate_array_levels(
-    array: &ArrayRef,
-    field: &Field,
-) -> Result<Vec<LevelInfo>> {
+pub fn calculate_array_levels(array: &ArrayRef, field: &Field) -> Result<Vec<LevelInfo>> {
     let mut builder = LevelInfoBuilder::try_new(field, Default::default())?;
     builder.write(array, 0..array.len());
     Ok(builder.finish())
@@ -538,7 +535,7 @@ impl LevelInfoBuilder {
 /// The data necessary to write a primitive Arrow array to parquet, taking into account
 /// any non-primitive parents it may have in the arrow representation
 #[derive(Debug, Eq, PartialEq, Clone)]
-pub(crate) struct LevelInfo {
+pub struct LevelInfo {
     /// Array's definition levels
     ///
     /// Present if `max_def_level != 0`

parquet/src/arrow/arrow_writer/mod.rs

@@ -29,7 +29,9 @@ use thrift::protocol::{TCompactOutputProtocol, TSerializable};
 use arrow_array::cast::AsArray;
 use arrow_array::types::*;
 use arrow_array::{Array, FixedSizeListArray, RecordBatch, RecordBatchWriter};
-use arrow_schema::{ArrowError, DataType as ArrowDataType, IntervalUnit, SchemaRef};
+use arrow_schema::{
+    ArrowError, DataType as ArrowDataType, IntervalUnit, Schema, SchemaRef,
+};
 
 use super::schema::{
     add_encoded_arrow_schema_to_metadata, arrow_to_parquet_schema,
@@ -52,7 +54,7 @@ use crate::schema::types::{ColumnDescPtr, SchemaDescriptor};
 use levels::{calculate_array_levels, LevelInfo};
 
 mod byte_array;
-mod levels;
+pub mod levels;
 
 /// Arrow writer
 ///
@@ -299,9 +301,9 @@ impl Read for ArrowColumnChunkReader {
 
 /// A shared [`ArrowColumnChunk`]
 ///
-/// This allows it to be owned by [`ArrowPageWriter`] whilst allowing access via
+/// This allows it to be owned by lower level page writers whilst allowing access via
 /// [`ArrowRowGroupWriter`] on flush, without requiring self-referential borrows
-type SharedColumnChunk = Arc<Mutex<ArrowColumnChunk>>;
+pub type SharedColumnChunk = Arc<Mutex<ArrowColumnChunk>>;
 
 #[derive(Default)]
 struct ArrowPageWriter {
@@ -347,8 +349,8 @@ impl PageWriter for ArrowPageWriter {
     }
 }
 
-/// Encodes a leaf column to [`ArrowPageWriter`]
-enum ArrowColumnWriter {
+/// Can be passed to [write_leaves] function to serialize arrays to [SharedColumnChunk]s
+pub enum ArrowColumnWriter {
     ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
     Column(ColumnWriter<'static>),
 }
@@ -398,6 +400,17 @@ impl ArrowRowGroupWriter {
         Ok(())
     }
 
+    pub fn schema(&self) -> &Arc<Schema> {
+        &self.schema
+    }
+
+    /// Converts this [ArrowRowGroupWriter] into a collection of individual [ArrowColumnWriter]s
+    /// and associated [SharedColumnChunk]s. This permits the caller greater control over how
+    /// data is serialized, such as via parallel threads or async tasks.
+    pub fn into_col_writers(self) -> Vec<(SharedColumnChunk, ArrowColumnWriter)> {
+        self.writers
+    }
+
     pub fn close(self) -> Result<Vec<(ArrowColumnChunk, ColumnCloseResult)>> {
         self.writers
             .into_iter()
@@ -414,6 +427,25 @@ impl ArrowRowGroupWriter {
     }
 }
 
+/// Represents components of a [ArrowRowGroupWriter] which have been broken apart by passing ownership
+/// back out to the caller.
+type DeconstructedRowGroupWriterComponents =
+    (Arc<Schema>, Vec<(SharedColumnChunk, ArrowColumnWriter)>);
+
+impl From<DeconstructedRowGroupWriterComponents> for ArrowRowGroupWriter {
+    fn from(value: DeconstructedRowGroupWriterComponents) -> Self {
+        let schema = value.0;
+        let writers = value.1;
+        Self {
+            writers,
+            schema,
+            // The caller is responsible for tracking buffered_rows when dissasembling and
+            // reasembling ArrowRowGroupWriter
+            buffered_rows: 0,
+        }
+    }
+}
+
 /// Get an [`ArrowColumnWriter`] along with a reference to its [`SharedColumnChunk`]
 fn get_arrow_column_writer(
     data_type: &ArrowDataType,
@@ -479,7 +511,7 @@ fn get_arrow_column_writer(
 }
 
 /// Write the leaves of `array` in depth-first order to `writers` with `levels`
-fn write_leaves<'a, W>(
+pub fn write_leaves<'a, W>(
     writers: &mut W,
     levels: &mut IntoIter<LevelInfo>,
     array: &(dyn Array + 'static),