saadtajwar commented on code in PR #23571: URL: https://github.com/apache/datafusion/pull/23571#discussion_r3581391840
########## datafusion/datasource-parquet/src/type_narrowing.rs: ########## @@ -0,0 +1,326 @@ +// 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. + +//! Narrow integer Arrow types in embedded Parquet `ARROW:schema` metadata +//! based on exact column chunk min/max statistics. +//! +//! Physical page encoding is left unchanged; only the Arrow schema hint in +//! file key/value metadata is rewritten so readers can advertise a tighter +//! type when the value domain fits. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Fields, Schema}; +use parquet::file::metadata::RowGroupMetaData; +use parquet::file::statistics::Statistics; + +/// If any integer fields can be narrowed from exact row-group statistics, +/// returns a new schema with those field types updated. Returns `None` when +/// nothing changes. +pub(crate) fn try_narrow_integer_schema( + schema: &Schema, + row_groups: &[RowGroupMetaData], +) -> Option<Schema> { + let domains = merge_exact_integer_domains(row_groups)?; + let mut changed = false; + let fields = narrow_fields(schema.fields(), &[], &domains, &mut changed); + if !changed { + return None; + } + + Some(Schema::new_with_metadata(fields, schema.metadata().clone())) +} + +/// Merge exact Int32/Int64 min/max across all row groups, keyed by Parquet +/// column path string (e.g. `"id"` or `"tags.list.element"`). +/// +/// A column is omitted if any row group lacks exact min/max for it. +fn merge_exact_integer_domains( + row_groups: &[RowGroupMetaData], +) -> Option<HashMap<String, (i64, i64)>> { + let mut domains: Option<HashMap<String, (i64, i64)>> = None; + + for rg in row_groups { + let mut this_rg = HashMap::new(); + for column in rg.columns() { + let path = column.column_path().string(); + let Some(stats) = column.statistics() else { + continue; + }; + if !stats.min_is_exact() || !stats.max_is_exact() { + continue; + } + let Some((min, max)) = int_min_max(stats) else { + continue; + }; + this_rg.insert(path, (min, max)); + } + + match &mut domains { + None => domains = Some(this_rg), + Some(merged) => { + merged.retain(|path, (file_min, file_max)| { + let Some((min, max)) = this_rg.get(path) else { + return false; + }; + *file_min = (*file_min).min(*min); + *file_max = (*file_max).max(*max); + true + }); + } + } + } + + domains +} + +/// Note that we're ignoring types like `Statistics::Int96` because if we can't cast to i64 then it's too big to downcast +fn int_min_max(stats: &Statistics) -> Option<(i64, i64)> { + match stats { + Statistics::Int32(val) => { + let min = val.min_opt().copied().map(i64::from)?; + let max = val.max_opt().copied().map(i64::from)?; + Some((min, max)) + } + Statistics::Int64(val) => { + let min = *val.min_opt()?; + let max = *val.max_opt()?; + Some((min, max)) + } + _ => None, + } +} + +/// Narrow to the smallest possible integer type that can represent all values +pub(crate) fn narrowest_integer_type(min: i64, max: i64) -> DataType { Review Comment: This was taken from Vortex's implementation [here](https://github.com/vortex-data/vortex/blob/5abaf9823dee973dde7295a6a36234935f08d060/vortex-array/src/arrays/primitive/array/mod.rs#L151)! If I should add a comment in the code itself for attribution then happy to do so! ########## datafusion/datasource-parquet/src/type_narrowing.rs: ########## @@ -0,0 +1,326 @@ +// 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. + +//! Narrow integer Arrow types in embedded Parquet `ARROW:schema` metadata +//! based on exact column chunk min/max statistics. +//! +//! Physical page encoding is left unchanged; only the Arrow schema hint in +//! file key/value metadata is rewritten so readers can advertise a tighter +//! type when the value domain fits. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Fields, Schema}; +use parquet::file::metadata::RowGroupMetaData; +use parquet::file::statistics::Statistics; + +/// If any integer fields can be narrowed from exact row-group statistics, +/// returns a new schema with those field types updated. Returns `None` when +/// nothing changes. +pub(crate) fn try_narrow_integer_schema( Review Comment: This is currently very integer-specific - I could rename this to a generic narrow, and then make this more extensible with the integer narrow + the ability to easily add more narrowing functionality for other types? ########## datafusion/datasource-parquet/src/type_narrowing.rs: ########## @@ -0,0 +1,326 @@ +// 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. + +//! Narrow integer Arrow types in embedded Parquet `ARROW:schema` metadata +//! based on exact column chunk min/max statistics. +//! +//! Physical page encoding is left unchanged; only the Arrow schema hint in +//! file key/value metadata is rewritten so readers can advertise a tighter +//! type when the value domain fits. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Fields, Schema}; +use parquet::file::metadata::RowGroupMetaData; +use parquet::file::statistics::Statistics; + +/// If any integer fields can be narrowed from exact row-group statistics, +/// returns a new schema with those field types updated. Returns `None` when +/// nothing changes. +pub(crate) fn try_narrow_integer_schema( + schema: &Schema, + row_groups: &[RowGroupMetaData], +) -> Option<Schema> { + let domains = merge_exact_integer_domains(row_groups)?; + let mut changed = false; + let fields = narrow_fields(schema.fields(), &[], &domains, &mut changed); + if !changed { + return None; + } Review Comment: We could just not use this `changed` flag at all and we would just return the same schema if no narrowing could be done? The main downside I saw was increasing the size of the Parquet file given that we append this schema if it's changed ########## datafusion/datasource-parquet/src/sink.rs: ########## @@ -306,12 +332,25 @@ impl FileSink for ParquetSink { .await?; let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) .register(context.memory_pool()); + let write_arrow_metadata = !parquet_opts.global.skip_arrow_metadata; + let writer_schema = get_writer_schema(&self.config); file_write_tasks.spawn( async move { while let Some(batch) = rx.recv().await { writer.write(&batch).await?; reservation.try_resize(writer.memory_size())?; } + + // Flushing here to get a complete view of the statistics for potential type narrowing before writer is closed Review Comment: I'm not sure if there's a better way to do this without flushing - it looked like the writer disallows any modifications after `close` is invoked -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
