diff --git a/cpp/src/arrow/dataset/partition.cc b/cpp/src/arrow/dataset/partition.cc index 4d5c6e8ce7a7..a2aa58565e80 100644 --- a/cpp/src/arrow/dataset/partition.cc +++ b/cpp/src/arrow/dataset/partition.cc @@ -37,6 +37,7 @@ #include "arrow/scalar.h" #include "arrow/util/int_util_overflow.h" #include "arrow/util/logging.h" +#include "arrow/util/string.h" #include "arrow/util/uri.h" #include "arrow/util/utf8.h" @@ -61,6 +62,9 @@ Result SafeUriUnescape(std::string_view encoded) { return decoded; } +/// Percent-encode a single character, eg. '.' becomes "%2E". +std::string PercentEncode(char c) { return "%" + ::arrow::HexEncode(&c, 1); } + // Remove parts of the path after the last filename partition // separator. For example, 12_a_part-0.parquet will become 12_a_ std::string StripNonPrefix(const std::string& path) { @@ -316,12 +320,76 @@ inline std::optional NextValid(const ScalarVector& values, int first_null) } Result> KeyValuePartitioning::FormatPartitionSegments( - const ScalarVector& values) const { + const ScalarVector& values, bool escape_filename_separator) const { std::vector segments(static_cast(schema_->num_fields())); for (int i = 0; i < schema_->num_fields(); ++i) { if (values[i] != nullptr && values[i]->is_valid) { - segments[i] = values[i]->ToString(); + std::string value = values[i]->ToString(); + // Empty segments are dropped when joining path segments, so an empty value + // would not be parsed back correctly + if (value.empty()) { + return Status::Invalid("Partition value for '", schema_->field(i)->name(), + "' is empty, which is not supported with ", type_name(), + " partitioning"); + } + switch (options_.segment_encoding) { + case SegmentEncoding::None: { + // Without encoding, it's an error to write values that are not a valid path + // segment. Backslashes are rejected on all platforms as they are a path + // separator on Windows. Other values that are only invalid on some + // platforms (eg. ending in "." on Windows) are not rejected. + if (value.find_first_of("/\\") != std::string::npos || value == "." || + value == "..") { + return Status::Invalid("Partition value for '", schema_->field(i)->name(), + "' is not a valid path segment: '", value, + "'. Use SegmentEncoding::Uri to write this value."); + } + // The value can't contain the separator between segments, or it would + // not be parsed back correctly + if (escape_filename_separator && + value.find(kFilenamePartitionSep) != std::string::npos) { + return Status::Invalid("Partition value for '", schema_->field(i)->name(), + "' contains the partition separator '", + kFilenamePartitionSep, "': '", value, + "'. Use SegmentEncoding::Uri to write this value."); + } + segments[i] = std::move(value); + break; + } + case SegmentEncoding::Uri: { + std::string escaped = ::arrow::util::UriEscape(value); + if (escape_filename_separator && + escaped.find(kFilenamePartitionSep) != std::string::npos) { + const std::string encoded_separator = PercentEncode(kFilenamePartitionSep); + std::string with_separator; + for (char c : escaped) { + if (c == kFilenamePartitionSep) { + with_separator += encoded_separator; + } else { + with_separator += c; + } + } + escaped = std::move(with_separator); + } + // A leading "." or "_" isn't normally required to be escaped, but "." and + // ".." are not valid directory names, and dataset discovery skips paths + // starting with "." or "_" by default. + if (!escaped.empty() && (escaped[0] == '.' || escaped[0] == '_')) { + escaped.replace(0, 1, PercentEncode(escaped[0])); + } + // Windows strips trailing "." characters from path names, so escape a + // trailing "." too. (A trailing space is already escaped by UriEscape.) + if (!escaped.empty() && escaped.back() == '.') { + escaped.replace(escaped.size() - 1, 1, PercentEncode('.')); + } + segments[i] = std::move(escaped); + break; + } + default: + return Status::NotImplemented("Unknown segment encoding: ", + options_.segment_encoding); + } continue; } @@ -410,7 +478,8 @@ Result DirectoryPartitioning::FormatValues( Result FilenamePartitioning::FormatValues( const ScalarVector& values) const { std::vector segments; - ARROW_ASSIGN_OR_RAISE(segments, FormatPartitionSegments(values)); + ARROW_ASSIGN_OR_RAISE( + segments, FormatPartitionSegments(values, /*escape_filename_separator=*/true)); return PartitionPathFormat{ "", fs::internal::JoinAbstractPath(std::move(segments), kFilenamePartitionSep) + kFilenamePartitionSep}; diff --git a/cpp/src/arrow/dataset/partition.h b/cpp/src/arrow/dataset/partition.h index 315a3d384d28..0c36fdbbf5e8 100644 --- a/cpp/src/arrow/dataset/partition.h +++ b/cpp/src/arrow/dataset/partition.h @@ -118,6 +118,16 @@ std::ostream& operator<<(std::ostream& os, SegmentEncoding segment_encoding); struct ARROW_DS_EXPORT KeyValuePartitioningOptions { /// After splitting a path into components, decode the path components /// before parsing according to this scheme. + /// + /// For directory and filename partitioning, partition values are also encoded + /// according to this scheme when formatting paths for writing. With + /// SegmentEncoding::None, values are written as-is, and values containing "/" + /// or "\", or equal to "." or "..", raise errors. Values starting with "." or + /// "_" are written, but are skipped by dataset discovery unless + /// selector_ignore_prefixes is changed. Other values may not be valid path + /// segments on all platforms, eg. on Windows, values ending in "." or " " are + /// not written correctly. For directory and filename partitioning, empty values + /// raise errors with either encoding. SegmentEncoding segment_encoding = SegmentEncoding::Uri; }; @@ -209,8 +219,16 @@ class ARROW_DS_EXPORT KeyValuePartitioning : public Partitioning { /// Convert a Key to a full expression. Result ConvertKey(const Key& key) const; + /// Format partition values as path segments, encoding them according to + /// options_.segment_encoding. + /// + /// \param[in] values the partition values for each schema field + /// \param[in] escape_filename_separator whether to escape the + /// filename partitioning separator ("_"). If no encoding is used + /// (SegmentEncoding::None), values containing the separator result in an error. Result> FormatPartitionSegments( - const ScalarVector& values) const; + const ScalarVector& values, bool escape_filename_separator = false) const; + Result> ParsePartitionSegments( const std::vector& segments) const; diff --git a/cpp/src/arrow/dataset/partition_test.cc b/cpp/src/arrow/dataset/partition_test.cc index 9f0bd7c0be04..12a7345df69a 100644 --- a/cpp/src/arrow/dataset/partition_test.cc +++ b/cpp/src/arrow/dataset/partition_test.cc @@ -30,6 +30,7 @@ #include "arrow/compute/api_vector.h" #include "arrow/compute/cast.h" #include "arrow/dataset/dataset.h" +#include "arrow/dataset/discovery.h" #include "arrow/dataset/file_ipc.h" #include "arrow/dataset/test_util_internal.h" #include "arrow/filesystem/path_util.h" @@ -245,6 +246,140 @@ TEST_F(TestPartitioning, DirectoryPartitioningEquals) { EXPECT_TRUE(another->Equals(*some_other)); } +TEST_F(TestPartitioning, DirectoryPartitioningEscapesValues) { + partitioning_ = + std::make_shared(schema({field("part", utf8())})); + written_schema_ = partitioning_->schema(); + + AssertFormat(equal(field_ref("part"), literal("plain")), "plain"); + AssertFormat(equal(field_ref("part"), literal("a b")), "a%20b"); + AssertFormat(equal(field_ref("part"), literal("100%")), "100%25"); + AssertFormat(equal(field_ref("part"), literal("a/b")), "a%2Fb"); + AssertFormat(equal(field_ref("part"), literal("ab\\cd")), "ab%5Ccd"); + AssertFormat(equal(field_ref("part"), literal("")), "%3Ctest%3E"); + AssertFormat(equal(field_ref("part"), literal(":|*")), "%3A%7C%2A"); + // "." is only escaped when it is the first or last character + AssertFormat(equal(field_ref("part"), literal("a.b")), "a.b"); + AssertFormat(equal(field_ref("part"), literal(".")), "%2E"); + AssertFormat(equal(field_ref("part"), literal("..")), "%2E%2E"); + AssertFormat(equal(field_ref("part"), literal("...")), "%2E.%2E"); + AssertFormat(equal(field_ref("part"), literal("a.")), "a%2E"); + AssertFormat(equal(field_ref("part"), literal("a ")), "a%20"); + AssertFormat(equal(field_ref("part"), literal(".hidden")), "%2Ehidden"); + AssertFormat(equal(field_ref("part"), literal("../escaped")), "%2E.%2Fescaped"); + AssertFormat(equal(field_ref("part"), literal("..\\escaped")), "%2E.%5Cescaped"); + // "_" is only escaped when it is the first character + AssertFormat(equal(field_ref("part"), literal("a_b")), "a_b"); + AssertFormat(equal(field_ref("part"), literal("_unknown")), "%5Funknown"); + AssertFormat(equal(field_ref("part"), literal("__x_")), "%5F_x_"); + // Empty values would be dropped from the path, so are rejected + AssertFormatError(equal(field_ref("part"), literal(""))); +} + +TEST_F(TestPartitioning, DirectoryPartitioningEscapedValuesRoundTrip) { + // Write and scan a dataset, to check that escaped directories are created + // inside the base directory and aren't skipped by discovery when starting + // with a ".". + 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 = "dataset"; + ASSERT_OK(filesystem->CreateDir(base_path)); + auto schema = arrow::schema( + {arrow::field("part", arrow::utf8()), arrow::field("x", arrow::int32())}); + + auto table = TableFromJSON(schema, { + R"([ + ["../escaped", 0], + [".", 1], + ["..", 2], + ["_unknown", 3] + ])", + }); + + 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())}); + auto partitioning = std::make_shared(partition_schema); + 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); + std::vector all_dirs; + for (const auto& dir : mockfs->AllDirs()) { + all_dirs.push_back(dir.full_path); + } + std::sort(all_dirs.begin(), all_dirs.end()); + std::vector expected_dirs = {"dataset", "dataset/%2E", "dataset/%2E%2E", + "dataset/%2E.%2Fescaped", + "dataset/%5Funknown"}; + ASSERT_EQ(all_dirs, expected_dirs); + + // Read dataset back to verify partition values round-trip correctly + fs::FileSelector selector; + selector.base_dir = base_path; + selector.recursive = true; + FileSystemFactoryOptions factory_options; + factory_options.partitioning = partitioning; + factory_options.partition_base_dir = base_path; + ASSERT_OK_AND_ASSIGN( + auto factory, + FileSystemDatasetFactory::Make(filesystem, selector, ipc_format, factory_options)); + ASSERT_OK_AND_ASSIGN(auto read_dataset, factory->Finish()); + ASSERT_OK_AND_ASSIGN(auto read_scanner_builder, read_dataset->NewScan()); + ASSERT_OK(read_scanner_builder->Project({"part", "x"})); + ASSERT_OK_AND_ASSIGN(auto read_scanner, read_scanner_builder->Finish()); + ASSERT_OK_AND_ASSIGN(auto read_table, read_scanner->ToTable()); + + ASSERT_OK_AND_ASSIGN( + auto sort_indices, + compute::SortIndices(read_table, compute::SortOptions({compute::SortKey("x")}))); + ASSERT_OK_AND_ASSIGN(auto sorted, compute::Take(read_table, sort_indices)); + AssertTablesEqual(*table, *sorted.table(), /*same_chunk_layout=*/false); +} + +TEST_F(TestPartitioning, DirectoryPartitioningNoSegmentEncoding) { + KeyValuePartitioningOptions options; + options.segment_encoding = SegmentEncoding::None; + partitioning_ = std::make_shared(schema({field("part", utf8())}), + ArrayVector(), options); + written_schema_ = partitioning_->schema(); + + // Values are written unescaped and round-trip without decoding + AssertFormat(equal(field_ref("part"), literal("a b")), "a b"); + AssertFormat(equal(field_ref("part"), literal("a%20b")), "a%20b"); + AssertFormat(equal(field_ref("part"), literal("")), ""); + AssertFormat(equal(field_ref("part"), literal("a_b")), "a_b"); + AssertFormat(equal(field_ref("part"), literal("a.b")), "a.b"); + // Values that aren't portable to all platforms are allowed + AssertFormat(equal(field_ref("part"), literal("a.")), "a."); + AssertFormat(equal(field_ref("part"), literal("a ")), "a "); + // Values starting with "." or "_" are allowed, although they are skipped by + // dataset discovery with the default selector_ignore_prefixes + AssertFormat(equal(field_ref("part"), literal(".hidden")), ".hidden"); + AssertFormat(equal(field_ref("part"), literal("_unknown")), "_unknown"); + + // Values that aren't a single valid path segment are rejected + AssertFormatError(equal(field_ref("part"), literal("../escaped"))); + AssertFormatError(equal(field_ref("part"), literal("a/b"))); + AssertFormatError(equal(field_ref("part"), literal("."))); + AssertFormatError(equal(field_ref("part"), literal(".."))); + // Backslash is a path separator on Windows, so is rejected on all platforms + AssertFormatError(equal(field_ref("part"), literal("..\\escaped"))); + AssertFormatError(equal(field_ref("part"), literal("a\\b"))); + // Empty values would be dropped from the path, so are rejected + AssertFormatError(equal(field_ref("part"), literal(""))); +} + TEST_F(TestPartitioning, FilenamePartitioning) { partitioning_ = std::make_shared( schema({field("alpha", int32()), field("beta", utf8())})); @@ -553,6 +688,18 @@ TEST_F(TestPartitioning, HivePartitioningFormat) { partitioning_ = std::make_shared(schema({field("x", large_utf8())})); AssertFormat(equal(field_ref("x"), literal("hello")), "x=hello"); + + // Empty string values round-trip + partitioning_ = std::make_shared( + schema({field("alpha", utf8()), field("beta", utf8())})); + written_schema_ = partitioning_->schema(); + AssertFormat(equal(field_ref("alpha"), literal("")), "alpha="); + AssertFormat(and_(equal(field_ref("alpha"), literal("")), + equal(field_ref("beta"), literal("x"))), + "alpha=/beta=x"); + AssertFormat(and_(equal(field_ref("alpha"), literal("x")), + equal(field_ref("beta"), literal(""))), + "alpha=x/beta="); } TEST_F(TestPartitioning, FilenamePartitioningFormat) { @@ -567,6 +714,87 @@ TEST_F(TestPartitioning, FilenamePartitioningFormat) { AssertFormat(equal(field_ref("alpha"), literal(0)), "", "0_"); } +TEST_F(TestPartitioning, FilenamePartitioningEscapesValues) { + partitioning_ = std::make_shared( + schema({field("alpha", utf8()), field("beta", utf8())})); + written_schema_ = partitioning_->schema(); + + AssertFormat(equal(field_ref("alpha"), literal("plain")), "", "plain_"); + AssertFormat(equal(field_ref("alpha"), literal("a b")), "", "a%20b_"); + AssertFormat(equal(field_ref("alpha"), literal("100%")), "", "100%25_"); + AssertFormat(equal(field_ref("alpha"), literal("a/b")), "", "a%2Fb_"); + AssertFormat(equal(field_ref("alpha"), literal("ab\\cd")), "", "ab%5Ccd_"); + AssertFormat(equal(field_ref("alpha"), literal("")), "", "%3Ctest%3E_"); + AssertFormat(equal(field_ref("alpha"), literal(":|*")), "", "%3A%7C%2A_"); + // "." is only escaped when it is the first or last character of a segment + AssertFormat(equal(field_ref("alpha"), literal("a.b")), "", "a.b_"); + AssertFormat(equal(field_ref("alpha"), literal(".")), "", "%2E_"); + AssertFormat(equal(field_ref("alpha"), literal("..")), "", "%2E%2E_"); + AssertFormat(equal(field_ref("alpha"), literal("a.")), "", "a%2E_"); + AssertFormat(equal(field_ref("alpha"), literal(".hidden")), "", "%2Ehidden_"); + AssertFormat(equal(field_ref("alpha"), literal("../escaped")), "", "%2E.%2Fescaped_"); + AssertFormat(equal(field_ref("alpha"), literal("..\\escaped")), "", "%2E.%5Cescaped_"); + // "_" is the segment separator, so it is always escaped + AssertFormat(equal(field_ref("alpha"), literal("a_b")), "", "a%5Fb_"); + AssertFormat(equal(field_ref("alpha"), literal("_unknown")), "", "%5Funknown_"); + AssertFormat(equal(field_ref("alpha"), literal("__x_")), "", "%5F%5Fx%5F_"); + // Each segment is escaped independently + AssertFormat(and_(equal(field_ref("alpha"), literal("a/b")), + equal(field_ref("beta"), literal(".hidden"))), + "", "a%2Fb_%2Ehidden_"); + AssertFormat(and_(equal(field_ref("alpha"), literal("a_b")), + equal(field_ref("beta"), literal("c_d"))), + "", "a%5Fb_c%5Fd_"); + // Empty values would be dropped from the filename, so are rejected + AssertFormatError(equal(field_ref("alpha"), literal(""))); + AssertFormatError(and_(equal(field_ref("alpha"), literal("")), + equal(field_ref("beta"), literal("x")))); + AssertFormatError(and_(equal(field_ref("alpha"), literal("x")), + equal(field_ref("beta"), literal("")))); +} + +TEST_F(TestPartitioning, FilenamePartitioningNoSegmentEncoding) { + KeyValuePartitioningOptions options; + options.segment_encoding = SegmentEncoding::None; + partitioning_ = std::make_shared( + schema({field("alpha", utf8()), field("beta", utf8())}), ArrayVector(), options); + written_schema_ = partitioning_->schema(); + + // Values are written unescaped and round-trip without decoding + AssertFormat(equal(field_ref("alpha"), literal("a b")), "", "a b_"); + AssertFormat(equal(field_ref("alpha"), literal("a%20b")), "", "a%20b_"); + AssertFormat(equal(field_ref("alpha"), literal("")), "", "_"); + AssertFormat(and_(equal(field_ref("alpha"), literal("a b")), + equal(field_ref("beta"), literal("c.d"))), + "", "a b_c.d_"); + // Values that aren't portable to all platforms are allowed + AssertFormat(equal(field_ref("alpha"), literal("a.")), "", "a._"); + // Values starting with "." are allowed, although they are skipped by dataset + // discovery with the default selector_ignore_prefixes + AssertFormat(equal(field_ref("alpha"), literal(".hidden")), "", ".hidden_"); + + // Values that aren't a single valid path segment are rejected + AssertFormatError(equal(field_ref("alpha"), literal("../escaped"))); + AssertFormatError(equal(field_ref("alpha"), literal("a/b"))); + AssertFormatError(equal(field_ref("alpha"), literal("."))); + AssertFormatError(equal(field_ref("alpha"), literal(".."))); + AssertFormatError(equal(field_ref("alpha"), literal("..\\escaped"))); + AssertFormatError(equal(field_ref("alpha"), literal("a\\b"))); + AssertFormatError(and_(equal(field_ref("alpha"), literal("ok")), + equal(field_ref("beta"), literal("a/b")))); + + // Values containing the "_" segment separator are rejected + AssertFormatError(equal(field_ref("alpha"), literal("a_b"))); + AssertFormatError(equal(field_ref("alpha"), literal("_unknown"))); + AssertFormatError(and_(equal(field_ref("alpha"), literal("ok")), + equal(field_ref("beta"), literal("c_d")))); + + // Empty values would be dropped from the filename, so are rejected + AssertFormatError(equal(field_ref("alpha"), literal(""))); + AssertFormatError(and_(equal(field_ref("alpha"), literal("x")), + equal(field_ref("beta"), literal("")))); +} + TEST_F(TestPartitioning, DiscoverHiveSchema) { auto options = HivePartitioningFactoryOptions(); options.null_fallback = "xyz"; diff --git a/python/pyarrow/_dataset.pyx b/python/pyarrow/_dataset.pyx index d40614a61fc0..8889090801c8 100644 --- a/python/pyarrow/_dataset.pyx +++ b/python/pyarrow/_dataset.pyx @@ -2758,8 +2758,8 @@ cdef class DirectoryPartitioning(KeyValuePartitioning): every value which may be taken by the corresponding column or an error will be raised in parsing. 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). + Method to use to encode and decode path segments. Valid values are + "uri" (URI-encode and decode segments) and "none" (leave as-is). Returns ------- @@ -3037,8 +3037,8 @@ cdef class FilenamePartitioning(KeyValuePartitioning): every value which may be taken by the corresponding column or an error will be raised in parsing. 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). + Method to use to encode and decode path segments. Valid values are + "uri" (URI-encode and decode segments) and "none" (leave as-is). Returns ------- diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index b50441a1ae05..3b575e24d5e1 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -853,8 +853,8 @@ def test_partitioning_pickling(pickle_module): "flavor, expected_defined_partition, expected_undefined_partition", [ ("HivePartitioning", (r"foo=A/bar=ant%20bee", ""), ("", "")), - ("DirectoryPartitioning", (r"A/ant bee", ""), ("", "")), - ("FilenamePartitioning", ("", r"A_ant bee_"), ("", "_")), + ("DirectoryPartitioning", (r"A/ant%20bee", ""), ("", "")), + ("FilenamePartitioning", ("", r"A_ant%20bee_"), ("", "_")), ], ) def test_dataset_partitioning_format(