diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 360586b0e9bae..3cebb01cc0763 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1413,6 +1413,17 @@ config_namespace! { /// Defaults to 20. pub max_in_list_size: usize, default = 20 + /// (reading) If true, top-level string and binary Parquet columns with + /// dictionary pages are inferred and scanned as + /// `Dictionary` / `Dictionary` instead of + /// their plain value type. + /// + /// This applies only when DataFusion infers the table schema. Tables with + /// a user-supplied schema are not promoted because the Parquet footer is + /// not read at DDL time, so dictionary pages cannot be detected per column. + /// See + pub enable_rle_to_dictionary: bool, default = false + // The following options affect writing to parquet files // and map to parquet::file::properties::WriterProperties diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs index c50bb42a38ef7..25302b0dedae9 100644 --- a/datafusion/common/src/file_options/parquet_writer.rs +++ b/datafusion/common/src/file_options/parquet_writer.rs @@ -249,6 +249,7 @@ impl ParquetOptions { skip_arrow_metadata: _, max_predicate_cache_size: _, max_in_list_size: _, + enable_rle_to_dictionary: _, } = self; let mut builder = WriterProperties::builder() @@ -508,6 +509,7 @@ mod tests { coerce_int96_tz: None, max_predicate_cache_size: defaults.max_predicate_cache_size, content_defined_chunking: defaults.content_defined_chunking.clone(), + enable_rle_to_dictionary: defaults.enable_rle_to_dictionary, } } @@ -630,6 +632,8 @@ mod tests { coerce_int96: None, coerce_int96_tz: None, content_defined_chunking: props.content_defined_chunking().into(), + enable_rle_to_dictionary: global_options_defaults + .enable_rle_to_dictionary, }, column_specific_options, key_value_metadata, diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 18f2b5a650c8d..f5b62e1de4b72 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -358,14 +358,16 @@ impl FileFormat for ParquetFormat { &object.location, ) .await?; - let result = DFParquetMetadata::new(store.as_ref(), object) + let meta = DFParquetMetadata::new(store.as_ref(), object) .with_metadata_size_hint(self.metadata_size_hint()) .with_decryption_properties(file_decryption_properties) .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) .with_coerce_int96(coerce_int96) .with_coerce_int96_tz(coerce_int96_tz.clone()) - .fetch_schema_with_location() - .await?; + .with_enable_rle_to_dictionary( + self.options.global.enable_rle_to_dictionary, + ); + let result = meta.fetch_schema_with_location().await?; Ok::<_, DataFusionError>(result) }) .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 @@ -400,7 +402,12 @@ impl FileFormat for ParquetFormat { } drop(seen); - let schemas = schemas.into_iter().map(|(_, schema)| schema); + // Normalize dict-promoted schemas before merging so mixed dict/plain files merge cleanly. + let mut schemas: Vec = + schemas.into_iter().map(|(_, schema)| schema).collect(); + if self.options.global.enable_rle_to_dictionary { + schemas = crate::schema_coercion::uniform_dict_schemas(schemas); + } let schema = if self.skip_metadata() { Schema::try_merge(clear_metadata(schemas)) @@ -522,7 +529,7 @@ impl FileFormat for ParquetFormat { source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); if let Some(metadata_size_hint) = metadata_size_hint { - source = source.with_metadata_size_hint(metadata_size_hint) + source = source.with_metadata_size_hint(metadata_size_hint); } source = self.set_source_encryption_factory(source, state)?; @@ -736,6 +743,7 @@ impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions { compression_opt: global_options.global.compression.map(|compression| { parquet_options::CompressionOpt::Compression(compression) }), + enable_rle_to_dictionary: global_options.global.enable_rle_to_dictionary, dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| { parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) }), diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 4781a309f5bcc..be7b303bbb889 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -23,11 +23,11 @@ use crate::{Int96Coercer, apply_file_schema_type_coercions}; use arrow::array::{Array, ArrayRef, BooleanArray}; use arrow::compute::kernels::cmp::eq; use arrow::compute::{and, sum}; -use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::stats::Precision; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, Result, ScalarValue, Statistics, + ColumnStatistics, DataFusionError, HashMap, HashSet, Result, ScalarValue, Statistics, internal_datafusion_err, }; use datafusion_execution::cache::cache_manager::{ @@ -146,6 +146,8 @@ pub struct DFParquetMetadata<'a> { pub coerce_int96: Option, /// Optional timezone applied to INT96-coerced timestamps. pub coerce_int96_tz: Option>, + /// If true, promote string/binary columns with dictionary pages to `Dictionary(Int32, ...)`. + enable_rle_to_dictionary: bool, } impl<'a> DFParquetMetadata<'a> { @@ -163,9 +165,16 @@ impl<'a> DFParquetMetadata<'a> { page_index_policy: None, coerce_int96: None, coerce_int96_tz: None, + enable_rle_to_dictionary: false, } } + /// Promote string/binary columns with dictionary pages to `Dictionary(Int32, ...)`. + pub fn with_enable_rle_to_dictionary(mut self, enable: bool) -> Self { + self.enable_rle_to_dictionary = enable; + self + } + /// Set a hint for the number of trailing bytes to prefetch from the end /// of the file, equivalent to /// [`ParquetMetaDataReader::with_prefetch_hint`]. @@ -439,6 +448,68 @@ impl<'a> DFParquetMetadata<'a> { .coerce() }) .unwrap_or(schema); + + let schema = if self.enable_rle_to_dictionary { + let schema_descr = file_metadata.schema_descr(); + // Top-level columns that have a dictionary page in at least one row group. + let dict_cols: HashSet = metadata + .row_groups() + .iter() + .flat_map(|rg| { + rg.columns() + .iter() + .enumerate() + .filter_map(|(col_idx, col)| { + col.dictionary_page_offset()?; + let col_desc = schema_descr.column(col_idx); + let parts = col_desc.path().parts(); + // Skip nested columns: their leaf name doesn't match the + // Arrow top-level field name. + (parts.len() == 1).then(|| parts[0].clone()) + }) + }) + .collect(); + if dict_cols.is_empty() { + schema + } else { + let promoted: Vec<_> = schema + .fields() + .iter() + .map(|field| { + if !dict_cols.contains(field.name()) { + return Arc::clone(field); + } + let dict_value_type = match field.data_type() { + DataType::Utf8 => Some(DataType::Utf8), + DataType::LargeUtf8 => Some(DataType::LargeUtf8), + DataType::Binary => Some(DataType::Binary), + DataType::LargeBinary => Some(DataType::LargeBinary), + _ => None, + }; + dict_value_type.map_or_else( + || Arc::clone(field), + |value_type| { + Arc::new( + Field::new( + field.name(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(value_type), + ), + field.is_nullable(), + ) + .with_metadata(field.metadata().clone()), + ) + }, + ) + }) + .collect(); + Schema::new_with_metadata(promoted, schema.metadata().clone()) + } + } else { + schema + }; + Ok(schema) } diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index ec24462db0564..3d010665570b5 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -35,7 +35,7 @@ use crate::row_group_filter::{RowGroupAccessPlanFilter, row_group_in_range}; use crate::{ BloomFilterStatistics, Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, ParquetRowSelection, ParquetVirtualColumn, - apply_file_schema_type_coercions, + schema_coercion::apply_file_schema_type_coercions_with_rle, }; use arrow::array::RecordBatch; use arrow::datatypes::DataType; @@ -294,6 +294,8 @@ pub(super) struct ParquetMorselizer { /// lists skip container-level pruning. Sourced from /// `datafusion.execution.parquet.max_in_list_size`. pub max_in_list_size: usize, + /// Whether to ask arrow-rs to read promoted dictionary columns directly. + pub enable_rle_to_dictionary: bool, /// Whether to read row groups in reverse order pub reverse_row_groups: bool, /// Optional sort order used to reorder row groups by their min/max statistics. @@ -467,6 +469,7 @@ struct PreparedParquetOpen { predicate_creation_errors: Count, max_predicate_cache_size: Option, max_in_list_size: usize, + enable_rle_to_dictionary: bool, reverse_row_groups: bool, sort_order_for_reorder: Option, preserve_order: bool, @@ -874,6 +877,7 @@ impl ParquetMorselizer { predicate_creation_errors, max_predicate_cache_size: self.max_predicate_cache_size, max_in_list_size: self.max_in_list_size, + enable_rle_to_dictionary: self.enable_rle_to_dictionary, reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: self.sort_order_for_reorder.clone(), preserve_order: self.preserve_order, @@ -984,9 +988,10 @@ impl MetadataLoadedParquetOpen { // desired schema (for example if we want to instruct the parquet // reader to read strings using Utf8View instead). Update if necessary let mut metadata_dirty = false; - if let Some(merged) = apply_file_schema_type_coercions( + if let Some(merged) = apply_file_schema_type_coercions_with_rle( &prepared.logical_file_schema, &physical_file_schema, + prepared.enable_rle_to_dictionary, ) { physical_file_schema = Arc::new(merged); options = options.with_schema(Arc::clone(&physical_file_schema)); @@ -1948,6 +1953,7 @@ mod test { coerce_int96: Option, max_predicate_cache_size: Option, max_in_list_size: usize, + enable_rle_to_dictionary: bool, reverse_row_groups: bool, preserve_order: bool, } @@ -2159,6 +2165,7 @@ mod test { coerce_int96: None, max_predicate_cache_size: None, max_in_list_size: MAX_IN_LIST_SIZE, + enable_rle_to_dictionary: false, reverse_row_groups: false, preserve_order: false, } @@ -2234,6 +2241,11 @@ mod test { self } + fn with_enable_rle_to_dictionary(mut self, enable: bool) -> Self { + self.enable_rle_to_dictionary = enable; + self + } + fn with_metrics(mut self, metrics: ExecutionPlanMetricsSet) -> Self { self.metrics = metrics; self @@ -2343,6 +2355,7 @@ mod test { encryption_factory: None, max_predicate_cache_size: self.max_predicate_cache_size, max_in_list_size: self.max_in_list_size, + enable_rle_to_dictionary: self.enable_rle_to_dictionary, reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: None, virtual_state, @@ -4390,4 +4403,59 @@ mod test { assert_eq!(rows, 5); } } + + async fn collect_batches( + morselizer: &ParquetMorselizer, + file: PartitionedFile, + ) -> Vec { + let mut stream = open_file(morselizer, file).await.unwrap(); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + batches + } + + // Proves the opener passes a promoted binary Dictionary schema to arrow-rs. + #[tokio::test] + async fn test_rle_binary_column_promotion() { + let store = Arc::new(InMemory::new()) as Arc; + let bin_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Binary, + true, + )])); + let values = + Arc::new(arrow::array::BinaryArray::from_vec(vec![b"a", b"b", b"a"])); + let batch = RecordBatch::try_new(Arc::clone(&bin_schema), vec![values]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(true) + .build(); + let bin_size = write_parquet_batches( + Arc::clone(&store), + "binary.parquet", + vec![batch], + Some(props), + ) + .await; + let dict_bin_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)), + true, + )])); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&dict_bin_schema)) + .with_enable_rle_to_dictionary(true) + .build(); + let batches = collect_batches( + &morselizer, + PartitionedFile::new("binary.parquet".to_string(), bin_size as u64), + ) + .await; + assert_eq!( + batches[0].schema().field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)) + ); + } } diff --git a/datafusion/datasource-parquet/src/schema_coercion.rs b/datafusion/datasource-parquet/src/schema_coercion.rs index ce51ab33fd8ff..4a681c4597c20 100644 --- a/datafusion/datasource-parquet/src/schema_coercion.rs +++ b/datafusion/datasource-parquet/src/schema_coercion.rs @@ -17,7 +17,7 @@ //! Arrow-schema coercion utilities used by the Parquet reader to make a //! file schema match the table schema (binary→string, regular→view, -//! INT96→Timestamp). +//! INT96→Timestamp, plain->Dictionary). //! //! These helpers are independent of the [`ParquetFormat`](crate::file_format::ParquetFormat) //! type and several have been re-exported at the crate root for use by @@ -50,32 +50,44 @@ use parquet::schema::types::SchemaDescriptor; pub fn apply_file_schema_type_coercions( table_schema: &Schema, file_schema: &Schema, +) -> Option { + apply_file_schema_type_coercions_with_rle(table_schema, file_schema, false) +} + +/// Like [`apply_file_schema_type_coercions`], but also coerces compatible +/// string/binary file fields to dictionary types already present in the table +/// schema. +pub(crate) fn apply_file_schema_type_coercions_with_rle( + table_schema: &Schema, + file_schema: &Schema, + enable_rle_to_dictionary: bool, ) -> Option { let mut needs_view_transform = false; let mut needs_string_transform = false; let mut needs_nested_transform = false; + let mut needs_dict_transform = false; // Create a mapping of table field names to their data types for fast lookup // and simultaneously check if we need any transformations let table_fields: HashMap<_, _> = table_schema .fields() .iter() - .map(|f| { - let dt = f.data_type(); + .map(|field| { + let data_type = field.data_type(); // Check if we need view type transformation - if matches!(dt, &DataType::Utf8View | &DataType::BinaryView) { + if matches!(data_type, &DataType::Utf8View | &DataType::BinaryView) { needs_view_transform = true; } // Check if we need string type transformation if matches!( - dt, + data_type, &DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View ) { needs_string_transform = true; } // Nested fields can need transformations even when their parent does not. if matches!( - dt, + data_type, DataType::Struct(_) | DataType::List(_) | DataType::LargeList(_) @@ -86,13 +98,22 @@ pub fn apply_file_schema_type_coercions( ) { needs_nested_transform = true; } + if enable_rle_to_dictionary + && matches!(data_type, &DataType::Dictionary(_, _)) + { + needs_dict_transform = true; + } - (f.name(), dt) + (field.name(), data_type) }) .collect(); // Early return if no transformation needed - if !needs_view_transform && !needs_string_transform && !needs_nested_transform { + if !needs_view_transform + && !needs_string_transform + && !needs_nested_transform + && !needs_dict_transform + { return None; } @@ -186,10 +207,15 @@ pub fn apply_file_schema_type_coercions( return field_with_new_type(field, new_type); } } + (DataType::Dictionary(_, _), _) + if enable_rle_to_dictionary + && can_promote_to_dictionary_type(field_type, table_type) => + { + return field_with_new_type(field, (*table_type).clone()); + } _ => {} } } - // If no transformation is needed, keep the original field Arc::clone(field) }) @@ -205,6 +231,170 @@ pub fn apply_file_schema_type_coercions( )) } +// Find the value type that can represent both sides without narrowing offsets +// or crossing string/binary families. +fn common_dictionary_value_type( + field_type: &DataType, + dictionary_value_type: &DataType, +) -> Option { + let field_type = match field_type { + DataType::Dictionary(_, field_value_type) => field_value_type.as_ref(), + _ => field_type, + }; + + match (field_type, dictionary_value_type) { + (DataType::Utf8, DataType::Utf8) => Some(DataType::Utf8), + (DataType::Utf8 | DataType::LargeUtf8, DataType::Utf8 | DataType::LargeUtf8) => { + Some(DataType::LargeUtf8) + } + (DataType::Binary, DataType::Binary) => Some(DataType::Binary), + ( + DataType::Binary | DataType::LargeBinary, + DataType::Binary | DataType::LargeBinary, + ) => Some(DataType::LargeBinary), + _ => None, + } +} + +// Same family (both signed or both unsigned): return the wider member. +// Mixed: return the type with the larger capacity. For the UInt64+Int64 case +// this is UInt64 since arrow-rs supports both signed and unsigned dictionary keys. +// Returns None only for non-integer key types. +fn common_dictionary_key_type(a: &DataType, b: &DataType) -> Option { + fn key_capacity(dt: &DataType) -> Option { + match dt { + DataType::Int8 => Some(1u128 << 7), + DataType::Int16 => Some(1u128 << 15), + DataType::Int32 => Some(1u128 << 31), + DataType::Int64 => Some(1u128 << 63), + DataType::UInt8 => Some(1u128 << 8), + DataType::UInt16 => Some(1u128 << 16), + DataType::UInt32 => Some(1u128 << 32), + DataType::UInt64 => Some(1u128 << 64), + _ => None, + } + } + fn is_signed(dt: &DataType) -> bool { + matches!( + dt, + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) + } + let cap_a = key_capacity(a)?; + let cap_b = key_capacity(b)?; + if is_signed(a) == is_signed(b) { + return Some(if cap_a >= cap_b { a.clone() } else { b.clone() }); + } + let max_cap = cap_a.max(cap_b); + Some(match max_cap { + c if c <= (1u128 << 7) => DataType::Int8, + c if c <= (1u128 << 15) => DataType::Int16, + c if c <= (1u128 << 31) => DataType::Int32, + c if c <= (1u128 << 63) => DataType::Int64, + _ => DataType::UInt64, + }) +} + +fn can_promote_to_dictionary_type( + file_field_type: &DataType, + table_dictionary_type: &DataType, +) -> bool { + let DataType::Dictionary(table_key, table_value) = table_dictionary_type else { + return false; + }; + if common_dictionary_value_type(file_field_type, table_value) + .is_none_or(|common| &common != table_value.as_ref()) + { + return false; + } + if let DataType::Dictionary(file_key, _) = file_field_type { + return common_dictionary_key_type(file_key, table_key) + .is_some_and(|common| &common == table_key.as_ref()); + } + // Plain file field: only allow promotion if the table key is at least Int32 wide. + // A plain column can have any number of distinct values, so narrow keys (Int8, Int16) + // are unsafe — uniform_dict_schemas widens to Int32 when a plain file is present, so + // this guard handles the case where can_promote_to_dictionary_type is called directly. + common_dictionary_key_type(&DataType::Int32, table_key) + .is_some_and(|common| &common == table_key.as_ref()) +} + +/// Normalize per-file schemas so that a column promoted to `Dictionary` in +/// *any* file is promoted to the same `Dictionary` type in *all* files. +/// +/// This lets [`Schema::try_merge`] accept directories that mix dictionary and +/// plain encodings for the same column. +pub(crate) fn uniform_dict_schemas(schemas: Vec) -> Vec { + // First pass: record the dictionary type for every column that is Dictionary in + // at least one schema. + let mut dict_types: HashMap = HashMap::new(); + for schema in &schemas { + for field in schema.fields() { + if matches!(field.data_type(), DataType::Dictionary(_, _)) { + dict_types + .entry(field.name().clone()) + .or_insert_with(|| field.data_type().clone()); + } + } + } + if dict_types.is_empty() { + return schemas; + } + + for schema in &schemas { + for field in schema.fields() { + let Some(dict_type) = dict_types.get_mut(field.name()) else { + continue; + }; + let DataType::Dictionary(key_type, value_type) = dict_type else { + continue; + }; + let key_type = key_type.clone(); + let value_type = value_type.as_ref().clone(); + let new_key = if let DataType::Dictionary(file_key, _) = field.data_type() { + common_dictionary_key_type(&key_type, file_key) + .map(Box::new) + .unwrap_or_else(|| key_type.clone()) + } else { + // Plain field: treat as Int32 key capacity since we don't know how + // many distinct values it has. This prevents a narrow key (e.g. Int8) + // from being selected as the common type when one file uses a narrow + // dictionary and another uses a plain encoding with potentially more + // values than the key can represent. + common_dictionary_key_type(&key_type, &DataType::Int32) + .map(Box::new) + .unwrap_or_else(|| Box::new(DataType::Int32)) + }; + if let Some(common_type) = + common_dictionary_value_type(field.data_type(), &value_type) + { + *dict_type = DataType::Dictionary(new_key, Box::new(common_type)); + } + } + } + + // Only promote fields whose type family matches the dictionary value type, + // so schema normalization does not reinterpret binary bytes as UTF-8. + schemas + .into_iter() + .map(|schema| { + let fields: Vec> = schema + .fields() + .iter() + .map(|field| { + if let Some(dict_type) = dict_types.get(field.name()) + && can_promote_to_dictionary_type(field.data_type(), dict_type) + { + return field_with_new_type(field, dict_type.clone()); + } + Arc::clone(field) + }) + .collect(); + Schema::new_with_metadata(fields, schema.metadata().clone()) + }) + .collect() +} + /// Coerces the file schema's Timestamps to the provided TimeUnit if the /// Parquet schema contains INT96. /// @@ -517,7 +707,11 @@ pub fn transform_schema_to_view(schema: &Schema) -> Schema { Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) } -/// Transform a schema so that any binary types are strings +/// Transform a schema so that any binary types are strings. +/// +/// Also handles `Dictionary(key, binary)` produced by `enable_rle_to_dictionary`, +/// converting the dictionary value type so `binary_as_string` applies consistently +/// regardless of whether the physical encoding triggered dictionary promotion. pub fn transform_binary_to_string(schema: &Schema) -> Schema { let transformed_fields: Vec> = schema .fields @@ -526,6 +720,21 @@ pub fn transform_binary_to_string(schema: &Schema) -> Schema { DataType::Binary => field_with_new_type(field, DataType::Utf8), DataType::LargeBinary => field_with_new_type(field, DataType::LargeUtf8), DataType::BinaryView => field_with_new_type(field, DataType::Utf8View), + DataType::Dictionary(key_type, value_type) => match value_type.as_ref() { + DataType::Binary => field_with_new_type( + field, + DataType::Dictionary(key_type.clone(), Box::new(DataType::Utf8)), + ), + DataType::LargeBinary => field_with_new_type( + field, + DataType::Dictionary(key_type.clone(), Box::new(DataType::LargeUtf8)), + ), + DataType::BinaryView => field_with_new_type( + field, + DataType::Dictionary(key_type.clone(), Box::new(DataType::Utf8View)), + ), + _ => Arc::clone(field), + }, _ => Arc::clone(field), }) .collect(); @@ -1186,4 +1395,389 @@ mod tests { ), } } + + fn dict(value_type: DataType) -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(value_type)) + } + + fn one_field_schema(data_type: DataType) -> Schema { + Schema::new(vec![Field::new("col", data_type, true)]) + } + + #[test] + fn uniform_dict_schemas_respects_value_type_families() { + // String/binary dictionary promotion must preserve the concrete value type + // family so schema merging does not silently reinterpret bytes as UTF-8. + let cases = vec![ + ( + "utf8", + dict(DataType::Utf8), + DataType::Utf8, + dict(DataType::Utf8), + dict(DataType::Utf8), + ), + ( + "large_utf8", + dict(DataType::LargeUtf8), + DataType::LargeUtf8, + dict(DataType::LargeUtf8), + dict(DataType::LargeUtf8), + ), + ( + "utf8_large_utf8", + dict(DataType::Utf8), + DataType::LargeUtf8, + dict(DataType::LargeUtf8), + dict(DataType::LargeUtf8), + ), + ( + "binary", + dict(DataType::Binary), + DataType::Binary, + dict(DataType::Binary), + dict(DataType::Binary), + ), + ( + "large_binary", + dict(DataType::LargeBinary), + DataType::LargeBinary, + dict(DataType::LargeBinary), + dict(DataType::LargeBinary), + ), + ( + "binary_large_binary", + dict(DataType::Binary), + DataType::LargeBinary, + dict(DataType::LargeBinary), + dict(DataType::LargeBinary), + ), + ( + "binary_not_utf8", + dict(DataType::Utf8), + DataType::Binary, + dict(DataType::Utf8), + DataType::Binary, + ), + ( + "utf8_not_binary", + dict(DataType::Binary), + DataType::Utf8, + dict(DataType::Binary), + DataType::Utf8, + ), + ]; + + for (name, dict_type, plain_type, expected_dict_type, expected_plain_type) in + cases + { + let result = uniform_dict_schemas(vec![ + one_field_schema(dict_type.clone()), + one_field_schema(plain_type), + ]); + + assert_eq!( + result[0].field(0).data_type(), + &expected_dict_type, + "{name}" + ); + assert_eq!( + result[1].field(0).data_type(), + &expected_plain_type, + "{name}" + ); + } + } + + #[test] + fn rle_schema_coercion_respects_dictionary_value_type() { + // Scan-time schema coercion can only use dictionary types already chosen by + // schema inference, and must leave incompatible file fields unchanged. + let cases = vec![ + ( + "utf8", + dict(DataType::Utf8), + DataType::Utf8, + dict(DataType::Utf8), + ), + ( + "large_utf8", + dict(DataType::LargeUtf8), + DataType::LargeUtf8, + dict(DataType::LargeUtf8), + ), + ( + "large_utf8_to_utf8_dict_not_safe", + dict(DataType::Utf8), + DataType::LargeUtf8, + DataType::LargeUtf8, + ), + ( + "utf8_to_large_utf8_dict", + dict(DataType::LargeUtf8), + DataType::Utf8, + dict(DataType::LargeUtf8), + ), + ( + "binary", + dict(DataType::Binary), + DataType::Binary, + dict(DataType::Binary), + ), + ( + "large_binary", + dict(DataType::LargeBinary), + DataType::LargeBinary, + dict(DataType::LargeBinary), + ), + ( + "large_binary_to_binary_dict_not_safe", + dict(DataType::Binary), + DataType::LargeBinary, + DataType::LargeBinary, + ), + ( + "binary_to_large_binary_dict", + dict(DataType::LargeBinary), + DataType::Binary, + dict(DataType::LargeBinary), + ), + ( + "binary_not_utf8", + dict(DataType::Utf8), + DataType::Binary, + DataType::Binary, + ), + ( + "utf8_not_binary", + dict(DataType::Binary), + DataType::Utf8, + DataType::Utf8, + ), + ( + "binary_not_int64_dict", + dict(DataType::Int64), + DataType::Binary, + DataType::Binary, + ), + ]; + + for (name, table_type, file_type, expected_type) in cases { + let table_schema = one_field_schema(table_type); + let file_schema = one_field_schema(file_type); + + let output_type = apply_file_schema_type_coercions_with_rle( + &table_schema, + &file_schema, + true, + ) + .as_ref() + .map(|s| s.field(0).data_type().clone()) + .unwrap_or_else(|| file_schema.field(0).data_type().clone()); + + assert_eq!(output_type, expected_type, "{name}"); + } + + let table_schema = Schema::new(vec![ + Field::new("dict_col", dict(DataType::Utf8), true), + Field::new("string_col", DataType::Utf8, true), + ]); + let file_schema = Schema::new(vec![ + Field::new("dict_col", DataType::Utf8, true), + Field::new("string_col", DataType::Binary, true), + ]); + let result = + apply_file_schema_type_coercions_with_rle(&table_schema, &file_schema, false) + .unwrap(); + + assert_eq!(result.field(0).data_type(), &DataType::Utf8); + assert_eq!(result.field(1).data_type(), &DataType::Utf8); + } + + fn dk(key: DataType) -> DataType { + DataType::Dictionary(Box::new(key), Box::new(DataType::Utf8)) + } + + #[test] + fn uniform_dict_schemas_key_type_widening() { + // (name, file_a, file_b, expected_a, expected_b) + let cases = [ + ( + "signed wider wins", + dk(DataType::Int8), + dk(DataType::Int32), + dk(DataType::Int32), + dk(DataType::Int32), + ), + ( + "signed wider wins reversed", + dk(DataType::Int32), + dk(DataType::Int8), + dk(DataType::Int32), + dk(DataType::Int32), + ), + ( + "same unsigned stays", + dk(DataType::UInt8), + dk(DataType::UInt8), + dk(DataType::UInt8), + dk(DataType::UInt8), + ), + ( + "uint8+int32→int32", + dk(DataType::UInt8), + dk(DataType::Int32), + dk(DataType::Int32), + dk(DataType::Int32), + ), + ( + "uint8+int8→int16", + dk(DataType::UInt8), + dk(DataType::Int8), + dk(DataType::Int16), + dk(DataType::Int16), + ), + ( + "uint64+int64→uint64", + dk(DataType::UInt64), + dk(DataType::Int64), + dk(DataType::UInt64), + dk(DataType::UInt64), + ), + ]; + for (name, a, b, exp_a, exp_b) in cases { + let result = + uniform_dict_schemas(vec![one_field_schema(a), one_field_schema(b)]); + assert_eq!(result[0].field(0).data_type(), &exp_a, "{name}"); + assert_eq!(result[1].field(0).data_type(), &exp_b, "{name}"); + } + } + + #[test] + fn uniform_dict_schemas_plain_field_widens_key_to_int32() { + // A plain field mixed with a narrow-key dictionary must widen the key to at + // least Int32 so that the plain file (with unknown cardinality) can be + // represented safely. + let cases = [ + ( + "int8_dict + plain → int32", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + DataType::Utf8, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ), + ( + "int16_dict + plain → int32", + DataType::Dictionary(Box::new(DataType::Int16), Box::new(DataType::Utf8)), + DataType::Utf8, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ), + ( + "int32_dict + plain stays int32", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Utf8, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ), + ( + "int64_dict + plain stays int64", + DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)), + DataType::Utf8, + DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)), + DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8)), + ), + ]; + for (name, dict_type, plain_type, exp_dict, exp_plain) in cases { + let result = uniform_dict_schemas(vec![ + one_field_schema(dict_type), + one_field_schema(plain_type), + ]); + assert_eq!(result[0].field(0).data_type(), &exp_dict, "{name}"); + assert_eq!(result[1].field(0).data_type(), &exp_plain, "{name}"); + } + } + + #[test] + fn rle_schema_coercion_rejects_key_narrowing() { + let coerce = |table: DataType, file: DataType| { + let t = one_field_schema(table); + let f = one_field_schema(file.clone()); + apply_file_schema_type_coercions_with_rle(&t, &f, true) + .as_ref() + .map(|s| s.field(0).data_type().clone()) + .unwrap_or(file) + }; + // Dict(Int32) file must not be narrowed to Dict(Int8) at scan time. + assert_eq!( + coerce(dk(DataType::Int8), dk(DataType::Int32)), + dk(DataType::Int32) + ); + // Plain Utf8 file must NOT be promoted to a narrow-key dict: we cannot know how + // many distinct values the file has, so Int8 capacity is unsafe. + assert_eq!(coerce(dk(DataType::Int8), DataType::Utf8), DataType::Utf8); + } + + #[test] + fn transform_binary_to_string_handles_dictionary_value_types() { + let schema = Schema::new(vec![ + Field::new( + "a", + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Binary), + ), + true, + ), + Field::new( + "b", + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::LargeBinary), + ), + true, + ), + Field::new( + "c", + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::BinaryView), + ), + true, + ), + Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + ), + Field::new("e", DataType::Binary, true), + ]); + + let result = transform_binary_to_string(&schema); + + assert_eq!( + result.field(0).data_type(), + &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ); + assert_eq!( + result.field(1).data_type(), + &DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::LargeUtf8) + ), + ); + assert_eq!( + result.field(2).data_type(), + &DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8View) + ), + ); + // Dictionary(_, Utf8) is left unchanged + assert_eq!( + result.field(3).data_type(), + &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ); + // Plain Binary is converted as before + assert_eq!(result.field(4).data_type(), &DataType::Utf8); + } } diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 4872db9fd3329..8fc6aa24c0c0f 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -660,6 +660,10 @@ impl FileSource for ParquetSource { encryption_factory: self.get_encryption_factory_with_config(), max_predicate_cache_size: self.max_predicate_cache_size(), max_in_list_size: self.max_in_list_size(), + enable_rle_to_dictionary: self + .table_parquet_options + .global + .enable_rle_to_dictionary, reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: self.sort_order_for_reorder.clone(), virtual_state, diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 712212f6b6ae5..a4a49d99ff180 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -620,6 +620,8 @@ message ParquetOptions { uint64 max_in_list_size = 38; + bool enable_rle_to_dictionary = 39; + string created_by = 16; oneof coerce_int96_opt { diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 92506ad92bad0..ffda86fd11e41 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1180,6 +1180,7 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { } }).transpose()?, content_defined_chunking: value.content_defined_chunking.map(ParquetCdcOptions::try_from).transpose()?.unwrap_or_default(), + enable_rle_to_dictionary: value.enable_rle_to_dictionary, }) } } diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 2e94368bfd01e..7a30a41873a8e 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -3999,7 +3999,7 @@ impl serde::Serialize for ExplainAnalyzeCategoriesNode { if !self.only.is_empty() { let v = self.only.iter().copied().map(|v| { MetricCategory::try_from(v) - .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", v))) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {v}"))) }).collect::, _>>()?; struct_ser.serialize_field("only", &v)?; } @@ -6432,6 +6432,9 @@ impl serde::Serialize for ParquetOptions { if self.max_in_list_size != 0 { len += 1; } + if self.enable_rle_to_dictionary != false { + len += 1; + } if !self.created_by.is_empty() { len += 1; } @@ -6557,6 +6560,9 @@ impl serde::Serialize for ParquetOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("maxInListSize", ToString::to_string(&self.max_in_list_size).as_str())?; } + if self.enable_rle_to_dictionary != false { + struct_ser.serialize_field("enableRleToDictionary", &self.enable_rle_to_dictionary)?; + } if !self.created_by.is_empty() { struct_ser.serialize_field("createdBy", &self.created_by)?; } @@ -6717,6 +6723,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "maxRowGroupSize", "max_in_list_size", "maxInListSize", + "enable_rle_to_dictionary", + "enableRleToDictionary", "created_by", "createdBy", "content_defined_chunking", @@ -6770,6 +6778,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { DataPageRowCountLimit, MaxRowGroupSize, MaxInListSize, + EnableRleToDictionary, CreatedBy, ContentDefinedChunking, MetadataSizeHint, @@ -6827,6 +6836,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit), "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize), "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), + "enableRleToDictionary" | "enable_rle_to_dictionary" => Ok(GeneratedField::EnableRleToDictionary), "createdBy" | "created_by" => Ok(GeneratedField::CreatedBy), "contentDefinedChunking" | "content_defined_chunking" => Ok(GeneratedField::ContentDefinedChunking), "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), @@ -6882,6 +6892,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut data_page_row_count_limit__ = None; let mut max_row_group_size__ = None; let mut max_in_list_size__ = None; + let mut enable_rle_to_dictionary__ = None; let mut created_by__ = None; let mut content_defined_chunking__ = None; let mut metadata_size_hint_opt__ = None; @@ -7041,6 +7052,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::EnableRleToDictionary => { + if enable_rle_to_dictionary__.is_some() { + return Err(serde::de::Error::duplicate_field("enableRleToDictionary")); + } + enable_rle_to_dictionary__ = Some(map_.next_value()?); + } GeneratedField::CreatedBy => { if created_by__.is_some() { return Err(serde::de::Error::duplicate_field("createdBy")); @@ -7155,6 +7172,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(), max_row_group_size: max_row_group_size__.unwrap_or_default(), max_in_list_size: max_in_list_size__.unwrap_or_default(), + enable_rle_to_dictionary: enable_rle_to_dictionary__.unwrap_or_default(), created_by: created_by__.unwrap_or_default(), content_defined_chunking: content_defined_chunking__, metadata_size_hint_opt: metadata_size_hint_opt__, diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index c0e79aec6d873..303e84f8d6aac 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -866,6 +866,8 @@ pub struct ParquetOptions { pub max_row_group_size: u64, #[prost(uint64, tag = "38")] pub max_in_list_size: u64, + #[prost(bool, tag = "39")] + pub enable_rle_to_dictionary: bool, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 7f95e03f41db4..8853b7cdc2213 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -940,6 +940,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { max_predicate_cache_size_opt: value.max_predicate_cache_size.map(|v| protobuf::parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(v as u64)), max_row_group_bytes_opt: value.max_row_group_bytes.map(|v| protobuf::parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(v.get() as u64)), content_defined_chunking: Some((&value.content_defined_chunking).into()), + enable_rle_to_dictionary: value.enable_rle_to_dictionary, }) } } diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index b2340a09678ac..a5a282c630688 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -385,6 +385,7 @@ impl TryFrom<&ParquetOptionsProto> for ParquetOptions { compression.clone() } }), + enable_rle_to_dictionary: proto.enable_rle_to_dictionary, dictionary_enabled: proto.dictionary_enabled_opt.as_ref().map(|opt| { match opt { parquet_options::DictionaryEnabledOpt::DictionaryEnabled( diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index c0e79aec6d873..303e84f8d6aac 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -866,6 +866,8 @@ pub struct ParquetOptions { pub max_row_group_size: u64, #[prost(uint64, tag = "38")] pub max_in_list_size: u64, + #[prost(bool, tag = "39")] + pub enable_rle_to_dictionary: bool, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] diff --git a/datafusion/proto/src/logical_plan/file_formats.rs b/datafusion/proto/src/logical_plan/file_formats.rs index 7c7a8ef639457..b12db665fc107 100644 --- a/datafusion/proto/src/logical_plan/file_formats.rs +++ b/datafusion/proto/src/logical_plan/file_formats.rs @@ -334,6 +334,33 @@ mod parquet { ParquetOptions::default().writer_version ); } + + #[test] + fn enable_rle_to_dictionary_round_trips_through_codec() { + use datafusion_common::config::TableParquetOptions; + let mut options = TableParquetOptions::default(); + options.global.enable_rle_to_dictionary = true; + let original: Arc = Arc::new(ParquetFormatFactory { + options: Some(options), + }); + + let mut buf = Vec::new(); + ParquetLogicalExtensionCodec + .try_encode_file_format(&mut buf, Arc::clone(&original)) + .expect("encode parquet options"); + + let decoded = ParquetLogicalExtensionCodec + .try_decode_file_format(&buf, &TaskContext::default()) + .expect("decode parquet options"); + let decoded_options = decoded + .downcast_ref::() + .expect("parquet format factory") + .options + .as_ref() + .expect("parquet options"); + + assert!(decoded_options.global.enable_rle_to_dictionary); + } } } #[cfg(feature = "parquet")] diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b270eba99d7b0..0685acbf46095 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -251,6 +251,7 @@ datafusion.execution.parquet.data_pagesize_limit 1048576 datafusion.execution.parquet.dictionary_enabled true datafusion.execution.parquet.dictionary_page_size_limit 1048576 datafusion.execution.parquet.enable_page_index true +datafusion.execution.parquet.enable_rle_to_dictionary false datafusion.execution.parquet.encoding NULL datafusion.execution.parquet.force_filter_selections false datafusion.execution.parquet.max_in_list_size 20 @@ -412,6 +413,7 @@ datafusion.execution.parquet.data_pagesize_limit 1048576 (writing) Sets best eff datafusion.execution.parquet.dictionary_enabled true (writing) Sets if dictionary encoding is enabled. If NULL, uses default parquet writer setting datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets best effort maximum dictionary page size, in bytes datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. +datafusion.execution.parquet.enable_rle_to_dictionary false (reading) If true, top-level string and binary Parquet columns with dictionary pages are inferred and scanned as `Dictionary` / `Dictionary` instead of their plain value type. This applies only when DataFusion infers the table schema. Tables with a user-supplied schema are not promoted because the Parquet footer is not read at DDL time, so dictionary pages cannot be detected per column. See datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting 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 lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. diff --git a/datafusion/sqllogictest/test_files/parquet_rle_to_dictionary.slt b/datafusion/sqllogictest/test_files/parquet_rle_to_dictionary.slt new file mode 100644 index 0000000000000..a303f5eaec5dc --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_rle_to_dictionary.slt @@ -0,0 +1,355 @@ +# 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. + +# Tests for datafusion.execution.parquet.enable_rle_to_dictionary. The flag +# applies only when DataFusion infers the table schema; explicit schemas are +# not promoted. + +# Write with dictionary_enabled=true to guarantee RLE_DICTIONARY encoding. +query I +COPY ( + SELECT column1 AS product_category, column2 AS status + FROM (VALUES + ('electronics', 'active'), + ('clothing', 'active'), + ('electronics', 'inactive'), + ('books', 'inactive'), + ('furniture', 'pending'), + ('books', 'active') + ) +) +TO 'test_files/scratch/parquet_rle_to_dictionary/products.parquet' +STORED AS PARQUET +OPTIONS ('format.dictionary_enabled' true); +---- +6 + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = false; + +statement ok +CREATE EXTERNAL TABLE products_utf8 +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/products.parquet'; + +query TT +SELECT DISTINCT arrow_typeof(product_category), arrow_typeof(status) FROM products_utf8; +---- +Utf8View Utf8View + +query TI rowsort +SELECT product_category, COUNT(*) FROM products_utf8 GROUP BY product_category; +---- +books 2 +clothing 1 +electronics 2 +furniture 1 + +statement ok +DROP TABLE products_utf8; + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = true; + +statement ok +CREATE EXTERNAL TABLE products_dict +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/products.parquet'; + +query TT +SELECT DISTINCT arrow_typeof(product_category), arrow_typeof(status) FROM products_dict; +---- +Dictionary(Int32, Utf8) Dictionary(Int32, Utf8) + +query TI rowsort +SELECT product_category, COUNT(*) FROM products_dict GROUP BY product_category; +---- +books 2 +clothing 1 +electronics 2 +furniture 1 + +# Predicates and aggregates must work against dictionary scan output. +query TT rowsort +SELECT product_category, status FROM products_dict WHERE status = 'active'; +---- +books active +clothing active +electronics active + +query TI rowsort +SELECT product_category, COUNT(*) FROM products_dict WHERE status != 'inactive' GROUP BY product_category; +---- +books 1 +clothing 1 +electronics 1 +furniture 1 + +query TT +SELECT DISTINCT arrow_typeof(product_category), arrow_typeof(product_count) +FROM ( + SELECT product_category, COUNT(*) AS product_count + FROM products_dict + WHERE status != 'inactive' + GROUP BY product_category +); +---- +Dictionary(Int32, Utf8) Int64 + +statement ok +DROP TABLE products_dict; + +# Per-column writer options: only columns with dictionary pages are promoted. + +statement ok +COPY ( + SELECT column1 AS env, column2 AS region, column3 AS build + FROM (VALUES + ('prod', 'us-east', 'debug'), + ('staging', 'us-west', 'release'), + ('prod', 'us-east', 'debug') + ) +) +TO 'test_files/scratch/parquet_rle_to_dictionary/selective.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.dictionary_enabled' false, + 'format.dictionary_enabled::env' true, + 'format.dictionary_enabled::region' true +); + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = false; + +statement ok +CREATE EXTERNAL TABLE selective_utf8 +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/selective.parquet'; + +query TTT +SELECT DISTINCT arrow_typeof(env), arrow_typeof(region), arrow_typeof(build) +FROM selective_utf8; +---- +Utf8View Utf8View Utf8View + +statement ok +DROP TABLE selective_utf8; + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = true; + +statement ok +CREATE EXTERNAL TABLE selective_dict +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/selective.parquet'; + +query TTT +SELECT DISTINCT arrow_typeof(env), arrow_typeof(region), arrow_typeof(build) +FROM selective_dict; +---- +Dictionary(Int32, Utf8) Dictionary(Int32, Utf8) Utf8View + +statement ok +DROP TABLE selective_dict; + +# Cross-file mixed encoding: one file has dictionary pages and one is plain. +# Schema inference must merge both as Dictionary(Int32, Utf8). + +statement ok +COPY (SELECT column1 AS category FROM (VALUES ('electronics'), ('electronics'))) +TO 'test_files/scratch/parquet_rle_to_dictionary/mixed/rle.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' true); + +statement ok +COPY (SELECT column1 AS category FROM (VALUES ('books'), ('books'))) +TO 'test_files/scratch/parquet_rle_to_dictionary/mixed/plain.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' false); + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = true; + +statement ok +CREATE EXTERNAL TABLE mixed_encoding +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/mixed/'; + +query TI +SELECT arrow_typeof(category), COUNT(*) FROM mixed_encoding GROUP BY arrow_typeof(category); +---- +Dictionary(Int32, Utf8) 4 + +statement ok +DROP TABLE mixed_encoding; + +# Mixed encoding with high cardinality: the plain file contains more than 128 +# distinct values so a narrow (Int8) dictionary key would overflow. After the +# fix, uniform_dict_schemas widens the shared key to at least Int32. +# Both files must be scannable and return the correct row count. + +statement ok +COPY ( + SELECT v AS category + FROM ( + VALUES + ('v000'),('v001'),('v002'),('v003'),('v004'),('v005'),('v006'),('v007'), + ('v008'),('v009'),('v010'),('v011'),('v012'),('v013'),('v014'),('v015'), + ('v016'),('v017'),('v018'),('v019'),('v020'),('v021'),('v022'),('v023'), + ('v024'),('v025'),('v026'),('v027'),('v028'),('v029'),('v030'),('v031'), + ('v032'),('v033'),('v034'),('v035'),('v036'),('v037'),('v038'),('v039'), + ('v040'),('v041'),('v042'),('v043'),('v044'),('v045'),('v046'),('v047'), + ('v048'),('v049'),('v050'),('v051'),('v052'),('v053'),('v054'),('v055'), + ('v056'),('v057'),('v058'),('v059'),('v060'),('v061'),('v062'),('v063'), + ('v064'),('v065'),('v066'),('v067'),('v068'),('v069'),('v070'),('v071'), + ('v072'),('v073'),('v074'),('v075'),('v076'),('v077'),('v078'),('v079'), + ('v080'),('v081'),('v082'),('v083'),('v084'),('v085'),('v086'),('v087'), + ('v088'),('v089'),('v090'),('v091'),('v092'),('v093'),('v094'),('v095'), + ('v096'),('v097'),('v098'),('v099'),('v100'),('v101'),('v102'),('v103'), + ('v104'),('v105'),('v106'),('v107'),('v108'),('v109'),('v110'),('v111'), + ('v112'),('v113'),('v114'),('v115'),('v116'),('v117'),('v118'),('v119'), + ('v120'),('v121'),('v122'),('v123'),('v124'),('v125'),('v126'),('v127'), + ('v128'),('v129') + ) AS t(v) +) +TO 'test_files/scratch/parquet_rle_to_dictionary/high_card/plain.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' false); + +# Write rle.parquet with an explicit Dictionary(Int8, Utf8) Arrow schema so +# that the file carries Int8 dictionary keys in its embedded Arrow schema +# metadata. This exercises the real regression path: without the key-widening +# fix, uniform_dict_schemas would select Int8 as the merged key type, which +# cannot represent the 130 distinct values in the plain file. +statement ok +COPY (SELECT arrow_cast(column1, 'Dictionary(Int8, Utf8)') AS category FROM (VALUES ('rle_a'), ('rle_b'), ('rle_a'))) +TO 'test_files/scratch/parquet_rle_to_dictionary/high_card/rle.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' true); + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = true; + +# Verify that rle.parquet alone reads back as Dictionary(Int8, Utf8) — +# confirming the Arrow schema metadata round-trips and was not widened on write. +statement ok +CREATE EXTERNAL TABLE rle_int8_only +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/high_card/rle.parquet'; + +query T +SELECT DISTINCT arrow_typeof(category) FROM rle_int8_only; +---- +Dictionary(Int8, Utf8) + +statement ok +DROP TABLE rle_int8_only; + +statement ok +CREATE EXTERNAL TABLE high_card_mixed +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/high_card/'; + +# Both files are scanned as Dictionary(Int32, Utf8) — uniform_dict_schemas +# widens the Int8 key from rle.parquet to Int32 to avoid overflow for the +# 130 distinct values in the plain file. +query TI +SELECT arrow_typeof(category), COUNT(*) FROM high_card_mixed GROUP BY arrow_typeof(category); +---- +Dictionary(Int32, Utf8) 133 + +statement ok +DROP TABLE high_card_mixed; + +# Explicit-schema tables are not promoted. + +statement ok +CREATE EXTERNAL TABLE explicit_schema ( + product_category VARCHAR, + status VARCHAR +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/products.parquet'; + +query TT +SELECT DISTINCT arrow_typeof(product_category), arrow_typeof(status) FROM explicit_schema; +---- +Utf8View Utf8View + +# Predicates must still compile against the explicit plain schema. +query TT rowsort +SELECT product_category, status FROM explicit_schema WHERE status = 'active'; +---- +books active +clothing active +electronics active + +statement ok +DROP TABLE explicit_schema; + +# Family-compatibility safety: a plain Binary field must not be promoted to +# Dictionary(Int32, Utf8). Schema inference should reject this mixed dataset. + +statement ok +COPY (SELECT 'hello' AS payload FROM (VALUES (1), (2))) +TO 'test_files/scratch/parquet_rle_to_dictionary/compat/utf8_rle.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' true); + +statement ok +COPY (SELECT arrow_cast(X'6865', 'Binary') AS payload FROM (VALUES (1), (2))) +TO 'test_files/scratch/parquet_rle_to_dictionary/compat/binary_plain.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' false); + +# Registering incompatible Dict Utf8 and plain Binary files must fail. +statement error Arrow error: Schema error: Fail to merge schema field 'payload' because the from data_type = Dictionary\(Int32, Utf8\) does not equal Binary +CREATE EXTERNAL TABLE mixed_str_binary +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/compat/'; + +# binary_as_string + enable_rle_to_dictionary: Dictionary(_, Binary) value type +# must be promoted to Utf8 so the binary_as_string contract holds regardless of +# physical encoding. + +query I +COPY (SELECT arrow_cast(X'6865', 'Binary') AS payload FROM (VALUES (1), (2))) +TO 'test_files/scratch/parquet_rle_to_dictionary/binary_rle.parquet' +STORED AS PARQUET OPTIONS ('format.dictionary_enabled' true); +---- +2 + +statement ok +set datafusion.execution.parquet.enable_rle_to_dictionary = true; + +statement ok +set datafusion.execution.parquet.binary_as_string = true; + +statement ok +CREATE EXTERNAL TABLE binary_rle_as_string +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_rle_to_dictionary/binary_rle.parquet'; + +# Dictionary(Int32, Binary) must be coerced to Dictionary(Int32, Utf8) when +# binary_as_string is true, not left as Dictionary(Int32, Binary). +query T +SELECT DISTINCT arrow_typeof(payload) FROM binary_rle_as_string; +---- +Dictionary(Int32, Utf8) + +statement ok +DROP TABLE binary_rle_as_string; + +statement ok +reset datafusion.execution.parquet.binary_as_string; + +statement ok +reset datafusion.execution.parquet.enable_rle_to_dictionary; 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 7658ce60eddf2..6430d172d67cf 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -168,6 +168,42 @@ let mut config = SessionConfig::new(); config.options_mut().execution.enable_nlj_coordinated_fallback = false; ``` +### New field `enable_rle_to_dictionary` added to `ParquetOptions` + +`ParquetOptions` gained a new field `enable_rle_to_dictionary: bool` (default +`false`). It controls whether top-level string and binary Parquet columns with +dictionary pages are read as `Dictionary` / `Dictionary` instead of their plain type. + +`ParquetOptions` is a public struct without `#[non_exhaustive]`, so downstream +crates that construct it with an exhaustive struct literal will fail to compile +with `E0063: missing field 'enable_rle_to_dictionary'`. + +**Migration guide:** + +Add `enable_rle_to_dictionary: false` to any exhaustive `ParquetOptions { .. }` +literal, or use `..Default::default()` to future-proof against further additions: + +```rust,ignore +// Before +ParquetOptions { + enable_page_index: true, + // ... every other field ... +} + +// After: set it explicitly +ParquetOptions { + enable_page_index: true, + enable_rle_to_dictionary: false, + // ... every other field ... +} + +// After: or let remaining fields come from Default +ParquetOptions { + enable_page_index: true, + ..Default::default() +} +``` + ### `datafusion.optimizer.use_statistics_registry` is deprecated and ignored The `datafusion.optimizer.use_statistics_registry` config flag is deprecated and diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 0085d4ac7c1fa..5915c59daa1d8 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -94,6 +94,7 @@ The following configuration settings are available: | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | 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_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 lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and 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.enable_rle_to_dictionary | false | (reading) If true, top-level string and binary Parquet columns with dictionary pages are inferred and scanned as `Dictionary` / `Dictionary` instead of their plain value type. This applies only when DataFusion infers the table schema. Tables with a user-supplied schema are not promoted because the Parquet footer is not read at DDL time, so dictionary pages cannot be detected per column. See | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" |