DictionaryKeyOverflowError on DataFrame.write_parquet

[valkum]

Describe the bug

I tried to come up with a repro case for a slow query (compared to polars) and encountered this bug. As the error is originating in https://github.com/apache/arrow-rs, I assume this is more of a tracking issue, but it might also be caused by how datafusion calls parquet. I haven't looked for the root cause yet.

With a sample dataset of 10m rows, I run into the following issue:
pyo3_runtime.PanicException: MutableArrayData::new is infallible: DictionaryKeyOverflowError

What this does is it takes an unnested schema, nests it and aggregates the newly nested struct into a list.

This does not happen for 1m rows. The input parquet file only contains a Dict with 10 keys. I am not sure why this blows up. When nesting the dict key, the same dict can be reused.

To Reproduce

For a repro case see here https://github.com/valkum/polars-datafusion-comparison/tree/datafusion_bug

Expected behavior

DataFrame.write_parquet should succeed.

Additional context

Failing plan:

DataFrame()
+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan  | Projection: df.name, df.group, CASE WHEN array_element(markets, Int64(1)) IS NOT NULL THEN markets ELSE List() END AS CASE WHEN markets[Int64(1)] IS NOT NULL THEN markets ELSE NULL END                                                                                                                                                                                                                                                                                                                                                                                                                                                   |
|               |   Aggregate: groupBy=[[df.name, df.group]], aggr=[[array_agg(market) AS markets]]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
|               |     Projection: df.name, df.group, CASE WHEN df.market IS NOT NULL AND df.price IS NOT NULL THEN named_struct(Utf8("market"), CAST(df.market AS Dictionary(UInt16, Utf8)), Utf8("price"), df.price) ELSE Struct({market:,price:}) END AS market                                                                                                                                                                                                                                                                                                                                                                                            |
|               |       TableScan: df projection=[name, group, market, price]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                |
| physical_plan | ProjectionExec: expr=[name@0 as name, group@1 as group, CASE WHEN array_element(markets@2, 1) IS NOT NULL THEN markets@2 END as CASE WHEN markets[Int64(1)] IS NOT NULL THEN markets ELSE NULL END]                                                                                                                                                                                                                                                                                                                                                                                                                                        |
|               |   AggregateExec: mode=FinalPartitioned, gby=[name@0 as name, group@1 as group], aggr=[markets], ordering_mode=Sorted                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
|               |     SortExec: expr=[name@0 ASC NULLS LAST, group@1 ASC NULLS LAST], preserve_partitioning=[true]                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           |
|               |       CoalesceBatchesExec: target_batch_size=8192                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
|               |         RepartitionExec: partitioning=Hash([name@0, group@1], 10), input_partitions=10                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                     |
|               |           AggregateExec: mode=Partial, gby=[name@0 as name, group@1 as group], aggr=[markets], ordering_mode=Sorted                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
|               |             ProjectionExec: expr=[name@0 as name, group@1 as group, CASE WHEN market@2 IS NOT NULL AND price@3 IS NOT NULL THEN named_struct(market, CAST(market@2 AS Dictionary(UInt16, Utf8)), price, price@3) END as market]                                                                                                                                                                                                                                                                                                                                                                                                            |
|               |               DataSourceExec: file_groups={10 groups: [[Users/valkum/git/polars-datafusion-comparison/sample-10m.parquet:0..39293391], [Users/valkum/git/polars-datafusion-comparison/sample-10m.parquet:39293391..78586782], [Users/valkum/git/polars-datafusion-comparison/sample-10m.parquet:78586782..117880173], [Users/valkum/git/polars-datafusion-comparison/sample-10m.parquet:117880173..157173564], [Users/valkum/git/polars-datafusion-comparison/sample-10m.parquet:157173564..196466955], ...]}, projection=[name, group, market, price], output_ordering=[name@0 ASC NULLS LAST, group@1 ASC NULLS LAST], file_type=parquet |
|               |                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            |
+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

[valkum]

If you stick to the Utf8 datatype for market, the automatic dict encoding done by the parquet writer works. So assume some operation is not properly merging the dict keys. The weird thing is that this only happens with 10m rows not with 1m even if all rows in the sample file have a market set.

[alamb]

I suspect that whatever is writing dictionary arrays is using a key type that is too small for the number of distinct values. For example, using a Dict(Int8, Utf8) to store a column with 1000 distinct values (where the type can only store 256 distinct values)

[valkum]

But the example is only using 10 distinct keys. If the writer is using a smaller Dict (Uint16 vs. Int8), a dict somewhere is collecting values as distinct that aren't distinct.

[duongcongtoai]

take

[duongcongtoai]

I also came across the same thing working on this PR

[duongcongtoai]

        let (dictionary, dict_concat) = match &data_type {
            DataType::Dictionary(_, _) => {
                // If more than one dictionary, concatenate dictionaries together
                let dict_concat = !arrays
                    .windows(2)
                    .all(|a| a[0].child_data()[0].ptr_eq(&a[1].child_data()[0]));

                match dict_concat {
                    false => (Some(arrays[0].child_data()[0].clone()), false),
                    true => {
                        if let Capacities::Dictionary(_, _) = capacities {
                            panic!("dictionary capacity not yet supported")
                        }
                        let dictionaries: Vec<_> =
                            arrays.iter().map(|array| &array.child_data()[0]).collect();
                        let lengths: Vec<_> = dictionaries
                            .iter()
                            .map(|dictionary| dictionary.len())
                            .collect();
                        let capacity = lengths.iter().sum();

                        let mut mutable = MutableArrayData::new(dictionaries, false, capacity);

                        for (i, len) in lengths.iter().enumerate() {
                            mutable.extend(i, 0, *len)
                        }

                        (Some(mutable.freeze()), true)
                    }
                }
            }
            _ => (None, false),
        };

I found this logic relevant, if dict_concat is set to true, the dict belongs to each children array is considered distinct dictionary and the function MutableArrayData::with_capacities will try to concat them, which may result into DictionaryKeyOverFlowError

        let extend_values = match &data_type {
            DataType::Dictionary(_, _) => {
                let mut next_offset = 0;
                let extend_values: Result<Vec<_>, _> = arrays
                    .iter()
                    .map(|array| {
                        let offset = next_offset;
                        let dict_len = array.child_data()[0].len();

                        if dict_concat {
                            next_offset += dict_len;
                        }

                        build_extend_dictionary(array, offset, offset + dict_len)
                            .ok_or(ArrowError::DictionaryKeyOverflowError)
                    })
                    .collect();

                extend_values.expect("MutableArrayData::new is infallible")
            }

[duongcongtoai]

                let dict_concat = !arrays
                    .windows(2)
                    .all(|a| a[0].child_data()[0].ptr_eq(&a[1].child_data()[0]));

something during execution change the address of the child_data array, or this is the side effect of spilling

[valkum]

Thanks @duongcongtoai.

If I understand the arrow c++ impl correctly, it does not concatenate on reading. You can explicitly merge dicts for a chunked array using DictionaryUnifier::UnifyChunkedArray. See tests for examples.

build_extend_dictionary is in the direct call chain of ArrayReader::consume_batch for ListArrayReader via MutableArrayData::new -> MutableArrayData::with_capacity.

The respective PR that added the concat to arrow-rs is apache/arrow-rs#15. For arrow c++ see apache/arrow#3165.

So IIUC this means that arrow-rs currently can inflate the dictionaries when reading and writing a parquet file without any other ops.

[duongcongtoai]

i think so, is only checking if the ptrs are equal between dictionary, this will be broken when:

• spilling
• repartition (if you set the original script to only use 1 partition the issue disappear)

I found this issue relevant: https://github.com/apache/arrow-rs/pull/7519/files but maybe the implementation does not consider logically equivalent dictionary (different addr by the same dataset)

[duongcongtoai]

although

fn interleave_dictionaries<K: ArrowDictionaryKeyType>(
    arrays: &[&dyn Array],
    indices: &[(usize, usize)],
) -> Result<ArrayRef, ArrowError> {
    let dictionaries: Vec<_> = arrays.iter().map(|x| x.as_dictionary::<K>()).collect();

    if !should_merge_dictionary_values::<K>(&dictionaries, indices.len()) {
        return interleave_fallback(arrays, indices);

Current interleave actually check the function should_merge_dictionary_values, which will returns true if the total value in the dictionary will overflow, and will try to merge the them, however it is not the case for for this dataset, and it continue calling interleave_fallback which cause the panic

[duongcongtoai]

alright, interleave_dictionaries is not called because the data type is List of dict, i guess we need to check overflow for this case as well