kumarUjjawal commented on code in PR #24227:
URL: https://github.com/apache/datafusion/pull/24227#discussion_r3940452913
##########
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:
the dictionary/dictionary case is now widened, but I think the plain-file
issue remains. uniform_dict_schemas retains the existing dictionary key for a
plain field, and can_promote_to_dictionary_type returns true without
considering key capacity. The new test explicitly expects plain Utf8 to become
Dictionary(Int8, Utf8).
With one Dictionary(Int8, Utf8) file and one plain Utf8 file containing at
least 129 distinct values, the merged table schema remains Dictionary(Int8,
Utf8), so Arrow cannot represent the 129th key. Could plain fields participate
as the feature’s default Int32 key when selecting the common type, with an
execution-level mixed-file regression test containing more than 128 distinct
values?
--
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]
