[C++][Dataset] Cannot write dataset with dictionary column as partition field

[asfimport]

When the column to use for partitioning is dictionary encoded, we get this error:

In [9]: import pyarrow.dataset as ds

In [10]: part = ["xxx"] * 3 + ["yyy"] * 3
    ...: table = pa.table([
    ...:     pa.array(range(len(part))),
    ...:     pa.array(part).dictionary_encode(),
    ...: ], names=['col', 'part'])

In [11]: part = ds.partitioning(table.select(["part"]).schema)

In [12]: ds.write_dataset(table, "test_dataset_dict_part", format="parquet", partitioning=part)
---------------------------------------------------------------------------
ArrowTypeError                            Traceback (most recent call last)
<ipython-input-12-c7b81c9b0bda> in <module>
----> 1 ds.write_dataset(table, "test_dataset_dict_part", format="parquet", partitioning=part)

~/scipy/repos/arrow/python/pyarrow/dataset.py in write_dataset(data, base_dir, basename_template, format, partitioning, schema, filesystem, file_options, use_threads)
    773     _filesystemdataset_write(
    774         data, base_dir, basename_template, schema,
--> 775         filesystem, partitioning, file_options, use_threads,
    776     )

~/scipy/repos/arrow/python/pyarrow/_dataset.pyx in pyarrow._dataset._filesystemdataset_write()

~/scipy/repos/arrow/python/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowTypeError: scalar xxx (of type string) is invalid for part: dictionary<values=string, indices=int32, ordered=0>
In ../src/arrow/dataset/filter.cc, line 1082, code: VisitConjunctionMembers(*and_.left_operand(), visitor)
In ../src/arrow/dataset/partition.cc, line 257, code: VisitKeys(expr, [&](const std::string& name, const std::shared_ptr<Scalar>& value) { auto&& _error_or_value28 = (FieldRef(name).FindOneOrNone(*schema_)); do { ::arrow::Status __s = ::arrow::internal::GenericToStatus((_error_or_value28).status()); do { if ((__builtin_expect(!!(!__s.ok()), 0))) { ::arrow::Status _st = (__s); _st.AddContextLine("../src/arrow/dataset/partition.cc", 257, "(_error_or_value28).status()"); return _st; } } while (0); } while (false); auto match = std::move(_error_or_value28).ValueUnsafe();;; if (match) { const auto& field = schema_->field(match[0]); if (!value->type->Equals(field->type())) { return Status::TypeError("scalar ", value->ToString(), " (of type ", *value->type, ") is invalid for ", field->ToString()); } values[match[0]] = value.get(); } return Status::OK(); })
In ../src/arrow/dataset/file_base.cc, line 321, code: (_error_or_value24).status()
In ../src/arrow/dataset/file_base.cc, line 354, code: task_group->Finish()

While this seems a quit normal use case, as this column will typically be repeated many times (and we also support reading it as such with dictionary type, so a roundtrip is currently not possible in that case)

I tagged it for 2.0.0 for a moment in case it's possible today, but I didn't yet look into how easy it would be to fix.

Reporter: Joris Van den Bossche / @jorisvandenbossche
Assignee: Ben Kietzman / @bkietz

PRs and other links:

• GitHub Pull Request #9130

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

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
cc @bkietz

[asfimport]

Lance Dacey / @ldacey:
What is the best workaround for this issue right now? I was playing around with making a new partition schema if a dictionary type was found in my partition columns:

 

partitioning = None
part_schema = t.select(["project", "date"]).schema
fields = []
for part in part_schema:
    if pa.types.is_dictionary(part.type):
        fields.append(pa.field(part.name, part.type.value_type))
    else:
        fields.append(pa.field(part.name, part.type))
new_schema = pa.schema(fields)
partitioning = ds.partitioning(new_schema, flavor="hive")

This seems to work for me. My only issue is if I have multiple partition columns with different types.

This would return an error when I read the dataset with ds.dataset():

partitioning = ds.partitioning(pa.schema([('date', pa.date32()), ("project", pa.dictionary(index_type=pa.int32(), value_type=pa.string()))]), flavor="hive")

ArrowInvalid: No dictionary provided for dictionary field project: dictionary<values=string, indices=int32, ordered=0>

 

And this returns dictionaries for both partitions (instead of date being pa.date32()) which is not ideal:

partitioning=ds.HivePartitioning.discover(infer_dictionary=True)

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
@ldacey yes, this problem of not being able to specify one column with a dictionary type (and discover(infer_dictionary=True) doing it for all columns) is indeed a open issue in the user API. I just noted that on the PR as well: #9130 (comment)

So with the imminent pyarrow 3.0, you will be able to do:

partitioning = ds.partitioning(
    pa.schema([('date', pa.date32()), ("project", pa.dictionary(index_type=pa.int32(), value_type=pa.string()))]),
    flavor="hive", dictionaries={'project': pa.array([... all possible values for the project field ...]}
)

Note that the above is only needed for reading, for writing it will work without specifying those dictionaries keyword manually. I will open a new JIRA to further improve this.

[asfimport]

Lance Dacey / @ldacey:
Nice - how would you generally go about finding the array of values? Would it be detected from the file paths, or would I need  store it externally somewhere (sometimes new categories could be added into the field without me being aware so explicitly listing them in my code might be weird)?

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:

how would you generally go about finding the array of values?

Well, that's up to you .. (parsing it from the file paths, storing that information somewhere, ..). But, so, my hunch is that we shouldn't actually require the user to pass it, since pyarrow can infer that itself from parsing the file paths if not provided. I opened ARROW-11260 for this.

[asfimport]

Ben Kietzman / @bkietz:
Issue resolved by pull request 9130
#9130