Rich-T-kid commented on code in PR #24227:
URL: https://github.com/apache/datafusion/pull/24227#discussion_r3946482880
##########
datafusion/datasource-parquet/src/schema_coercion.rs:
##########
@@ -135,6 +157,114 @@ pub fn apply_file_schema_type_coercions(
))
}
+fn dictionary_value_type(data_type: &DataType) -> Option<&DataType> {
+ match data_type {
+ DataType::Dictionary(_, value_type) => Some(value_type.as_ref()),
+ _ => None,
+ }
+}
+
+// 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<DataType> {
+ 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,
+ }
+}
+
+/// Allows safe widening into the table dictionary type.
+/// - `Utf8` to `Dictionary(Int32, LargeUtf8)`: allowed
+/// - `LargeBinary` to `Dictionary(Int32, Binary)`: rejected
+fn can_promote_to_dictionary_type(
+ file_field_type: &DataType,
+ table_dictionary_type: &DataType,
+) -> bool {
+
dictionary_value_type(table_dictionary_type).is_some_and(|dictionary_value_type|
{
+ common_dictionary_value_type(file_field_type, dictionary_value_type)
+ .is_some_and(|common_type| &common_type == dictionary_value_type)
+ })
+}
+
+/// 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<Schema>) -> Vec<Schema> {
+ // First pass: record the dictionary type for every column that is
Dictionary in
+ // at least one schema.
+ let mut dict_types: HashMap<String, DataType> = 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;
+ }
+
+ // Widen the recorded dictionary value type before promoting plain fields.
+ 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 {
Review Comment:
https://github.com/apache/datafusion/pull/24227/changes/ffd6230868958af8ee2840bcb310a62dbabbcc12;
`uniform_dict_schemas` now treats plain fields as having Int32 key capacity
when computing the common key type (so Dict(Int8) + plain → Dict(Int32)), and
`can_promote_to_dictionary_type` now rejects promotion of plain files to
narrow-key dicts. Added `uniform_dict_schemas_plain_field_widens_key_to_int32`
unit test and an execution-level test with 130 distinct values in the plain
file.
--
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]
