Support Partitioning Data by Dictionary Encoded String Array Types

[devinjdangelo]

Which issue does this PR close?

Closes #7891

Rationale for this change

The initial implementation of hive style partition demux code only supported plain UTF8 arrays. This PR adds support for dictionary encoded UTF8 arrays.

What changes are included in this PR?

• Extracts string values from dictionary encoded UTF8 arrays to determine hive partitions
• Adds additional cases to schema semantic equivalence checking: Dictionary(_, dtype) is now considered semantically equivalent to dtype.

Are these changes tested?

Yes using dictionary encoded values from parquet-testing submodule

Are there any user-facing changes?

No, just (hopefully) more robust support for inserting inferred schema string data as a partition column.

[devinjdangelo]

It would be nice to not have to enumerate every possible key_type here. It would be nice if there was a way to cast any Dictionary(_, V) -> V. Or at least iterate over any Dictionary(_, V) to pull out the V values in order without caring what type the key is.

[suremarc]

Hey Devin, there is a downcast_dictionary_array! macro that I encountered while attempting to fix this issue. (I was going to submit a PR but was figuring out if/how to add tests!) It lets you write code once that works for all key types, if I am not mistaken. I saw it used elsewhere in DataFusion code: https://github.com/apache/arrow-datafusion/blob/9fde5c4282fd9f0e3332fb40998bf1562c17fcda/datafusion/common/src/hash_utils.rs#L326-L329

You can see the code I used here:
https://github.com/polygon-io/arrow-datafusion/blob/8c955891a6e5b6002faa07e54a67a6ec93347ade/datafusion/core/src/datasource/file_format/write/demux.rs#L316-L328

            DataType::Dictionary(_, value_type)
                if value_type.as_ref() == &DataType::Utf8 =>
            {
                downcast_dictionary_array!(
                    col_array =>  {
                        let array = col_array.downcast_dict::<StringArray>().unwrap();
                        for i in 0..rb.num_rows() {
                            partition_values.push(array.value(i));
                        }
                    },
                    _ => unreachable!(),
                )
            }

[devinjdangelo]

Thanks @suremarc for the pointer to that macro! I updated this PR and it seems to work as expected!

I also added a new test based on a dictionary encoded parquet file from one of the testing submodules. It is working without error, but the values look a bit off. Not sure at this point if related to the downcasting or something else.

[devinjdangelo]

Traditionally null values are sent to __HIVE_DEFAULT_PARTITION__ imo it is also a reasonable choice to throw an error instead. If the user desired the traditional behavior, they could do COALESCE(part_col, '__HIVE_DEFAULT_PARTITION__') as part_col

https://cwiki.apache.org/confluence/display/Hive/Tutorial#Tutorial-Dynamic-PartitionInsert

[devinjdangelo]

date_string_col and string_col look off. I'm not sure if that is due to how the values are being downcast or something else.

[devinjdangelo]

datafusion-cli shows the values slightly different but also odd:

select * from './parquet-testing/data/alltypes_dictionary.parquet';
+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
| id | bool_col | tinyint_col | smallint_col | int_col | bigint_col | float_col | double_col | date_string_col  | string_col | timestamp_col       |
+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
| 0  | true     | 0           | 0            | 0       | 0          | 0.0       | 0.0        | 30312f30312f3039 | 30         | 2009-01-01T00:00:00 |
| 1  | false    | 1           | 1            | 1       | 10         | 1.1       | 10.1       | 30312f30312f3039 | 31         | 2009-01-01T00:01:00 |
+----+----------+-------------+--------------+---------+------------+-----------+------------+------------------+------------+---------------------+
2 rows in set. Query took 0.001 seconds.

[alamb]

This might be an easier way to create a data with dictionary encoded data:

❯ create table test as values (arrow_cast('foo', 'Dictionary(Int32, Utf8)')), (arrow_cast('bar', 'Dictionary(Int32, Utf8)'));
0 rows in set. Query took 0.002 seconds.

❯ select arrow_typeof(column1) from test;
+----------------------------+
| arrow_typeof(test.column1) |
+----------------------------+
| Dictionary(Int32, Utf8)    |
| Dictionary(Int32, Utf8)    |
+----------------------------+
2 rows in set. Query took 0.001 seconds.

❯ select * from test;
+---------+
| column1 |
+---------+
| foo     |
| bar     |
+---------+
2 rows in set. Query took 0.000 seconds.

[devinjdangelo]

Thanks this worked great for the test

[alamb]

Thank you very much @devinjdangelo and @suremarc -- it is really appreciated to get this case working 🙏

[alamb]

I don't see any "dictionary" here -- why is it called dictionary_encoded_parquet?

[devinjdangelo]

I assumed (incorrectly it seems) that alltypes_dictionary.parquet in the parquet-testing submodule was dictionary encoded. https://github.com/apache/parquet-testing/tree/e45cd23f784aab3d6bf0701f8f4e621469ed3be7/data

I will use your suggestion to generate dictionary encoded values instead.

[alamb]

I am somewhat worried about this change (as it effectively changes the meaning of the function to also include logical types).

I think we should either update the comments to reflect this change, or (my preference) make a new function that is explicit. Perhaps something like :

   ///  Returns true of two [`DataType`]s have the same logical data type:
   /// Either the same name and type, or if one is dictionary encoded, then 
   /// the same name and type of the values.
   /// E.g. Dictionary(_, Utf8) is semantically equivalent to Utf8 since both represent an array of strings
  fn datatype_has_same_logical_type(
        match (dt1, dt2) {
            (DataType::Dictionary(_, v1), othertype) => v1.as_ref() == othertype,
            (othertype, DataType::Dictionary(_, v1)) => v1.as_ref() == othertype,
             _ => Self::datatype_is_semantically_equal(dt1, dt2)
        }
  }

[devinjdangelo]

I think your concern is justified as the optimizer also relies on this function and might have stricter equivalence requirements. I created a separate method for logical equivalence which allows for different dictionary encodings as long as values can ultimately be resolved to the same type.

The optimizer continues to use the original semantic equivalence method, and I've updated the insert_into methods to use the softer logical equivalence check.

[alamb]

This might be an easier way to create a data with dictionary encoded data:

❯ create table test as values (arrow_cast('foo', 'Dictionary(Int32, Utf8)')), (arrow_cast('bar', 'Dictionary(Int32, Utf8)'));
0 rows in set. Query took 0.002 seconds.

❯ select arrow_typeof(column1) from test;
+----------------------------+
| arrow_typeof(test.column1) |
+----------------------------+
| Dictionary(Int32, Utf8)    |
| Dictionary(Int32, Utf8)    |
+----------------------------+
2 rows in set. Query took 0.001 seconds.

❯ select * from test;
+---------+
| column1 |
+---------+
| foo     |
| bar     |
+---------+
2 rows in set. Query took 0.000 seconds.

[devinjdangelo]

Checks currently failing, but I don't think it is related to the changes in this PR:

Run cargo check --no-default-features -p datafusion
    Updating crates.io index
error: failed to select a version for the requirement `ahash = "^0.7.0"`
candidate versions found which didn't match: 0.8.5, 0.8.4, 0.4.1, ...
location searched: crates.io index
required by package `hashbrown v0.11.0`
    ... which satisfies dependency `hashbrown = "^0.11"` of package `rkyv v0.7.0`
    ... which satisfies dependency `rkyv = "^0.7"` of package `rust_decimal v1.27.0`
    ... which satisfies dependency `rust_decimal = "^1.27.0"` of package `datafusion v32.0.0 (/__w/arrow-datafusion/arrow-datafusion/datafusion/core)`
    ... which satisfies path dependency `datafusion` of package `datafusion-benchmarks v32.0.0 (/__w/arrow-datafusion/arrow-datafusion/benchmarks)`
perhaps a crate was updated and forgotten to be re-vendored?

[alamb]

I think ahash 0.8.3 just got yanked... Maybe we also need to update hashbrown 🤔

[alamb]

Thank you @devinjdangelo -- this looks great now