jorisvandenbossche commented on a change in pull request #9323:
URL: https://github.com/apache/arrow/pull/9323#discussion_r575328856
##########
File path: python/pyarrow/tests/test_dataset.py
##########
@@ -1587,33 +1587,54 @@ def test_open_dataset_non_existing_file():
@pytest.mark.parquet
@pytest.mark.parametrize('partitioning', ["directory", "hive"])
+@pytest.mark.parametrize('null_fallback', ['xyz', None])
@pytest.mark.parametrize('partition_keys', [
(["A", "B", "C"], [1, 2, 3]),
([1, 2, 3], ["A", "B", "C"]),
(["A", "B", "C"], ["D", "E", "F"]),
([1, 2, 3], [4, 5, 6]),
+ ([1, None, 3], ["A", "B", "C"]),
+ ([1, 2, 3], ["A", None, "C"]),
+ ([None, 2, 3], [None, 2, 3]),
])
-def test_open_dataset_partitioned_dictionary_type(tempdir, partitioning,
- partition_keys):
+def test_open_dataset_partitioned_dictionary_type(
+ tempdir, partitioning, null_fallback, partition_keys
+):
# ARROW-9288 / ARROW-9476
import pyarrow.parquet as pq
- table = pa.table({'a': range(9), 'b': [0.] * 4 + [1.] * 5})
+
+ table = pa.table({'a': range(9), 'b': [0.0] * 4 + [1.0] * 5})
+
+ if None in partition_keys[0] or None in partition_keys[1]:
+ # Directory partitioning can't handle the first part being null
+ return
Review comment:
only return here if `partitioning == "directory"` ?
##########
File path: python/pyarrow/tests/test_dataset.py
##########
@@ -1587,33 +1587,54 @@ def test_open_dataset_non_existing_file():
@pytest.mark.parquet
@pytest.mark.parametrize('partitioning', ["directory", "hive"])
+@pytest.mark.parametrize('null_fallback', ['xyz', None])
Review comment:
What does `null_fallback=None` mean? (based on the docstring above it
seems it can only be a string?)
##########
File path: cpp/src/arrow/dataset/partition.cc
##########
@@ -74,15 +74,26 @@ Status KeyValuePartitioning::SetDefaultValuesFromKeys(const
Expression& expr,
RecordBatchProjector*
projector) {
ARROW_ASSIGN_OR_RAISE(auto known_values, ExtractKnownFieldValues(expr));
for (const auto& ref_value : known_values) {
- if (!ref_value.second.is_scalar()) {
- return Status::Invalid("non-scalar partition key ",
ref_value.second.ToString());
+ const auto& known_value = ref_value.second;
+ if (known_value.concrete() && !known_value.datum.is_scalar()) {
+ return Status::Invalid("non-scalar partition key ",
known_value.datum.ToString());
}
ARROW_ASSIGN_OR_RAISE(auto match,
ref_value.first.FindOneOrNone(*projector->schema()));
if (match.empty()) continue;
- RETURN_NOT_OK(projector->SetDefaultValue(match,
ref_value.second.scalar()));
+
+ const auto& field = projector->schema()->field(match[0]);
+ if (known_value.concrete()) {
+ RETURN_NOT_OK(projector->SetDefaultValue(match,
known_value.datum.scalar()));
+ } else if (known_value.valid) {
+ return Status::Invalid(
+ "Partition expression not defined enough to set default value for ",
Review comment:
What does "not defined enough" in practice mean? (or what would be an
example?)
##########
File path: python/pyarrow/tests/test_dataset.py
##########
@@ -1587,33 +1587,54 @@ def test_open_dataset_non_existing_file():
@pytest.mark.parquet
@pytest.mark.parametrize('partitioning', ["directory", "hive"])
+@pytest.mark.parametrize('null_fallback', ['xyz', None])
@pytest.mark.parametrize('partition_keys', [
(["A", "B", "C"], [1, 2, 3]),
([1, 2, 3], ["A", "B", "C"]),
(["A", "B", "C"], ["D", "E", "F"]),
([1, 2, 3], [4, 5, 6]),
+ ([1, None, 3], ["A", "B", "C"]),
+ ([1, 2, 3], ["A", None, "C"]),
+ ([None, 2, 3], [None, 2, 3]),
])
-def test_open_dataset_partitioned_dictionary_type(tempdir, partitioning,
- partition_keys):
+def test_open_dataset_partitioned_dictionary_type(
Review comment:
you added this to a test that is specifically about reading partitioned
datasets while inferring the partition fields as dictionary. Which is fine (as
this case also needs to be able to hand that), but this should also work (and
so be tested) in the default case not inferring dictionary type?
And should we also have a test for the writing part? (this one only tests
reading)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
