Initial commit

Author: rok

Cargo.lock

@@ -89,20 +89,22 @@ ahash = { version = "0.8", default-features = false, features = [
     "runtime-rng",
 ] }
 apache-avro = { version = "0.17", default-features = false }
-arrow = { version = "55.2.0", features = [
+arrow = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing" , features = [
     "prettyprint",
     "chrono-tz",
 ] }
-arrow-buffer = { version = "55.2.0", default-features = false }
-arrow-flight = { version = "55.2.0", features = [
+
+
+arrow-buffer = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing", default-features = false }
+arrow-flight = {git = "https://github.com/rok/arrow-rs.git", features = [
     "flight-sql-experimental",
 ] }
-arrow-ipc = { version = "55.2.0", default-features = false, features = [
+arrow-ipc = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing" , default-features = false, features = [
     "lz4",
     "zstd",
 ] }
-arrow-ord = { version = "55.2.0", default-features = false }
-arrow-schema = { version = "55.2.0", default-features = false }
+arrow-ord = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing" , default-features = false }
+arrow-schema = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing" , default-features = false }
 async-trait = "0.1.88"
 bigdecimal = "0.4.8"
 bytes = "1.10"
@@ -155,7 +157,7 @@ itertools = "0.14"
 log = "^0.4"
 object_store = { version = "0.12.2", default-features = false }
 parking_lot = "0.12"
-parquet = { version = "55.2.0", default-features = false, features = [
+parquet = { git = "https://github.com/rok/arrow-rs.git", branch = "multi-threaded_encrypted_writing" , default-features = false, features = [
     "arrow",
     "async",
     "object_store",

Cargo.toml

@@ -65,7 +65,7 @@ log = { workspace = true }
 object_store = { workspace = true, optional = true }
 parquet = { workspace = true, optional = true, default-features = true }
 paste = "1.0.15"
-pyo3 = { version = "0.24.2", optional = true }
+pyo3 = { version = "0.25.1", optional = true }
 recursive = { workspace = true, optional = true }
 sqlparser = { workspace = true }
 tokio = { workspace = true }

datafusion/common/Cargo.toml

@@ -25,6 +25,8 @@ use crate::{
     DataFusionError, Result, _internal_datafusion_err,
 };
 
+pub const DEFAULT_MAX_STATISTICS_SIZE: usize = 4096;
+
 use arrow::datatypes::Schema;
 // TODO: handle once deprecated
 #[allow(deprecated)]
@@ -35,7 +37,7 @@ use parquet::{
         metadata::KeyValue,
         properties::{
             EnabledStatistics, WriterProperties, WriterPropertiesBuilder, WriterVersion,
-            DEFAULT_MAX_STATISTICS_SIZE, DEFAULT_STATISTICS_ENABLED,
+            DEFAULT_STATISTICS_ENABLED,
         },
     },
     schema::types::ColumnPath,
@@ -167,13 +169,13 @@ impl TryFrom<&TableParquetOptions> for WriterPropertiesBuilder {
 
             // max_statistics_size is deprecated, currently it is not being used
             // TODO: remove once deprecated
-            #[allow(deprecated)]
-            if let Some(max_statistics_size) = options.max_statistics_size {
-                builder = {
-                    #[allow(deprecated)]
-                    builder.set_column_max_statistics_size(path, max_statistics_size)
-                }
-            }
+            // #[allow(deprecated)]
+            // if let Some(max_statistics_size) = options.max_statistics_size {
+            //     builder = {
+            //         #[allow(deprecated)]
+            //         builder.set_column_max_statistics_size(path, max_statistics_size)
+            //     }
+            // }
         }
 
         Ok(builder)
@@ -222,7 +224,7 @@ impl ParquetOptions {
             dictionary_enabled,
             dictionary_page_size_limit,
             statistics_enabled,
-            max_statistics_size,
+            max_statistics_size: _max_statistics_size,
             max_row_group_size,
             created_by,
             column_index_truncate_length,
@@ -268,12 +270,12 @@ impl ParquetOptions {
             .set_data_page_row_count_limit(*data_page_row_count_limit)
             .set_bloom_filter_enabled(*bloom_filter_on_write);
 
-        builder = {
-            #[allow(deprecated)]
-            builder.set_max_statistics_size(
-                max_statistics_size.unwrap_or(DEFAULT_MAX_STATISTICS_SIZE),
-            )
-        };
+        // builder = {
+        //     #[allow(deprecated)]
+        //     builder.set_max_statistics_size(
+        //         max_statistics_size.unwrap_or(DEFAULT_MAX_STATISTICS_SIZE),
+        //     )
+        // };
 
         if let Some(bloom_filter_fpp) = bloom_filter_fpp {
             builder = builder.set_bloom_filter_fpp(*bloom_filter_fpp);

datafusion/common/src/file_options/parquet_writer.rs

@@ -2134,7 +2134,9 @@ impl ScalarValue {
             | DataType::Time64(TimeUnit::Millisecond)
             | DataType::RunEndEncoded(_, _)
             | DataType::ListView(_)
-            | DataType::LargeListView(_) => {
+            | DataType::LargeListView(_)
+            | DataType::Decimal32(_, _)
+            | DataType::Decimal64(_, _) => {
                 return _not_impl_err!(
                     "Unsupported creation of {:?} array from ScalarValue {:?}",
                     data_type,

datafusion/common/src/scalar/mod.rs

@@ -407,6 +407,7 @@ impl From<DataType> for NativeType {
             DataType::Union(union_fields, _) => {
                 Union(LogicalUnionFields::from(&union_fields))
             }
+            DataType::Decimal32(p, s) | DataType::Decimal64(p, s) |
             DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => Decimal(p, s),
             DataType::Map(field, _) => Map(Arc::new(field.as_ref().into())),
             DataType::Dictionary(_, data_type) => data_type.as_ref().clone().into(),

datafusion/common/src/types/native.rs

@@ -277,6 +277,7 @@ mod tests {
         // Write encrypted parquet using write_parquet
         let mut options = TableParquetOptions::default();
         options.crypto.file_encryption = Some((&encrypt).into());
+        options.global.allow_single_file_parallelism = true;
 
         df.write_parquet(
             tempfile_str.as_str(),

datafusion/core/src/dataframe/parquet.rs

@@ -237,6 +237,8 @@ fn default_field_name(dt: &DataType) -> &str {
         }
         DataType::Decimal128(_, _) => "decimal",
         DataType::Decimal256(_, _) => "decimal",
+        DataType::Decimal32(_, _) => "decimal",
+        DataType::Decimal64(_, _) => "decimal",
     }
 }
 

datafusion/datasource-avro/src/avro_to_arrow/schema.rs

@@ -72,8 +72,8 @@ use object_store::path::Path;
 use object_store::{ObjectMeta, ObjectStore};
 use parquet::arrow::arrow_reader::statistics::StatisticsConverter;
 use parquet::arrow::arrow_writer::{
-    compute_leaves, get_column_writers, ArrowColumnChunk, ArrowColumnWriter,
-    ArrowLeafColumn, ArrowWriterOptions,
+    compute_leaves, ArrowColumnChunk, ArrowColumnWriter, ArrowLeafColumn,
+    ArrowRowGroupWriterFactory, ArrowWriterOptions,
 };
 use parquet::arrow::async_reader::MetadataFetch;
 use parquet::arrow::{parquet_to_arrow_schema, ArrowSchemaConverter, AsyncArrowWriter};
@@ -1306,14 +1306,6 @@ impl FileSink for ParquetSink {
         object_store: Arc<dyn ObjectStore>,
     ) -> Result<u64> {
         let parquet_opts = &self.parquet_options;
-        let mut allow_single_file_parallelism =
-            parquet_opts.global.allow_single_file_parallelism;
-
-        if parquet_opts.crypto.file_encryption.is_some() {
-            // For now, arrow-rs does not support parallel writes with encryption
-            // See https://github.com/apache/arrow-rs/issues/7359
-            allow_single_file_parallelism = false;
-        }
 
         let mut file_write_tasks: JoinSet<
             std::result::Result<(Path, FileMetaData), DataFusionError>,
@@ -1330,7 +1322,7 @@ impl FileSink for ParquetSink {
         };
 
         while let Some((path, mut rx)) = file_stream_rx.recv().await {
-            if !allow_single_file_parallelism {
+            if !parquet_opts.global.allow_single_file_parallelism {
                 let mut writer = self
                     .create_async_arrow_writer(
                         &path,
@@ -1458,13 +1450,13 @@ type ColSender = Sender<ArrowLeafColumn>;
 /// Returns join handles for each columns serialization task along with a send channel
 /// to send arrow arrays to each serialization task.
 fn spawn_column_parallel_row_group_writer(
-    schema: Arc<Schema>,
-    parquet_props: Arc<WriterProperties>,
+    arrow_row_group_writer_factory: Arc<ArrowRowGroupWriterFactory>,
     max_buffer_size: usize,
     pool: &Arc<dyn MemoryPool>,
 ) -> Result<(Vec<ColumnWriterTask>, Vec<ColSender>)> {
-    let schema_desc = ArrowSchemaConverter::new().convert(&schema)?;
-    let col_writers = get_column_writers(&schema_desc, &parquet_props, &schema)?;
+    let arrow_row_group_writer =
+        arrow_row_group_writer_factory.create_row_group_writer(0)?;
+    let col_writers = arrow_row_group_writer.into_column_writers();
     let num_columns = col_writers.len();
 
     let mut col_writer_tasks = Vec::with_capacity(num_columns);
@@ -1559,6 +1551,7 @@ fn spawn_rg_join_and_finalize_task(
 /// across both columns and row_groups, with a theoretical max number of parallel tasks
 /// given by n_columns * num_row_groups.
 fn spawn_parquet_parallel_serialization_task(
+    arrow_row_group_writer_factory: Arc<ArrowRowGroupWriterFactory>,
     mut data: Receiver<RecordBatch>,
     serialize_tx: Sender<SpawnedTask<RBStreamSerializeResult>>,
     schema: Arc<Schema>,
@@ -1571,12 +1564,14 @@ fn spawn_parquet_parallel_serialization_task(
         let max_row_group_rows = writer_props.max_row_group_size();
         let (mut column_writer_handles, mut col_array_channels) =
             spawn_column_parallel_row_group_writer(
-                Arc::clone(&schema),
-                Arc::clone(&writer_props),
+                arrow_row_group_writer_factory.clone(),
                 max_buffer_rb,
                 &pool,
             )?;
         let mut current_rg_rows = 0;
+        // TODO: row_group_writer should use the correct row group index. Currently this would fail if
+        // multiple row groups were written.
+        // let mut rg_index = 0;
 
         while let Some(mut rb) = data.recv().await {
             // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case
@@ -1623,8 +1618,7 @@ fn spawn_parquet_parallel_serialization_task(
 
                     (column_writer_handles, col_array_channels) =
                         spawn_column_parallel_row_group_writer(
-                            Arc::clone(&schema),
-                            Arc::clone(&writer_props),
+                            arrow_row_group_writer_factory.clone(),
                             max_buffer_rb,
                             &pool,
                         )?;
@@ -1655,24 +1649,15 @@ fn spawn_parquet_parallel_serialization_task(
 /// Consume RowGroups serialized by other parallel tasks and concatenate them in
 /// to the final parquet file, while flushing finalized bytes to an [ObjectStore]
 async fn concatenate_parallel_row_groups(
+    mut parquet_writer: SerializedFileWriter<SharedBuffer>,
+    merged_buff: SharedBuffer,
     mut serialize_rx: Receiver<SpawnedTask<RBStreamSerializeResult>>,
-    schema: Arc<Schema>,
-    writer_props: Arc<WriterProperties>,
     mut object_store_writer: Box<dyn AsyncWrite + Send + Unpin>,
     pool: Arc<dyn MemoryPool>,
 ) -> Result<FileMetaData> {
-    let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES);
-
     let mut file_reservation =
         MemoryConsumer::new("ParquetSink(SerializedFileWriter)").register(&pool);
 
-    let schema_desc = ArrowSchemaConverter::new().convert(schema.as_ref())?;
-    let mut parquet_writer = SerializedFileWriter::new(
-        merged_buff.clone(),
-        schema_desc.root_schema_ptr(),
-        writer_props,
-    )?;
-
     while let Some(task) = serialize_rx.recv().await {
         let result = task.join_unwind().await;
         let mut rg_out = parquet_writer.next_row_group()?;
@@ -1723,28 +1708,47 @@ async fn output_single_parquet_file_parallelized(
     let (serialize_tx, serialize_rx) =
         mpsc::channel::<SpawnedTask<RBStreamSerializeResult>>(max_rowgroups);
 
+    let parquet_schema = ArrowSchemaConverter::new()
+        .with_coerce_types(parquet_props.coerce_types())
+        .convert(&output_schema)?;
+    let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES);
+    let parquet_writer = SerializedFileWriter::new(
+        merged_buff.clone(),
+        parquet_schema.root_schema_ptr(),
+        parquet_props.clone().into(),
+    )?;
+    let arrow_row_group_writer_factory = ArrowRowGroupWriterFactory::new(
+        &parquet_writer,
+        parquet_schema,
+        output_schema.clone(),
+        parquet_props.clone().into(),
+    );
+
     let arc_props = Arc::new(parquet_props.clone());
     let launch_serialization_task = spawn_parquet_parallel_serialization_task(
+        Arc::new(arrow_row_group_writer_factory),
         data,
         serialize_tx,
         Arc::clone(&output_schema),
         Arc::clone(&arc_props),
         parallel_options,
         Arc::clone(&pool),
     );
+
+    launch_serialization_task
+        .join_unwind()
+        .await
+        .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
+
     let file_metadata = concatenate_parallel_row_groups(
+        parquet_writer,
+        merged_buff,
         serialize_rx,
-        Arc::clone(&output_schema),
-        Arc::clone(&arc_props),
         object_store_writer,
         pool,
     )
     .await?;
 
-    launch_serialization_task
-        .join_unwind()
-        .await
-        .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
     Ok(file_metadata)
 }
 

datafusion/datasource-parquet/src/file_format.rs

@@ -816,6 +816,8 @@ pub fn can_hash(data_type: &DataType) -> bool {
         DataType::Float64 => true,
         DataType::Decimal128(_, _) => true,
         DataType::Decimal256(_, _) => true,
+        DataType::Decimal32(_, _) => true,
+        DataType::Decimal64(_, _) => true,
         DataType::Timestamp(_, _) => true,
         DataType::Utf8 => true,
         DataType::LargeUtf8 => true,