diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index 025b34eaf0393..bc45f7b18159b 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -1695,6 +1695,53 @@ mod tests { Ok(()) } + #[tokio::test] + async fn parquet_sink_write_without_retained_metadata() -> Result<()> { + let field_a = Field::new("a", DataType::Utf8, false); + let field_b = Field::new("b", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let object_store_url = ObjectStoreUrl::local_filesystem(); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: object_store_url.clone(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ListingTableUrl::parse("file:///")?], + output_schema: schema.clone(), + table_partition_cols: vec![], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: false, + file_extension: "parquet".into(), + file_output_mode: FileOutputMode::Automatic, + }; + let parquet_sink = Arc::new( + ParquetSink::new(file_sink_config, TableParquetOptions::default()) + .with_retain_written_metadata(false), + ); + + let col_a: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar"])); + let col_b: ArrayRef = Arc::new(StringArray::from(vec!["baz", "baz"])); + let batch = RecordBatch::try_from_iter(vec![("a", col_a), ("b", col_b)])?; + + let rows_written = 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()), + ) + .await?; + + assert_eq!(rows_written, 2, "the file is still written and counted"); + assert!( + parquet_sink.written().is_empty(), + "no file metadata is retained" + ); + + Ok(()) + } + #[tokio::test] async fn parquet_sink_write_memory_reservation() -> Result<()> { async fn test_memory_reservation(global: ParquetOptions) -> Result<()> { diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index cfced8bbb9903..c4ac1ee7afea8 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -710,7 +710,11 @@ impl FileFormat for ParquetFormat { let sink = Arc::new( ParquetSink::new(conf, self.options.clone()) - .with_sorting_columns(sorting_columns), + .with_sorting_columns(sorting_columns) + // Nothing on this path reads `ParquetSink::written`, and keeping + // every file's metadata would grow memory with the number of + // files a statement writes. + .with_retain_written_metadata(false), ); Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 86b80b9654ff2..35f18ac3c6eb8 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -61,15 +61,19 @@ const INITIAL_BUFFER_BYTES: usize = 1048576; /// this size, it is flushed to object store const BUFFER_FLUSH_BYTES: usize = 1024000; +/// Metadata of the files a [`ParquetSink`] has written, keyed by path. Shared +/// with the writer tasks, which insert into it as they close their files. +type WrittenFiles = Arc>>; + /// Implements [`DataSink`] for writing to a parquet file. pub struct ParquetSink { /// Config options for writing data 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>>, + /// Metadata of every file written so far, for [`Self::written`]. `None` + /// once retention was turned off with [`Self::with_retain_written_metadata`]. + written: Option, /// Optional sorting columns to write to Parquet metadata sorting_columns: Option>, /// Metrics for tracking write operations @@ -104,7 +108,7 @@ impl ParquetSink { Self { config, parquet_options, - written: Default::default(), + written: Some(Default::default()), sorting_columns: None, metrics: ExecutionPlanMetricsSet::new(), } @@ -119,10 +123,29 @@ impl ParquetSink { self } + /// Whether to keep the [`ParquetMetaData`] of every written file so that it + /// can be read back with [`Self::written`]. Defaults to `true`. + /// + /// The metadata carries every column chunk's statistics and page indexes + /// and is held until the sink is dropped, so a write that produces many + /// files pays for it in memory. Turn it off when nothing reads + /// [`Self::written`], as the sinks `ParquetFormat` creates for `COPY` and + /// `INSERT` do. + pub fn with_retain_written_metadata(mut self, retain: bool) -> Self { + self.written = retain.then(Default::default); + self + } + /// Retrieve the file metadata for the written files, keyed to the path /// which may be partitioned (in the case of hive style partitioning). + /// + /// Empty if retention was turned off with + /// [`Self::with_retain_written_metadata`]. pub fn written(&self) -> HashMap { - self.written.lock().clone() + self.written + .as_ref() + .map(|written| written.lock().clone()) + .unwrap_or_default() } /// Create writer properties based upon configuration settings, @@ -270,9 +293,10 @@ impl FileSink for ParquetSink { let parquet_opts = &self.parquet_options; - let mut file_write_tasks: JoinSet< - std::result::Result<(Path, ParquetMetaData), DataFusionError>, - > = JoinSet::new(); + // Writer tasks hand back only their file's row and byte counts. The + // file's metadata is stored or dropped by `record_written_file` inside + // the task rather than buffered here until every file has been written. + let mut file_write_tasks: JoinSet> = JoinSet::new(); let runtime = context.runtime_env(); let parallel_options = ParallelParquetWriterOptions { @@ -288,6 +312,7 @@ impl FileSink for ParquetSink { let parquet_props = self .create_writer_props(&runtime, &file_metadata.path) .await?; + let written = self.written.clone(); // CDC requires the sequential writer: the chunker state lives in ArrowWriter // and persists across row groups. The parallel path bypasses ArrowWriter entirely. if !parquet_opts.global.allow_single_file_parallelism @@ -312,7 +337,11 @@ impl FileSink for ParquetSink { .close() .await .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; - Ok((file_metadata.path, parquet_meta_data)) + record_written_file( + written, + file_metadata.path, + parquet_meta_data, + ) } .with_elapsed_compute(elapsed_compute.clone()), ); @@ -349,7 +378,7 @@ impl FileSink for ParquetSink { encoding_time, ) .await?; - Ok((file_metadata.path, parquet_meta_data)) + record_written_file(written, file_metadata.path, parquet_meta_data) }); } } @@ -357,20 +386,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() { @@ -391,6 +409,36 @@ impl FileSink for ParquetSink { } } +/// Accounts for a file the sink has finished writing: returns its row count +/// and the sum of its compressed row group sizes for the `rows_written` and +/// `bytes_written` metrics, and stores its [`ParquetMetaData`] in `written` +/// when the sink retains metadata. +/// +/// Runs inside the writer task so that, with retention off, the metadata is +/// dropped as soon as the file is closed and the sink holds no per-file state. +fn record_written_file( + written: Option, + path: Path, + parquet_meta_data: ParquetMetaData, +) -> Result<(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(); + if let Some(written) = written { + let mut written = written.lock(); + written.try_insert(path, parquet_meta_data).map_err(|e| { + internal_datafusion_err!( + "duplicate entry detected for partitioned file {}: {e}", + e.entry.key() + ) + })?; + } + Ok((rows, bytes)) +} + #[async_trait] impl DataSink for ParquetSink { fn metrics(&self) -> Option { @@ -423,7 +471,8 @@ impl DataSink for ParquetSink { let Self { config: _, parquet_options: _, - // Runtime output state, not part of the plan. + // Written-file metadata and its retention switch are runtime + // state, not part of the plan. written: _, sorting_columns: _, // Runtime metrics are recreated on decode. @@ -455,7 +504,8 @@ impl TryFrom<&ParquetSink> for datafusion_proto_models::protobuf::ParquetSink { let ParquetSink { config, parquet_options, - // Runtime output state, not part of the plan. + // Written-file metadata and its retention switch are runtime + // state, not part of the plan. written: _, sorting_columns, // Runtime metrics are recreated on decode. diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 769d3fdbeb1db..6304cf0bea6c9 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1498,3 +1498,17 @@ fn is_nullable(&self) -> bool { ``` See [issue #18882](https://github.com/apache/datafusion/issues/18882) for details. + +### `ParquetSink`s created by `ParquetFormat` no longer retain written file metadata + +`ParquetSink` keeps the `ParquetMetaData` of every file it writes so that it can +be read back with `ParquetSink::written()`. That is unchanged for sinks you +construct yourself. Retention can now be turned off with +`ParquetSink::with_retain_written_metadata(false)`, and the sinks that +`ParquetFormat::create_writer_physical_plan` builds for `COPY` and `INSERT` do +so: nothing on that path reads the metadata, and keeping it made the memory of a +statement grow with the number of files written (see [#27026]). If you relied on +downcasting the sink of such a plan to call `written()`, construct the +`ParquetSink` yourself instead. + +[#27026]: https://github.com/apache/datafusion/issues/27026