diff --git a/cpp/src/arrow/dataset/partition.cc b/cpp/src/arrow/dataset/partition.cc index 4d5c6e8ce7a7..2e6153bf2f99 100644 --- a/cpp/src/arrow/dataset/partition.cc +++ b/cpp/src/arrow/dataset/partition.cc @@ -755,7 +755,18 @@ Result HivePartitioning::FormatValues( // field_index <-> path nesting relation segments[i] = name + "=" + hive_options_.null_fallback; } else { - segments[i] = name + "=" + arrow::util::UriEscape(values[i]->ToString()); + std::string value = values[i]->ToString(); + switch (hive_options_.segment_encoding) { + case SegmentEncoding::None: + break; + case SegmentEncoding::Uri: + value = arrow::util::UriEscape(value); + break; + default: + return Status::NotImplemented("Unknown segment encoding: ", + hive_options_.segment_encoding); + } + segments[i] = name + "=" + std::move(value); } } diff --git a/cpp/src/arrow/dataset/partition.h b/cpp/src/arrow/dataset/partition.h index 315a3d384d28..c507a68cc5bf 100644 --- a/cpp/src/arrow/dataset/partition.h +++ b/cpp/src/arrow/dataset/partition.h @@ -116,8 +116,9 @@ std::ostream& operator<<(std::ostream& os, SegmentEncoding segment_encoding); /// \brief Options for key-value based partitioning (hive/directory). struct ARROW_DS_EXPORT KeyValuePartitioningOptions { - /// After splitting a path into components, decode the path components - /// before parsing according to this scheme. + /// How to encode and decode partition segment values. With SegmentEncoding::Uri, + /// values are URI-decoded when parsing and URI-encoded when formatting (Hive + /// partitioning only). With SegmentEncoding::None, values are used as-is. SegmentEncoding segment_encoding = SegmentEncoding::Uri; }; @@ -132,8 +133,9 @@ struct ARROW_DS_EXPORT PartitioningFactoryOptions { /// will only check discovered fields against the schema and update internal /// state (such as dictionaries). std::shared_ptr schema; - /// After splitting a path into components, decode the path components - /// before parsing according to this scheme. + /// How to encode and decode partition segment values. With SegmentEncoding::Uri, + /// values are URI-decoded when parsing and URI-encoded when formatting (Hive + /// partitioning only). With SegmentEncoding::None, values are used as-is. SegmentEncoding segment_encoding = SegmentEncoding::Uri; KeyValuePartitioningOptions AsPartitioningOptions() const; diff --git a/cpp/src/arrow/dataset/partition_test.cc b/cpp/src/arrow/dataset/partition_test.cc index 9f0bd7c0be04..77104aad9aa0 100644 --- a/cpp/src/arrow/dataset/partition_test.cc +++ b/cpp/src/arrow/dataset/partition_test.cc @@ -553,6 +553,14 @@ TEST_F(TestPartitioning, HivePartitioningFormat) { partitioning_ = std::make_shared(schema({field("x", large_utf8())})); AssertFormat(equal(field_ref("x"), literal("hello")), "x=hello"); + + // GH-34905: don't URI-encode values when segment_encoding is none + HivePartitioningOptions options; + options.segment_encoding = SegmentEncoding::None; + partitioning_ = std::make_shared(schema({field("x", utf8())}), + ArrayVector(), options); + written_schema_ = partitioning_->schema(); + AssertFormat(equal(field_ref("x"), literal("a b&c=d")), "x=a b&c=d"); } TEST_F(TestPartitioning, FilenamePartitioningFormat) { @@ -945,6 +953,50 @@ TEST_F(TestPartitioning, WriteHiveWithSlashesInValues) { } } +TEST_F(TestPartitioning, WriteHiveWithoutUriEncoding) { + // GH-34905: segment_encoding=none should write partition values verbatim + fs::TimePoint mock_now = std::chrono::system_clock::now(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr filesystem, + fs::internal::MockFileSystem::Make(mock_now, {})); + auto base_path = ""; + ASSERT_OK(filesystem->CreateDir(base_path)); + auto schema = arrow::schema( + {arrow::field("a", arrow::int64()), arrow::field("part", arrow::utf8())}); + auto table = TableFromJSON(schema, {R"([ + [0, "My Fancy Product"], + [1, "100% done"], + [2, "My Fancy Product"] + ])"}); + + auto dataset = std::make_shared(table); + ASSERT_OK_AND_ASSIGN(auto scanner_builder, dataset->NewScan()); + ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish()); + + auto partition_schema = arrow::schema({arrow::field("part", arrow::utf8())}); + HivePartitioningOptions options; + options.segment_encoding = SegmentEncoding::None; + auto partitioning = std::make_shared(partition_schema, + ArrayVector(), options); + auto ipc_format = std::make_shared(); + dataset::FileSystemDatasetWriteOptions write_options; + write_options.file_write_options = ipc_format->DefaultWriteOptions(); + write_options.filesystem = filesystem; + write_options.base_dir = base_path; + write_options.partitioning = partitioning; + write_options.basename_template = "part{i}.arrow"; + ASSERT_OK(dataset::FileSystemDataset::Write(write_options, scanner)); + + auto mockfs = + arrow::internal::checked_pointer_cast(filesystem); + auto all_dirs = mockfs->AllDirs(); + + std::vector expected_paths = {"part=100% done", "part=My Fancy Product"}; + ASSERT_EQ(all_dirs.size(), expected_paths.size()); + for (size_t i = 0; i < all_dirs.size(); i++) { + ASSERT_EQ(all_dirs[i].full_path, expected_paths[i]); + } +} + TEST_F(TestPartitioning, EtlThenHive) { FieldVector etl_fields{field("year", int16()), field("month", int8()), field("day", int8()), field("hour", int8())}; diff --git a/python/pyarrow/_dataset.pyx b/python/pyarrow/_dataset.pyx index d40614a61fc0..1b1e9ac8242b 100644 --- a/python/pyarrow/_dataset.pyx +++ b/python/pyarrow/_dataset.pyx @@ -2893,8 +2893,9 @@ cdef class HivePartitioning(KeyValuePartitioning): null_fallback : str, default "__HIVE_DEFAULT_PARTITION__" If any field is None then this fallback will be used as a label segment_encoding : str, default "uri" - After splitting paths into segments, decode the segments. Valid - values are "uri" (URI-decode segments) and "none" (leave as-is). + If "uri", URI-encode partition values when writing and URI-decode + them when reading. If "none", leave them as-is; values must not + contain path separators. Returns ------- diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 9ddfdfdc9c49..9978dec0aa95 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -4324,6 +4324,28 @@ def test_write_dataset_with_field_names_hive(tempdir): assert load_back_table.equals(table) +@pytest.mark.parametrize("segment_encoding,expected", [ + ("uri", ["part=100%25%20done", "part=My%20Fancy%20Product"]), + ("none", ["part=100% done", "part=My Fancy Product"]), +]) +def test_write_dataset_hive_segment_encoding(tempdir, segment_encoding, expected): + # GH-34905: segment_encoding controls how hive partition values are written + table = pa.table({'a': [1, 2], 'part': ["My Fancy Product", "100% done"]}) + partitioning = ds.HivePartitioning( + pa.schema([("part", pa.string())]), segment_encoding=segment_encoding) + + ds.write_dataset(table, tempdir, format='ipc', partitioning=partitioning) + + load_back = ds.dataset(tempdir, format='ipc', partitioning=partitioning) + partitioning_dirs = { + str(pathlib.Path(f).relative_to(tempdir).parent) for f in load_back.files + } + assert partitioning_dirs == set(expected) + + load_back_table = load_back.to_table() + assert load_back_table.sort_by('a').equals(table.sort_by('a')) + + def test_write_dataset_with_scanner(tempdir): table = pa.table({'a': ['x', 'y', None], 'b': ['x', 'y', 'z'], 'c': [1, 2, 3]})