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 {