From f5e5f512334c76412d6dcff18e5c1e2dcd7cac0b Mon Sep 17 00:00:00 2001 From: "xudong.w" Date: Tue, 3 Nov 2026 08:21:24 +0800 Subject: [PATCH] Stop ParquetSink from retaining every written file's ParquetMetaData ParquetSink kept each written file's ParquetMetaData in a HashMap for the lifetime of the sink, and its writer tasks also handed the metadata back through a JoinSet that is only drained once the last file has been announced. Nothing in the engine read either; only the datafusion-core unit tests used ParquetSink::written(). For Atlas this meant a 5-day options NBBO weekly rollup (~700 output files) grew ~1 GiB of RSS per 10 GB written outside the memory pool until the 64 GiB container limit killed it. Writer tasks now reduce the closed file's metadata to its row count and compressed byte count for the rows_written / bytes_written metrics and drop the ParquetMetaData inside the task. The written map, ParquetSink::written() and the parking_lot dependency go away. The core tests read the footers back from a temporary local object store instead. --- Cargo.lock | 1 - .../src/datasource/file_format/parquet.rs | 174 ++++++++++-------- datafusion/datasource-parquet/Cargo.toml | 1 - datafusion/datasource-parquet/src/sink.rs | 52 +++--- 4 files changed, 122 insertions(+), 106 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ffdba30cb5ae6..b900a1ece895b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2084,7 +2084,6 @@ dependencies = [ "itertools 0.15.0", "log", "object_store", - "parking_lot", "parquet", "tempfile", "tokio", diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index 0f5db4a057d76..6f0b814b3002a 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -154,11 +154,11 @@ mod tests { use async_trait::async_trait; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource_parquet::metadata::DFParquetMetadata; - use futures::StreamExt; use futures::stream::BoxStream; + use futures::{StreamExt, TryStreamExt}; use insta::assert_snapshot; use object_store::local::LocalFileSystem; - use object_store::{CopyOptions, ObjectMeta}; + use object_store::{CopyOptions, ObjectMeta, ObjectStoreExt}; use object_store::{ GetOptions, GetResult, ListResult, MultipartUpload, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, path::Path, @@ -167,7 +167,7 @@ mod tests { use parquet::arrow::arrow_reader::ArrowReaderOptions; use parquet::file::metadata::{ KeyValue, PageIndexPolicy, ParquetColumnIndex, ParquetMetaData, - ParquetOffsetIndex, + ParquetMetaDataReader, ParquetOffsetIndex, }; use parquet::file::page_index::column_index::ColumnIndexMetaData; use tokio::fs::File; @@ -1250,39 +1250,78 @@ mod tests { Ok(()) } - fn build_ctx(store_url: &url::Url) -> Arc { - let tmp_dir = tempfile::TempDir::new().unwrap(); - let local = Arc::new( - LocalFileSystem::new_with_prefix(&tmp_dir) - .expect("should create object store"), - ); + /// Temporary local object store for a test [`ParquetSink`] to write into. + /// The directory and the files written to it go away when this is dropped. + struct SinkTestStore { + store: Arc, + _tmp_dir: tempfile::TempDir, + } - let mut session = SessionConfig::default(); - let mut parquet_opts = ParquetOptions { - allow_single_file_parallelism: true, - ..Default::default() - }; - parquet_opts.allow_single_file_parallelism = true; - session.options_mut().execution.parquet = parquet_opts; - - let runtime = RuntimeEnv::default(); - runtime - .object_store_registry - .register_store(store_url, local); - - Arc::new( - TaskContext::default() - .with_session_config(session) - .with_runtime(Arc::new(runtime)), - ) + impl SinkTestStore { + fn new() -> Self { + let tmp_dir = tempfile::TempDir::new().unwrap(); + let store: Arc = Arc::new( + LocalFileSystem::new_with_prefix(&tmp_dir) + .expect("should create object store"), + ); + Self { + store, + _tmp_dir: tmp_dir, + } + } + + /// Task context with this store registered under `store_url`. + fn task_context(&self, store_url: &url::Url) -> Arc { + let mut session = SessionConfig::default(); + session.options_mut().execution.parquet = ParquetOptions { + allow_single_file_parallelism: true, + ..Default::default() + }; + + let runtime = RuntimeEnv::default(); + runtime + .object_store_registry + .register_store(store_url, Arc::clone(&self.store)); + + Arc::new( + TaskContext::default() + .with_session_config(session) + .with_runtime(Arc::new(runtime)), + ) + } + + /// The footer of every file in the store, keyed by the file's path within + /// the store (so including any hive partition directories). + async fn written(&self) -> Result> { + let files: Vec = self.store.list(None).try_collect().await?; + let mut written = Vec::with_capacity(files.len()); + for file in files { + let bytes = self.store.get(&file.location).await?.bytes().await?; + let metadata = ParquetMetaDataReader::new().parse_and_finish(&bytes)?; + written.push((file.location, metadata)); + } + Ok(written) + } + + /// [`Self::written`] for a sink that is expected to have written one file. + async fn single_written(&self) -> Result<(Path, ParquetMetaData)> { + let mut written = self.written().await?; + assert_eq!( + written.len(), + 1, + "expected a single parquet file to be written, instead found {}", + written.len() + ); + Ok(written.remove(0)) + } } #[tokio::test] async fn parquet_sink_write() -> Result<()> { - let parquet_sink = create_written_parquet_sink("file:///").await?; + let store = create_written_parquet_sink("file:///").await?; // assert written to proper path - let (path, file_metadata) = get_written(parquet_sink)?; + let (path, file_metadata) = store.single_written().await?; let path_parts = path.parts().collect::>(); assert_eq!(path_parts.len(), 1, "should not have path prefix"); @@ -1316,11 +1355,10 @@ mod tests { ..Default::default() }; - let parquet_sink = - create_written_parquet_sink_using_config("file:///", opts).await?; + let store = create_written_parquet_sink_using_config("file:///", opts).await?; // assert written to proper path - let (path, file_metadata) = get_written(parquet_sink)?; + let (path, file_metadata) = store.single_written().await?; let path_parts = path.parts().collect::>(); assert_eq!(path_parts.len(), 1, "should not have path prefix"); @@ -1424,9 +1462,8 @@ mod tests { skip_arrow_metadata: true, ..Default::default() }; - let parquet_sink = - create_written_parquet_sink_using_config("file:///", opts).await?; - let (_, file_metadata) = get_written(parquet_sink)?; + let store = create_written_parquet_sink_using_config("file:///", opts).await?; + let (_, file_metadata) = store.single_written().await?; assert_file_metadata(file_metadata, &expected_without); // single threaded write, do not skip insert @@ -1435,9 +1472,8 @@ mod tests { skip_arrow_metadata: false, ..Default::default() }; - let parquet_sink = - create_written_parquet_sink_using_config("file:///", opts).await?; - let (_, file_metadata) = get_written(parquet_sink)?; + let store = create_written_parquet_sink_using_config("file:///", opts).await?; + let (_, file_metadata) = store.single_written().await?; assert_file_metadata(file_metadata, &expected_with); // multithreaded write, skip insert @@ -1448,9 +1484,8 @@ mod tests { skip_arrow_metadata: true, ..Default::default() }; - let parquet_sink = - create_written_parquet_sink_using_config("file:///", opts).await?; - let (_, file_metadata) = get_written(parquet_sink)?; + let store = create_written_parquet_sink_using_config("file:///", opts).await?; + let (_, file_metadata) = store.single_written().await?; assert_file_metadata(file_metadata, &expected_without); // multithreaded write, do not skip insert @@ -1461,9 +1496,8 @@ mod tests { skip_arrow_metadata: false, ..Default::default() }; - let parquet_sink = - create_written_parquet_sink_using_config("file:///", opts).await?; - let (_, file_metadata) = get_written(parquet_sink)?; + let store = create_written_parquet_sink_using_config("file:///", opts).await?; + let (_, file_metadata) = store.single_written().await?; assert_file_metadata(file_metadata, &expected_with); Ok(()) @@ -1473,10 +1507,10 @@ mod tests { async fn parquet_sink_write_with_extension() -> Result<()> { let filename = "test_file.custom_ext"; let file_path = format!("file:///path/to/{filename}"); - let parquet_sink = create_written_parquet_sink(file_path.as_str()).await?; + let store = create_written_parquet_sink(file_path.as_str()).await?; // assert written to proper path - let (path, _) = get_written(parquet_sink)?; + let (path, _) = store.single_written().await?; let path_parts = path.parts().collect::>(); assert_eq!( path_parts.len(), @@ -1492,10 +1526,10 @@ mod tests { #[tokio::test] async fn parquet_sink_write_with_directory_name() -> Result<()> { let file_path = "file:///path/to"; - let parquet_sink = create_written_parquet_sink(file_path).await?; + let store = create_written_parquet_sink(file_path).await?; // assert written to proper path - let (path, _) = get_written(parquet_sink)?; + let (path, _) = store.single_written().await?; let path_parts = path.parts().collect::>(); assert_eq!( path_parts.len(), @@ -1511,10 +1545,10 @@ mod tests { #[tokio::test] async fn parquet_sink_write_with_folder_ending() -> Result<()> { let file_path = "file:///path/to/"; - let parquet_sink = create_written_parquet_sink(file_path).await?; + let store = create_written_parquet_sink(file_path).await?; // assert written to proper path - let (path, _) = get_written(parquet_sink)?; + let (path, _) = store.single_written().await?; let path_parts = path.parts().collect::>(); assert_eq!( path_parts.len(), @@ -1527,7 +1561,7 @@ mod tests { Ok(()) } - async fn create_written_parquet_sink(table_path: &str) -> Result> { + async fn create_written_parquet_sink(table_path: &str) -> Result { create_written_parquet_sink_using_config(table_path, ParquetOptions::default()) .await } @@ -1537,7 +1571,7 @@ mod tests { async fn create_written_parquet_sink_using_config( table_path: &str, global: ParquetOptions, - ) -> Result> { + ) -> Result { // schema should match the ENCODED_ARROW_SCHEMA bove let field_a = Field::new("a", DataType::Utf8, false); let field_b = Field::new("b", DataType::Utf8, false); @@ -1556,7 +1590,7 @@ mod tests { file_extension: "parquet".into(), file_output_mode: FileOutputMode::Automatic, }; - let parquet_sink = Arc::new(ParquetSink::new( + let parquet_sink = ParquetSink::new( file_sink_config, TableParquetOptions { key_value_metadata: std::collections::HashMap::from([ @@ -1566,7 +1600,7 @@ mod tests { global, ..Default::default() }, - )); + ); // create data let col_a: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar"])); @@ -1574,31 +1608,18 @@ mod tests { let batch = RecordBatch::try_from_iter(vec![("a", col_a), ("b", col_b)])?; // write stream + let store = SinkTestStore::new(); FileSink::write_all( - parquet_sink.as_ref(), + &parquet_sink, Box::pin(RecordBatchStreamAdapter::new( schema, futures::stream::iter(vec![Ok(batch)]), )), - &build_ctx(object_store_url.as_ref()), + &store.task_context(object_store_url.as_ref()), ) .await?; - Ok(parquet_sink) - } - - fn get_written(parquet_sink: Arc) -> Result<(Path, ParquetMetaData)> { - let mut written = parquet_sink.written(); - let written = written.drain(); - assert_eq!( - written.len(), - 1, - "expected a single parquet files to be written, instead found {}", - written.len() - ); - - let (path, parquet_meta_data) = written.take(1).next().unwrap(); - Ok((path, parquet_meta_data)) + Ok(store) } fn assert_file_metadata( @@ -1659,19 +1680,19 @@ mod tests { let batch = RecordBatch::try_from_iter(vec![("a", col_a), ("b", col_b)])?; // write stream + let store = SinkTestStore::new(); FileSink::write_all( parquet_sink.as_ref(), Box::pin(RecordBatchStreamAdapter::new( schema, futures::stream::iter(vec![Ok(batch)]), )), - &build_ctx(object_store_url.as_ref()), + &store.task_context(object_store_url.as_ref()), ) .await?; // assert written - let mut written = parquet_sink.written(); - let written = written.drain(); + let written = store.written().await?; assert_eq!( written.len(), 2, @@ -1681,7 +1702,7 @@ mod tests { // check the file metadata includes partitions let mut expected_partitions = std::collections::HashSet::from(["a=foo", "a=bar"]); - for (path, parquet_metadata) in written.take(2) { + for (path, parquet_metadata) in written { let file_metadata = parquet_metadata.file_metadata(); let schema = file_metadata.schema_descr(); let num_rows = file_metadata.num_rows(); @@ -1757,7 +1778,8 @@ mod tests { let batch = RecordBatch::try_from_iter(vec![("a", col_a), ("b", col_b)])?; // create task context - let task_context = build_ctx(object_store_url.as_ref()); + let store = SinkTestStore::new(); + let task_context = store.task_context(object_store_url.as_ref()); assert_eq!( task_context.memory_pool().reserved(), 0, diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a2589af19a6ee..7fe68a9dedebe 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -53,7 +53,6 @@ futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } object_store = { workspace = true } -parking_lot = { workspace = true } parquet = { workspace = true } tokio = { workspace = true } diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 53f6f1e6b4323..6a07296e0868b 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -27,7 +27,7 @@ use arrow::array::RecordBatch; use arrow::datatypes::{Schema, SchemaRef}; use async_trait::async_trait; use datafusion_common::config::TableParquetOptions; -use datafusion_common::{DataFusionError, HashMap, Result, internal_datafusion_err}; +use datafusion_common::{DataFusionError, Result}; use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_compression_type::FileCompressionType; @@ -81,9 +81,6 @@ pub struct ParquetSink { config: FileSinkConfig, /// Underlying parquet options parquet_options: TableParquetOptions, - /// File metadata from successfully produced parquet files. The Mutex is only used - /// to allow inserting to HashMap from behind borrowed reference in DataSink::write_all. - written: Arc>>, /// Optional sorting columns to write to Parquet metadata sorting_columns: Option>, /// Metrics for tracking write operations @@ -118,7 +115,6 @@ impl ParquetSink { Self { config, parquet_options, - written: Default::default(), sorting_columns: None, metrics: ExecutionPlanMetricsSet::new(), } @@ -133,12 +129,6 @@ impl ParquetSink { self } - /// Retrieve the file metadata for the written files, keyed to the path - /// which may be partitioned (in the case of hive style partitioning). - pub fn written(&self) -> HashMap { - self.written.lock().clone() - } - /// Create writer properties based upon configuration settings, /// including partitioning and the inclusion of arrow schema metadata. async fn create_writer_props( @@ -280,9 +270,9 @@ impl FileSink for ParquetSink { let parquet_opts = &self.parquet_options; - let mut file_write_tasks: JoinSet< - std::result::Result<(Path, ParquetMetaData), DataFusionError>, - > = JoinSet::new(); + // Each task reports only its file's row and byte counts; see + // `written_file_counts` for why the file's metadata itself stays in the task. + let mut file_write_tasks: JoinSet> = JoinSet::new(); let runtime = context.runtime_env(); let parallel_options = ParallelParquetWriterOptions { @@ -319,7 +309,7 @@ impl FileSink for ParquetSink { .close() .await .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; - Ok((path, parquet_meta_data)) + Ok(written_file_counts(&parquet_meta_data)) } .with_elapsed_compute(elapsed_compute.clone()), ); @@ -355,7 +345,7 @@ impl FileSink for ParquetSink { encoding_time, ) .await?; - Ok((path, parquet_meta_data)) + Ok(written_file_counts(&parquet_meta_data)) }); } } @@ -363,20 +353,9 @@ impl FileSink for ParquetSink { while let Some(result) = file_write_tasks.join_next().await { match result { Ok(r) => { - let (path, parquet_meta_data) = r?; - let file_rows = parquet_meta_data.file_metadata().num_rows() as usize; - let file_bytes: usize = parquet_meta_data - .row_groups() - .iter() - .map(|rg| rg.compressed_size() as usize) - .sum(); + let (file_rows, file_bytes) = r?; rows_written_counter.add(file_rows); bytes_written_counter.add(file_bytes); - let mut written_files = self.written.lock(); - written_files - .try_insert(path.clone(), parquet_meta_data) - .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; - drop(written_files); } Err(e) => { if e.is_panic() { @@ -397,6 +376,23 @@ impl FileSink for ParquetSink { } } +/// Reduces a finished file's metadata to its `(rows, compressed bytes)` for the +/// `rows_written` and `bytes_written` metrics. +/// +/// The writer task drops the [`ParquetMetaData`] right after this, so the sink +/// holds no per-file state once a file is closed. The metadata carries every +/// column chunk's statistics and page indexes, and keeping it for every file +/// until the sink finished made memory grow with the number of files written. +fn written_file_counts(parquet_meta_data: &ParquetMetaData) -> (usize, usize) { + let rows = parquet_meta_data.file_metadata().num_rows() as usize; + let bytes = parquet_meta_data + .row_groups() + .iter() + .map(|rg| rg.compressed_size() as usize) + .sum(); + (rows, bytes) +} + #[async_trait] impl DataSink for ParquetSink { fn metrics(&self) -> Option {