[C++][Dataset] Discovery of partition field as dictionary type segfaulting with HivePartitioning

[asfimport]

Testing new feature from ARROW-8647, python test that reproduces it:

@pytest.mark.parquet
@pytest.mark.parametrize('partitioning', ["directory", "hive"])
def test_open_dataset_partitioned_dictionary_type(tempdir, partitioning):
    import pyarrow.parquet as pq
    table = pa.table({'a': range(9), 'b': [0.] * 4 + [1.] * 5})

    path = tempdir / "dataset"
    path.mkdir()

    for part in ["A", "B", "C"]:
        fmt = "{}" if partitioning == "directory" else "part={}"
        part = path / fmt.format(part)
        part.mkdir()
        pq.write_table(table, part / "test.parquet")

    if partitioning == "directory":
        part = ds.DirectoryPartitioning.discover(["part"], max_partition_dictionary_size=-1)
    else:
        part = ds.HivePartitioning.discover(max_partition_dictionary_size=-1)
    
    dataset = ds.dataset(str(path), partitioning=part)
    expected_schema = table.schema.append(
        pa.field("part", pa.dictionary(pa.int32(), pa.string()))
    )
    assert dataset.schema.equals(expected_schema)

This test fails (segfaults) for HivePartitioning, but works for DirectoryPartitioning

Reporter: Joris Van den Bossche / @jorisvandenbossche
Assignee: Joris Van den Bossche / @jorisvandenbossche

Related issues:

• [C++][Dataset] Optionally encode partition field values as dictionary type (is related to)

PRs and other links:

• GitHub Pull Request #7608

_{Note: This issue was originally created as ARROW-9288. Please see the migration documentation for further details.}

[asfimport]

Antoine Pitrou / @pitrou:
Can you debug this by yourself or do you need some help?

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Yes, I already debugged it, but was trying to look for a solution as well before posting it here. But my notes up to now:

From time to time, it also gives "ArrowInvalid: No dictionary provided for dictionary field part: dictionary<values=string, indices=int32, ordered=0>" instead of segfaulting.

From debugging, it seems the problem is here:

arrow/cpp/src/arrow/dataset/partition.cc

Line 159 in f25a014

value.dictionary = dictionaries[field_index];

The field_index is different for hive vs directory partitioning (or maybe whether an explicit schema is provided or not). For DirectoryPartitioning, the schema passed to ConvertKey is the schema of only the partition fields, and thus field_index is 0-based into that. But for HivePartitioning, the schema passed is the schema of the full dataset, so the field_index typically doesn't start at 0 since we by default append partition fields at the end. This causes the field_index to be out of bounds for the dictionaries (which has a length depending on the number of partition fields)

[asfimport]

Wes McKinney / @wesm:
Issue resolved by pull request 7608
#7608