diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 347721c43d7c..023261ff4404 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1451,14 +1451,15 @@ config_namespace! { /// either limit is reached, whichever comes first. pub max_row_group_size: usize, default = 1024 * 1024 - /// (writing) Target maximum size of each row group in bytes. When set, - /// the writer flushes whenever either this limit or `max_row_group_size` - /// is reached, whichever comes first. Useful for bounding writer memory - /// on wide schemas where a row-count limit can map to very different - /// byte sizes. Matches the behavior of `parquet.block.size` in - /// parquet-mr. If `None` (the default), only the row-count limit - /// applies. Currently only honored when `allow_single_file_parallelism` - /// is `false`; by default the parallel file writer ignores this limit. + /// (writing) Target maximum estimated encoded size of each row group in bytes. + /// When set, either this target or `max_row_group_size` triggers a flush. The + /// first batch, subject to the row limit, is written before its size can be + /// estimated. Subsequent batches are split using the observed average row size, + /// so this is not a hard byte or memory limit. The parallel writer synchronizes + /// column feedback to match the single-threaded writer's boundaries for the + /// same batches. Columns within each slice still encode in parallel, but + /// synchronization reduces overlap between slices and may lower write + /// throughput. If `None` (the default), only the row-count limit applies. pub max_row_group_bytes: Option, default = None /// (writing) Sets "created by" property diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a2589af19a6e..df30b968b102 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -63,6 +63,7 @@ criterion = { workspace = true } datafusion-functions = { workspace = true } datafusion-functions-nested = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true, features = ["rt-multi-thread"] } # Note: add additional linter rules in lib.rs. # Rust does not support workspace + new linter rules in subcrates yet diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 1b79ae665bb1..b6340804b2ff 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -66,6 +66,7 @@ use parquet::file::properties::{ use parquet::file::writer::SerializedFileWriter; use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio::sync::mpsc::{self, Receiver, Sender}; +use tokio::sync::watch; /// Initial writing buffer size. Note this is just a size hint for efficiency. It /// will grow beyond the set value if needed. @@ -515,57 +516,48 @@ impl ParquetSink { } } -/// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] -/// Once the channel is exhausted, returns the ArrowColumnWriter. +/// A leaf array and its number of root records. Nested leaf arrays can have a +/// different length, so the dispatcher supplies the record count explicitly. +struct ColumnInput { + column: ArrowLeafColumn, + rows: usize, +} + +/// A coherent snapshot published only after encoding and reserving memory. +#[derive(Clone, Copy, Debug, Default)] +struct ColumnProgress { + rows: usize, + bytes: usize, +} + +/// Encodes one leaf column. Dropping the input channel finishes the row group. async fn column_serializer_task( - mut rx: Receiver, + mut rx: Receiver, mut writer: ArrowColumnWriter, reservation: MemoryReservation, encoding_time: Time, + progress: Option>, ) -> Result<(ArrowColumnWriter, MemoryReservation)> { - while let Some(col) = rx.recv().await { + let mut rows = 0; + while let Some(input) = rx.recv().await { let _timer = encoding_time.timer(); - writer.write(&col)?; + writer.write(&input.column)?; reservation.try_resize(writer.memory_size())?; + if let Some(progress) = &progress { + rows += input.rows; + // Observation ends when the dispatcher hands this group to the + // finalizer. The worker result still carries any encoding error. + let _ = progress.send(ColumnProgress { + rows, + bytes: writer.get_estimated_total_bytes(), + }); + } } Ok((writer, reservation)) } type ColumnWriterTask = SpawnedTask>; -type ColSender = Sender; - -/// Spawns a parallel serialization task for each column -/// 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( - col_writers: Vec, - max_buffer_size: usize, - pool: &Arc, - encoding_time: &Time, -) -> Result<(Vec, Vec)> { - let num_columns = col_writers.len(); - - let mut col_writer_tasks = Vec::with_capacity(num_columns); - let mut col_array_channels = Vec::with_capacity(num_columns); - for writer in col_writers.into_iter() { - // Buffer size of this channel limits the number of arrays queued up for column level serialization - let (send_array, receive_array) = - mpsc::channel::(max_buffer_size); - col_array_channels.push(send_array); - - let reservation = - MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(pool); - let task = SpawnedTask::spawn(column_serializer_task( - receive_array, - writer, - reservation, - encoding_time.clone(), - )); - col_writer_tasks.push(task); - } - - Ok((col_writer_tasks, col_array_channels)) -} +type ColSender = Sender; /// Settings related to writing parquet files in parallel #[derive(Clone)] @@ -594,28 +586,27 @@ struct ParquetFileWriteContext { /// i.e. the Vec of encoded columns which can be appended to a row group type RBStreamSerializeResult = Result<(Vec, MemoryReservation, usize)>; -/// Sends the ArrowArrays in passed [RecordBatch] through the channels to their respective -/// parallel column serializers. +/// Sends a batch to every leaf writer, returning the index of a closed input +/// channel so its owner can join the worker and propagate the original error. async fn send_arrays_to_col_writers( col_array_channels: &[ColSender], rb: &RecordBatch, - schema: Arc, -) -> Result<()> { - // Each leaf column has its own channel, increment next_channel for each leaf column sent. + schema: &Schema, +) -> Result> { let mut next_channel = 0; for (array, field) in rb.columns().iter().zip(schema.fields()) { - for c in compute_leaves(field, array)? { - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if col_array_channels[next_channel].send(c).await.is_err() { - return Ok(()); + for column in compute_leaves(field, array)? { + let input = ColumnInput { + column, + rows: rb.num_rows(), + }; + if col_array_channels[next_channel].send(input).await.is_err() { + return Ok(Some(next_channel)); } - next_channel += 1; } } - - Ok(()) + Ok(None) } /// Spawns a tokio task which joins the parallel column writer tasks, @@ -647,14 +638,154 @@ fn spawn_rg_join_and_finalize_task( }) } -/// This task coordinates the serialization of a parquet file in parallel. -/// As the query produces RecordBatches, these are written to a RowGroup -/// via parallel [ArrowColumnWriter] tasks. Once the desired max rows per -/// row group is reached, the parallel tasks are joined on another separate task -/// and sent to a concatenation task. This task immediately continues to work -/// on the next row group in parallel. So, parquet serialization is parallelized -/// across both columns and row_groups, with a theoretical max number of parallel tasks -/// given by n_columns * num_row_groups. +/// Owns all tasks and feedback for a single row group. Worker results must +/// either be joined here on failure or transferred to the finalizer. +struct InProgressRowGroup { + column_writer_handles: Vec, + col_array_channels: Vec, + progress: Vec>, + rows: usize, +} + +impl InProgressRowGroup { + fn new( + factory: &ArrowRowGroupWriterFactory, + index: usize, + ctx: &ParquetFileWriteContext, + encoding_time: &Time, + ) -> Result { + let writers = factory.create_column_writers(index)?; + let track_bytes = ctx.props.max_row_group_bytes().is_some(); + let mut group = Self { + column_writer_handles: Vec::with_capacity(writers.len()), + col_array_channels: Vec::with_capacity(writers.len()), + progress: Vec::with_capacity(if track_bytes { writers.len() } else { 0 }), + rows: 0, + }; + for writer in writers { + let (tx, rx) = mpsc::channel( + ctx.parallel_options.max_buffered_record_batches_per_stream, + ); + let progress = if track_bytes { + let (tx, rx) = watch::channel(ColumnProgress::default()); + group.progress.push(rx); + Some(tx) + } else { + None + }; + let reservation = + MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(&ctx.pool); + group + .column_writer_handles + .push(SpawnedTask::spawn(column_serializer_task( + rx, + writer, + reservation, + encoding_time.clone(), + progress, + ))); + group.col_array_channels.push(tx); + } + Ok(group) + } + + async fn write(&mut self, batch: &RecordBatch, schema: &Schema) -> Result<()> { + if let Some(index) = + send_arrays_to_col_writers(&self.col_array_channels, batch, schema).await? + { + return self.column_error(index).await; + } + self.rows += batch.num_rows(); + Ok(()) + } + + /// A closed worker channel is not evidence of successful shutdown. Join its + /// task before dropping the remaining workers so the original error survives. + async fn column_error(&mut self, index: usize) -> Result { + self.col_array_channels.clear(); + let _ = self + .column_writer_handles + .swap_remove(index) + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + Err(internal_datafusion_err!( + "Parquet column writer exited before completing its input" + )) + } + + async fn synchronize(&mut self) -> Result<()> { + for index in 0..self.progress.len() { + // Drop the watch borrow before any other await. + let closed = self.progress[index] + .wait_for(|p| p.rows >= self.rows) + .await + .is_err(); + if closed { + return self.column_error(index).await; + } + } + Ok(()) + } + + /// All columns have acknowledged the current prefix in `after_write`. + fn estimated_bytes(&self) -> usize { + self.progress + .iter() + .fold(0usize, |total, rx| total.saturating_add(rx.borrow().bytes)) + } + + async fn after_write(&mut self) -> Result<()> { + self.synchronize().await + } + + fn finish( + self, + pool: &Arc, + encoding_time: &Time, + ) -> SpawnedTask { + let Self { + column_writer_handles, + col_array_channels, + rows, + .. + } = self; + drop(col_array_channels); + spawn_rg_join_and_finalize_task( + column_writer_handles, + rows, + pool, + encoding_time.clone(), + ) + } +} + +async fn finish_and_restart_row_group( + group: &mut InProgressRowGroup, + index: &mut usize, + factory: &ArrowRowGroupWriterFactory, + ctx: &ParquetFileWriteContext, + encoding_time: &Time, + serialize_tx: &Sender>, +) -> Result { + group.col_array_channels.clear(); + let task = spawn_rg_join_and_finalize_task( + std::mem::take(&mut group.column_writer_handles), + group.rows, + &ctx.pool, + encoding_time.clone(), + ); + // The consumer owns the error when its output channel has closed. + if serialize_tx.send(task).await.is_err() { + return Ok(false); + } + *index += 1; + *group = InProgressRowGroup::new(factory, *index, ctx, encoding_time)?; + Ok(true) +} + +/// Selects common root-record boundaries for the parallel leaf writers. +/// Byte decisions use the same acknowledged estimates as ArrowWriter. fn spawn_parquet_parallel_serialization_task( row_group_writer_factory: ArrowRowGroupWriterFactory, mut data: Receiver, @@ -663,98 +794,102 @@ fn spawn_parquet_parallel_serialization_task( encoding_time: Time, ) -> SpawnedTask> { SpawnedTask::spawn(async move { - let max_buffer_rb = ctx.parallel_options.max_buffered_record_batches_per_stream; - let max_row_group_rows = ctx + let max_rows = ctx .props .max_row_group_row_count() .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT); - let mut row_group_index = 0; - let col_writers = - row_group_writer_factory.create_column_writers(row_group_index)?; - let (mut column_writer_handles, mut col_array_channels) = - spawn_column_parallel_row_group_writer( - col_writers, - max_buffer_rb, - &ctx.pool, - &encoding_time, - )?; - let mut current_rg_rows = 0; - - while let Some(mut rb) = data.recv().await { - // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case - // when max_row_group_rows < execution.batch_size as an alternative to a recursive async - // function. + let max_bytes = ctx.props.max_row_group_bytes(); + let mut index = 0; + let mut group = InProgressRowGroup::new( + &row_group_writer_factory, + index, + &ctx, + &encoding_time, + )?; + while let Some(mut batch) = data.recv().await { + if batch.num_rows() == 0 { + continue; + } + // Preserve ArrowWriter's row-first, then byte-based recursive + // slicing order without recursive async calls. No allocation or + // batch slice is needed when the whole batch fits. + let mut pending = Vec::new(); loop { - if current_rg_rows + rb.num_rows() < max_row_group_rows { - send_arrays_to_col_writers( - &col_array_channels, - &rb, - Arc::clone(&ctx.schema), - ) - .await?; - current_rg_rows += rb.num_rows(); - break; - } else { - let rows_left = max_row_group_rows - current_rg_rows; - let a = rb.slice(0, rows_left); - send_arrays_to_col_writers( - &col_array_channels, - &a, - Arc::clone(&ctx.schema), - ) - .await?; + let row_budget = max_rows - group.rows; + if batch.num_rows() > row_budget { + pending.push(batch.slice(row_budget, batch.num_rows() - row_budget)); + batch = batch.slice(0, row_budget); + } - // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup - // on a separate task, so that we can immediately start on the next RG before waiting - // for the current one to finish. - drop(col_array_channels); - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - max_row_group_rows, - &ctx.pool, - encoding_time.clone(), - ); - - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { - return Ok(()); + if let Some(limit) = max_bytes.filter(|_| group.rows > 0) { + let bytes = group.estimated_bytes(); + let avg = bytes / group.rows; + // A zero integer average disables prediction, just as in + // ArrowWriter. The accumulated byte check still applies. + let byte_budget = if bytes >= limit { + 0 + } else { + (limit - bytes).checked_div(avg).unwrap_or(usize::MAX) + }; + if byte_budget == 0 { + if !finish_and_restart_row_group( + &mut group, + &mut index, + &row_group_writer_factory, + &ctx, + &encoding_time, + &serialize_tx, + ) + .await? + { + return Ok(()); + } + continue; + } + if batch.num_rows() > byte_budget { + pending.push( + batch.slice(byte_budget, batch.num_rows() - byte_budget), + ); + batch = batch.slice(0, byte_budget); } + } - current_rg_rows = 0; - rb = rb.slice(rows_left, rb.num_rows() - rows_left); - - row_group_index += 1; - let col_writers = row_group_writer_factory - .create_column_writers(row_group_index)?; - (column_writer_handles, col_array_channels) = - spawn_column_parallel_row_group_writer( - col_writers, - max_buffer_rb, - &ctx.pool, - &encoding_time, - )?; + group.write(&batch, &ctx.schema).await?; + let full = if group.rows == max_rows { + // This boundary is already fixed. Finalization and the + // next group's encoding can overlap without a feedback wait. + true + } else if let Some(limit) = max_bytes { + group.after_write().await?; + group.estimated_bytes() >= limit + } else { + false + }; + if full + && !finish_and_restart_row_group( + &mut group, + &mut index, + &row_group_writer_factory, + &ctx, + &encoding_time, + &serialize_tx, + ) + .await? + { + return Ok(()); + } + match pending.pop() { + Some(next) => batch = next, + None => break, } } } - - drop(col_array_channels); - // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows - if current_rg_rows > 0 { - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - current_rg_rows, - &ctx.pool, - encoding_time.clone(), - ); - - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { + if group.rows > 0 { + let task = group.finish(&ctx.pool, &encoding_time); + if serialize_tx.send(task).await.is_err() { return Ok(()); } } - Ok(()) }) } @@ -854,3 +989,6 @@ async fn output_single_parquet_file_parallelized( .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; Ok(parquet_meta_data) } + +#[cfg(test)] +mod tests; diff --git a/datafusion/datasource-parquet/src/sink/tests.rs b/datafusion/datasource-parquet/src/sink/tests.rs new file mode 100644 index 000000000000..63bfe0822ecc --- /dev/null +++ b/datafusion/datasource-parquet/src/sink/tests.rs @@ -0,0 +1,467 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::*; +use arrow::array::{ + ArrayRef, BooleanArray, Int64Array, Int64Builder, ListBuilder, StringArray, +}; +use arrow::compute::concat_batches; +use arrow::datatypes::{DataType, Field}; +use bytes::Bytes; +use datafusion_execution::memory_pool::UnboundedMemoryPool; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::Compression; +use std::io::Write; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::Duration; + +struct TestOutput(SharedBuffer); + +impl AsyncWrite for TestOutput { + fn poll_write( + mut self: Pin<&mut Self>, + _: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + Poll::Ready(Write::write(&mut self.0, bytes)) + } + + fn poll_flush( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } +} + +fn properties(rows: usize, bytes: Option) -> WriterProperties { + WriterProperties::builder() + .set_max_row_group_row_count(Some(rows)) + .set_max_row_group_bytes(bytes) + .set_dictionary_enabled(false) + .set_compression(Compression::UNCOMPRESSED) + .build() +} + +fn context( + schema: SchemaRef, + props: WriterProperties, + pool: Arc, + capacity: usize, +) -> ParquetFileWriteContext { + ParquetFileWriteContext { + schema, + props: Arc::new(props), + skip_arrow_metadata: false, + parallel_options: Arc::new(ParallelParquetWriterOptions { + max_parallel_row_groups: 2, + max_buffered_record_batches_per_stream: capacity, + }), + pool, + } +} + +async fn parallel( + batches: &[RecordBatch], + props: WriterProperties, + pool: Arc, + capacity: usize, +) -> Result<(ParquetMetaData, Vec)> { + let ctx = context(batches[0].schema(), props, pool, capacity); + let output = SharedBuffer::new(0); + let (tx, rx) = mpsc::channel(2); + let batches = batches.to_vec(); + let feeder = SpawnedTask::spawn(async move { + for batch in batches { + if tx.send(batch).await.is_err() { + break; + } + } + }); + let result = output_single_parquet_file_parallelized( + Box::new(TestOutput(output.clone())), + rx, + ctx, + Time::new(), + ) + .await; + feeder.join_unwind().await.unwrap(); + let bytes = output.buffer.lock().await.clone(); + Ok((result?, bytes)) +} + +fn serial(batches: &[RecordBatch], props: WriterProperties) -> ParquetMetaData { + let mut writer = + ArrowWriter::try_new(Vec::new(), batches[0].schema(), Some(props)).unwrap(); + for batch in batches { + writer.write(batch).unwrap(); + } + writer.close().unwrap() +} + +fn row_counts(metadata: &ParquetMetaData) -> Vec { + metadata + .row_groups() + .iter() + .map(|group| group.num_rows()) + .collect() +} + +fn integers(rows: usize, start: i64) -> RecordBatch { + RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(start..start + rows as i64)) as ArrayRef, + )]) + .unwrap() +} + +fn strings(rows: usize, width: usize) -> RecordBatch { + RecordBatch::try_from_iter([( + "s", + Arc::new(StringArray::from_iter_values( + (0..rows).map(|i| format!("{i:0width$}")), + )) as ArrayRef, + )]) + .unwrap() +} + +async fn check_round_trip( + batches: &[RecordBatch], + props: WriterProperties, + capacity: usize, +) -> ParquetMetaData { + let pool = Arc::new(UnboundedMemoryPool::default()); + let (metadata, bytes) = tokio::time::timeout( + Duration::from_secs(10), + parallel(batches, props.clone(), pool.clone(), capacity), + ) + .await + .expect("parallel writer stalled") + .unwrap(); + assert_eq!( + metadata.file_metadata().num_rows() as usize, + batches.iter().map(RecordBatch::num_rows).sum::() + ); + assert!( + metadata + .row_groups() + .iter() + .all(|g| g.num_rows() as usize <= props.max_row_group_row_count().unwrap()) + ); + let read = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes)) + .unwrap() + .with_batch_size(137) + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + let schema = batches[0].schema(); + assert_eq!( + concat_batches(&schema, batches).unwrap(), + concat_batches(&schema, &read).unwrap() + ); + assert_eq!(pool.reserved(), 0); + metadata +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn row_and_byte_limits_round_trip() { + let batches = [integers(900, 0), integers(1024, 900)]; + for capacity in [1, 2, 8] { + for bytes in [None, Some(1), Some(7600), Some(1_000_000)] { + let props = properties(1000, bytes); + let actual = check_round_trip(&batches, props.clone(), capacity).await; + assert_eq!(row_counts(&actual), row_counts(&serial(&batches, props))); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn prediction_reconsiders_a_cheaper_suffix() { + let batches = [strings(100, 1000), strings(1000, 8)]; + let props = properties(20_000, Some(131_072)); + let expected = serial(&batches, props.clone()); + assert_eq!(row_counts(&expected), vec![1100]); + let actual = check_round_trip(&batches, props, 2).await; + assert_eq!(row_counts(&actual), row_counts(&expected)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn fractional_bytes_per_row_and_wide_booleans() { + for columns in [1, 100] { + let schema = Arc::new(Schema::new( + (0..columns) + .map(|i| Field::new(format!("b{i}"), DataType::Boolean, false)) + .collect::>(), + )); + let values = + Arc::new(BooleanArray::from_iter((0..1024).map(|i| i % 2 == 0))) as ArrayRef; + let batch = RecordBatch::try_new(schema, vec![values; columns]).unwrap(); + let batches = vec![batch; 4]; + let props = properties(100_000, Some(if columns == 1 { 256 } else { 4096 })); + let actual = check_round_trip(&batches, props.clone(), 2).await; + assert_eq!(row_counts(&actual), row_counts(&serial(&batches, props))); + if columns == 100 { + assert_eq!(row_counts(&actual), vec![1024; 4]); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn nested_columns_use_root_row_counts() { + let mut list = ListBuilder::new(Int64Builder::new()); + for row in 0..472 { + for value in 0..row % 13 { + list.values().append_value(value); + } + list.append(row % 7 != 0); + } + let batch = RecordBatch::try_from_iter([ + ("list", Arc::new(list.finish()) as ArrayRef), + ( + "id", + Arc::new(Int64Array::from_iter_values(0..472)) as ArrayRef, + ), + ]) + .unwrap(); + let batches = [ + batch.slice(0, 17), + batch.slice(17, 103), + batch.slice(120, 41), + batch.slice(161, 311), + ]; + let props = properties(200, Some(2048)); + let actual = check_round_trip(&batches, props.clone(), 1).await; + assert_eq!(row_counts(&actual), row_counts(&serial(&batches, props))); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn empty_batches_and_exact_final_boundary() { + let data = integers(64, 0); + let empty = RecordBatch::new_empty(data.schema()); + for batches in [vec![empty.clone()], vec![empty.clone(), data, empty]] { + for bytes in [None, Some(1), Some(1024)] { + let props = properties(64, bytes); + let actual = check_round_trip(&batches, props.clone(), 2).await; + assert_eq!(row_counts(&actual), row_counts(&serial(&batches, props))); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn dictionary_fallback_and_compression() { + let batches = [strings(100, 1000), strings(1000, 8), strings(1024, 200)]; + for compression in [ + Compression::UNCOMPRESSED, + Compression::SNAPPY, + Compression::ZSTD(Default::default()), + ] { + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(2000)) + .set_max_row_group_bytes(Some(131_072)) + .set_dictionary_enabled(true) + .set_dictionary_page_size_limit(1024) + .set_data_page_size_limit(4096) + .set_compression(compression) + .build(); + let actual = check_round_trip(&batches, props.clone(), 2).await; + assert_eq!(row_counts(&actual), row_counts(&serial(&batches, props))); + } +} + +#[derive(Debug)] +struct FailingPool { + inner: UnboundedMemoryPool, + calls: AtomicUsize, + fail_after: usize, +} + +impl fmt::Display for FailingPool { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "FailingPool") + } +} + +impl MemoryPool for FailingPool { + fn name(&self) -> &str { + "FailingPool" + } + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + } + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + } + fn reserved(&self) -> usize { + self.inner.reserved() + } + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + if reservation.consumer().name() == "ParquetSink(ArrowColumnWriter)" + && self.calls.fetch_add(1, Ordering::SeqCst) >= self.fail_after + { + return Err(DataFusionError::ResourcesExhausted( + "injected column allocation failure".into(), + )); + } + self.inner.try_grow(reservation, additional) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn column_failure_is_not_successful_eof() { + for (fail_after, max_bytes) in [ + (0, None), + (2, None), + (0, Some(10_000_000)), + (2, Some(10_000_000)), + ] { + let pool = Arc::new(FailingPool { + inner: Default::default(), + calls: AtomicUsize::new(0), + fail_after, + }); + let batches = vec![strings(128, 1024); 10]; + let error = tokio::time::timeout( + Duration::from_secs(10), + parallel(&batches, properties(100_000, max_bytes), pool.clone(), 1), + ) + .await + .expect("error propagation stalled") + .expect_err("column failure was discarded"); + assert!( + error + .to_string() + .contains("injected column allocation failure"), + "{error}" + ); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn dropping_a_group_cancels_workers_and_releases_memory() { + let batch = strings(128, 1024); + let pool = Arc::new(UnboundedMemoryPool::default()); + let props = properties(100_000, Some(10_000_000)); + let writer = + ArrowWriter::try_new(Vec::new(), batch.schema(), Some(props.clone())).unwrap(); + let (_, factory) = writer.into_serialized_writer().unwrap(); + let ctx = context(batch.schema(), props, pool.clone(), 2); + let mut group = InProgressRowGroup::new(&factory, 0, &ctx, &Time::new()).unwrap(); + group.write(&batch, &ctx.schema).await.unwrap(); + group.synchronize().await.unwrap(); + assert!(pool.reserved() > 0); + drop(group); + tokio::time::timeout(Duration::from_secs(5), async { + while pool.reserved() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled workers retained memory"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn byte_limit_finishes_a_group_before_more_input_arrives() { + let batch = integers(1024, 0); + let props = properties(2000, Some(1)); + let writer = + ArrowWriter::try_new(Vec::new(), batch.schema(), Some(props.clone())).unwrap(); + let (_, factory) = writer.into_serialized_writer().unwrap(); + let ctx = context( + batch.schema(), + props, + Arc::new(UnboundedMemoryPool::default()), + 2, + ); + let (data_tx, data_rx) = mpsc::channel(1); + let (serialize_tx, mut serialize_rx) = mpsc::channel(1); + let dispatcher = spawn_parquet_parallel_serialization_task( + factory, + data_rx, + serialize_tx, + ctx, + Time::new(), + ); + data_tx.send(batch).await.unwrap(); + // Keep the input open: the byte threshold must flush without a next batch + // or EOF, even though the first batch exceeds the target by itself. + let group = tokio::time::timeout(Duration::from_secs(5), serialize_rx.recv()) + .await + .expect("byte boundary waited for more input") + .unwrap(); + let (_, _, rows) = group.join_unwind().await.unwrap().unwrap(); + assert_eq!(rows, 1024); + drop(data_tx); + dispatcher.join_unwind().await.unwrap().unwrap(); + assert!(serialize_rx.recv().await.is_none()); +} + +#[cfg(feature = "parquet_encryption")] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn byte_boundaries_preserve_encrypted_row_group_ordinals() { + use parquet::arrow::arrow_reader::ArrowReaderOptions; + use parquet::encryption::decrypt::FileDecryptionProperties; + + let key = b"0123456789012345".to_vec(); + let encryption = FileEncryptionProperties::builder(key.clone()) + .build() + .unwrap(); + let decryption = FileDecryptionProperties::builder(key).build().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(3000)) + .set_max_row_group_bytes(Some(1)) + .with_file_encryption_properties(encryption) + .build(); + let batches = [ + integers(1024, 0), + integers(1024, 1024), + integers(1024, 2048), + ]; + let pool = Arc::new(UnboundedMemoryPool::default()); + let (metadata, bytes) = tokio::time::timeout( + Duration::from_secs(10), + parallel(&batches, props, pool.clone(), 2), + ) + .await + .expect("encrypted write stalled") + .unwrap(); + assert_eq!(row_counts(&metadata), vec![1024; 3]); + let options = ArrowReaderOptions::new().with_file_decryption_properties(decryption); + let decoded = ParquetRecordBatchReaderBuilder::try_new_with_options( + Bytes::from(bytes), + options, + ) + .unwrap() + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!( + concat_batches(&batches[0].schema(), &batches).unwrap(), + concat_batches(&batches[0].schema(), &decoded).unwrap(), + ); + assert_eq!(pool.reserved(), 0); +} diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index f2a587ed72ae..6acd90868bff 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -417,7 +417,7 @@ datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. -datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. +datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum estimated encoded size of each row group in bytes. When set, either this target or `max_row_group_size` triggers a flush. The first batch, subject to the row limit, is written before its size can be estimated. Subsequent batches are split using the observed average row size, so this is not a hard byte or memory limit. The parallel writer synchronizes column feedback to match the single-threaded writer's boundaries for the same batches. Columns within each slice still encode in parallel, but synchronization reduces overlap between slices and may lower write throughput. If `None` (the default), only the row-count limit applies. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.maximum_parallel_row_group_writers 1 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. diff --git a/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt b/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt index 8de83329ae07..e83da2ba31f3 100644 --- a/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt +++ b/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt @@ -90,13 +90,10 @@ OPTIONS ('format.max_row_group_bytes' 0); # changes how the writer splits row groups, and that combining it with # `max_row_group_size` flushes on whichever limit is reached first. # -# NOTE: byte-based flushing is currently honored only by the single-threaded -# Parquet writer (`AsyncArrowWriter`/`ArrowWriter`), which encodes inline and -# can therefore observe the in-progress row group's encoded size. The -# multi-threaded (parallel) writer decides row-group boundaries by row count -# only and ignores `max_row_group_bytes`, so these cases force the -# single-threaded path with `allow_single_file_parallelism = false`. Extending -# the parallel writer to honor the byte limit is a follow-up change. +# Both Parquet writers honor `max_row_group_bytes`: the single-threaded writer +# (`AsyncArrowWriter`/`ArrowWriter`) and the multi-threaded parallel writer. The +# cases below run under both paths (forcing `allow_single_file_parallelism` to +# `false`, then `true`) and produce the same row-group counts. # ----------------------------------------------------------------------------- statement ok @@ -108,6 +105,7 @@ set datafusion.execution.minimum_parallel_output_files = 1; statement ok set datafusion.execution.batch_size = 1024; +# Single-threaded writer. statement ok set datafusion.execution.parquet.allow_single_file_parallelism = false; @@ -173,6 +171,68 @@ Plan with Metrics 01)FilterExec: id@0 >= 0 02)--DataSourceExec: row_groups_pruned_statistics=4 total +# Parallel writer: the same assertions exercising the multi-threaded path. The +# row-group counts match the single-threaded writer's above. +statement ok +set datafusion.execution.parquet.allow_single_file_parallelism = true; + +# Row-count limit only -> 5 row groups (parallel). +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_only_parallel/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 1000); + +statement ok +CREATE EXTERNAL TABLE rg_count_size_only_parallel +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_only_parallel/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_size_only_parallel WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=5 total + +# Both limits set, whichever first -> 8 row groups (parallel). +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_and_bytes_parallel/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 1000, 'format.max_row_group_bytes' 1); + +statement ok +CREATE EXTERNAL TABLE rg_count_size_and_bytes_parallel +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_and_bytes_parallel/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_size_and_bytes_parallel WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=8 total + +# Byte limit drives alone -> 4 row groups (parallel). +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_bytes_only_parallel/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 100000, 'format.max_row_group_bytes' 1); + +statement ok +CREATE EXTERNAL TABLE rg_count_bytes_only_parallel +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_bytes_only_parallel/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_bytes_only_parallel WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=4 total + statement ok reset datafusion.execution.parquet.allow_single_file_parallelism; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ab344028cdab..8595c7ea4c12 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -103,7 +103,7 @@ The following configuration settings are available: | datafusion.execution.parquet.dictionary_page_size_limit | 1048576 | (writing) Sets best effort maximum dictionary page size, in bytes | | datafusion.execution.parquet.statistics_enabled | page | (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. | -| datafusion.execution.parquet.max_row_group_bytes | NULL | (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. | +| datafusion.execution.parquet.max_row_group_bytes | NULL | (writing) Target maximum estimated encoded size of each row group in bytes. When set, either this target or `max_row_group_size` triggers a flush. The first batch, subject to the row limit, is written before its size can be estimated. Subsequent batches are split using the observed average row size, so this is not a hard byte or memory limit. The parallel writer synchronizes column feedback to match the single-threaded writer's boundaries for the same batches. Columns within each slice still encode in parallel, but synchronization reduces overlap between slices and may lower write throughput. If `None` (the default), only the row-count limit applies. | | datafusion.execution.parquet.created_by | datafusion version 55.0.0 | (writing) Sets "created by" property | | datafusion.execution.parquet.column_index_truncate_length | 64 | (writing) Sets column index truncate length | | datafusion.execution.parquet.statistics_truncate_length | 64 | (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting | diff --git a/docs/source/user-guide/sql/format_options.md b/docs/source/user-guide/sql/format_options.md index 719fd5bd6b1e..569612722c15 100644 --- a/docs/source/user-guide/sql/format_options.md +++ b/docs/source/user-guide/sql/format_options.md @@ -142,7 +142,7 @@ The following options are available when reading or writing Parquet files. If an | BLOOM_FILTER_FPP | Yes | Sets bloom filter false positive probability (global or per column). | `'bloom_filter_fpp'` or `'bloom_filter_fpp::col'` | None | | BLOOM_FILTER_NDV | Yes | Sets bloom filter number of distinct values (global or per column). | `'bloom_filter_ndv'` or `'bloom_filter_ndv::col'` | None | | MAX_ROW_GROUP_SIZE | No | Sets the maximum number of rows per row group. Larger groups require more memory but can improve compression and scan efficiency. | `'max_row_group_size'` | 1048576 | -| MAX_ROW_GROUP_BYTES | No | Sets the maximum size of each row group in bytes. When both this and `MAX_ROW_GROUP_SIZE` are set, the row group flushes whenever either limit is reached. Mirrors `parquet.block.size` from parquet-mr. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores it. | `'max_row_group_bytes'` | None | +| MAX_ROW_GROUP_BYTES | No | Target maximum estimated encoded row-group size in bytes. Either this target or `MAX_ROW_GROUP_SIZE` triggers a flush. The first batch is accepted subject to the row limit; prediction can overshoot the byte target. | `'max_row_group_bytes'` | None | | ENABLE_PAGE_INDEX | No | If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce I/O and decoding. | `'enable_page_index'` | true | | PRUNING | No | If true, enables row group pruning based on min/max statistics. | `'pruning'` | true | | SKIP_METADATA | No | If true, skips optional embedded metadata in the file schema. | `'skip_metadata'` | true | @@ -170,6 +170,13 @@ The following options are available when reading or writing Parquet files. If an | CONTENT_DEFINED_CHUNKING_NORM_LEVEL | No | Controls how aggressively chunk boundaries are selected. Higher values can improve deduplication but increase fragmentation. The recommended range is `-3` through `3`. | `'content_defined_chunking.norm_level'` | 0 | | KEY_VALUE_METADATA | No (Key is specific) | Adds custom key-value pairs to the file metadata. Use the format `'metadata::your_key_name' 'your_value'`. Multiple entries allowed. | `'metadata::key_name'` | None | +When `MAX_ROW_GROUP_BYTES` is set, the parallel writer synchronizes column +progress between input slices. Columns within each slice still encode in +parallel, but reduced overlap between slices may lower write throughput. The +target applies to estimated encoded row-group size; it is not a hard byte limit +or a limit on total writer memory. When this option is unset (the default), only +the row-count limit applies and this synchronization is not needed. + **Example:** ```sql