diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index eead221dbda..9bf41b9f760 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -424,6 +424,7 @@ set(ARROW_SRCS extension_type.cc extension/bool8.cc extension/json.cc + extension/parquet_file.cc extension/parquet_variant.cc extension/uuid.cc pretty_print.cc diff --git a/cpp/src/arrow/extension/CMakeLists.txt b/cpp/src/arrow/extension/CMakeLists.txt index ae52bc32a99..15a27ec4a26 100644 --- a/cpp/src/arrow/extension/CMakeLists.txt +++ b/cpp/src/arrow/extension/CMakeLists.txt @@ -27,4 +27,10 @@ add_arrow_test(test PREFIX "arrow-canonical-extensions") +add_arrow_test(test + SOURCES + parquet_file_test.cc + PREFIX + "arrow-parquet-file-extension") + arrow_install_all_headers("arrow/extension") diff --git a/cpp/src/arrow/extension/meson.build b/cpp/src/arrow/extension/meson.build index 84dafe4bbe3..a5c0c1c4bae 100644 --- a/cpp/src/arrow/extension/meson.build +++ b/cpp/src/arrow/extension/meson.build @@ -31,12 +31,20 @@ exc = executable( ) test('arrow-canonical-extensions-test', exc) +file_extension_test = executable( + 'arrow-parquet-file-extension-test', + sources: ['parquet_file_test.cc'], + dependencies: [arrow_test_dep], +) +test('arrow-parquet-file-extension-test', file_extension_test) + install_headers( [ 'bool8.h', 'fixed_shape_tensor.h', 'json.h', 'opaque.h', + 'parquet_file.h', 'parquet_variant.h', 'uuid.h', 'variable_shape_tensor.h', diff --git a/cpp/src/arrow/extension/parquet_file.cc b/cpp/src/arrow/extension/parquet_file.cc new file mode 100644 index 00000000000..d10227d40f2 --- /dev/null +++ b/cpp/src/arrow/extension/parquet_file.cc @@ -0,0 +1,128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "arrow/extension/parquet_file.h" + +#include "arrow/extension_type.h" +#include "arrow/result.h" +#include "arrow/status.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/logging_internal.h" + +namespace arrow::extension { + +namespace { + +bool IsSupportedField(const std::shared_ptr& field) { + if (!field->nullable()) { + return false; + } + if (field->name() == "uri" || field->name() == "content_type" || + field->name() == "checksum") { + return ::arrow::is_string_or_string_view(field->type()->id()); + } + if (field->name() == "offset" || field->name() == "size") { + return field->type()->id() == Type::INT64; + } + if (field->name() == "inline") { + return ::arrow::is_binary_or_binary_view(field->type()->id()); + } + return false; +} + +} // namespace + +FileExtensionType::FileExtensionType(const std::shared_ptr& storage_type) + : ExtensionType(storage_type) { + for (const auto& field : storage_type->fields()) { + if (field->name() == "uri") { + uri_ = field; + } else if (field->name() == "offset") { + offset_ = field; + } else if (field->name() == "size") { + size_ = field; + } else if (field->name() == "content_type") { + content_type_ = field; + } else if (field->name() == "checksum") { + checksum_ = field; + } else if (field->name() == "inline") { + inline_bytes_ = field; + } + } +} + +bool FileExtensionType::ExtensionEquals(const ExtensionType& other) const { + return other.extension_name() == extension_name() && + other.storage_type()->Equals(*storage_type()); +} + +Result> FileExtensionType::Deserialize( + std::shared_ptr storage_type, const std::string& serialized) const { + if (!serialized.empty()) { + return Status::Invalid("Unexpected serialized metadata: '", serialized, "'"); + } + return FileExtensionType::Make(std::move(storage_type)); +} + +std::string FileExtensionType::Serialize() const { return ""; } + +std::shared_ptr FileExtensionType::MakeArray( + std::shared_ptr data) const { + DCHECK_EQ(data->type->id(), Type::EXTENSION); + DCHECK_EQ(kFileExtensionName, + internal::checked_cast(*data->type).extension_name()); + return std::make_shared(std::move(data)); +} + +bool FileExtensionType::IsSupportedStorageType( + const std::shared_ptr& storage_type) { + if (!storage_type || storage_type->id() != Type::STRUCT || + storage_type->fields().empty()) { + return false; + } + + for (const auto& field : storage_type->fields()) { + if (!IsSupportedField(field)) { + return false; + } + } + + for (int i = 0; i < storage_type->num_fields(); ++i) { + for (int j = i + 1; j < storage_type->num_fields(); ++j) { + if (storage_type->field(i)->name() == storage_type->field(j)->name()) { + return false; + } + } + } + return true; +} + +Result> FileExtensionType::Make( + std::shared_ptr storage_type) { + if (!IsSupportedStorageType(storage_type)) { + return Status::Invalid("Invalid storage type for FileExtensionType: ", + storage_type ? storage_type->ToString() : "null"); + } + return std::make_shared(std::move(storage_type)); +} + +std::shared_ptr file(std::shared_ptr storage_type) { + return FileExtensionType::Make(std::move(storage_type)).ValueOrDie(); +} + +} // namespace arrow::extension diff --git a/cpp/src/arrow/extension/parquet_file.h b/cpp/src/arrow/extension/parquet_file.h new file mode 100644 index 00000000000..f079b17562b --- /dev/null +++ b/cpp/src/arrow/extension/parquet_file.h @@ -0,0 +1,76 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include +#include +#include + +#include "arrow/extension_type.h" +#include "arrow/util/visibility.h" + +namespace arrow::extension { + +/// \brief The extension name for the File extension type. +inline constexpr std::string_view kFileExtensionName = "parquet.file.experimental.v1"; + +class ARROW_EXPORT FileArray : public ExtensionArray { + public: + using ExtensionArray::ExtensionArray; +}; + +class ARROW_EXPORT FileExtensionType : public ExtensionType { + public: + explicit FileExtensionType(const std::shared_ptr& storage_type); + + std::string extension_name() const override { return std::string(kFileExtensionName); } + + bool ExtensionEquals(const ExtensionType& other) const override; + + Result> Deserialize( + std::shared_ptr storage_type, + const std::string& serialized_data) const override; + + std::string Serialize() const override; + + std::shared_ptr MakeArray(std::shared_ptr data) const override; + + static Result> Make(std::shared_ptr storage_type); + + static bool IsSupportedStorageType(const std::shared_ptr& storage_type); + + std::shared_ptr uri() const { return uri_; } + std::shared_ptr offset() const { return offset_; } + std::shared_ptr size() const { return size_; } + std::shared_ptr content_type() const { return content_type_; } + std::shared_ptr checksum() const { return checksum_; } + std::shared_ptr inline_bytes() const { return inline_bytes_; } + + private: + std::shared_ptr uri_; + std::shared_ptr offset_; + std::shared_ptr size_; + std::shared_ptr content_type_; + std::shared_ptr checksum_; + std::shared_ptr inline_bytes_; +}; + +/// \brief Return a FileExtensionType instance. +ARROW_EXPORT std::shared_ptr file(std::shared_ptr storage_type); + +} // namespace arrow::extension diff --git a/cpp/src/arrow/extension/parquet_file_test.cc b/cpp/src/arrow/extension/parquet_file_test.cc new file mode 100644 index 00000000000..d1147ae9e5e --- /dev/null +++ b/cpp/src/arrow/extension/parquet_file_test.cc @@ -0,0 +1,58 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include + +#include "arrow/extension/parquet_file.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/type.h" + +namespace arrow::extension { + +TEST(FileType, InvalidStorage) { + // FILE storage must be a non-empty struct. + ASSERT_NOT_OK(FileExtensionType::Make(nullptr)); + ASSERT_NOT_OK(FileExtensionType::Make(utf8())); + ASSERT_NOT_OK(FileExtensionType::Make(struct_({}))); + + // Every FILE storage field must be nullable, including when other fields are valid. + ASSERT_NOT_OK(FileExtensionType::Make( + struct_({field("uri", utf8()), field("inline", binary(), /*nullable=*/false)}))); + + // FILE storage fields must use one of the six recognized names. + ASSERT_NOT_OK(FileExtensionType::Make( + struct_({field("uri", utf8()), field("unknown", binary())}))); + + // The uri, content_type, and checksum fields must use a string storage family. + ASSERT_NOT_OK(FileExtensionType::Make(struct_({field("uri", binary())}))); + ASSERT_NOT_OK( + FileExtensionType::Make(struct_({field("content_type", large_binary())}))); + ASSERT_NOT_OK(FileExtensionType::Make(struct_({field("checksum", binary_view())}))); + + // The offset and size fields must use INT64 storage. + ASSERT_NOT_OK(FileExtensionType::Make(struct_({field("offset", int32())}))); + ASSERT_NOT_OK(FileExtensionType::Make(struct_({field("size", int32())}))); + + // The inline field must use a binary storage family. + ASSERT_NOT_OK(FileExtensionType::Make(struct_({field("inline", utf8())}))); + + // FILE storage field names must be unique. + ASSERT_NOT_OK( + FileExtensionType::Make(struct_({field("uri", utf8()), field("uri", utf8())}))); +} + +} // namespace arrow::extension diff --git a/cpp/src/arrow/meson.build b/cpp/src/arrow/meson.build index fea26ef4e45..29c29f44bd5 100644 --- a/cpp/src/arrow/meson.build +++ b/cpp/src/arrow/meson.build @@ -157,6 +157,7 @@ arrow_components = { 'extension_type.cc', 'extension/bool8.cc', 'extension/json.cc', + 'extension/parquet_file.cc', 'extension/parquet_variant.cc', 'extension/uuid.cc', 'pretty_print.cc', diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 2bdbc38b364..546a0d10fe9 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -26,6 +26,7 @@ #include #include +#include #include #include #include @@ -41,6 +42,7 @@ #include "arrow/chunked_array.h" #include "arrow/compute/api.h" #include "arrow/extension/json.h" +#include "arrow/extension/parquet_file.h" #include "arrow/io/api.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" @@ -6079,5 +6081,84 @@ TEST(TestArrowReadWrite, AllNulls) { ASSERT_TRUE(expected_table->Equals(*read_table)); } +struct FileRoundtripTestCase { + std::shared_ptr<::arrow::DataType> storage_type; + + friend std::ostream& operator<<(std::ostream& os, + const FileRoundtripTestCase& test_case) { + return os << "{storage_type=" << test_case.storage_type->ToString() << "}"; + } +}; + +class TestArrowReadWriteFileType + : public ::testing::TestWithParam { + protected: + ::arrow::Result> OpenReader( + const std::shared_ptr<::arrow::Buffer>& buffer, + const std::shared_ptr<::arrow::DataType>& file_type) { + extension_guard_.emplace(::arrow::DataTypeVector{file_type}); + ArrowReaderProperties reader_properties; + reader_properties.set_arrow_extensions_enabled(true); + std::unique_ptr reader; + FileReaderBuilder builder; + RETURN_NOT_OK(builder.Open(std::make_shared(buffer))); + RETURN_NOT_OK(builder.properties(reader_properties)->Build(&reader)); + return reader; + } + + private: + std::optional<::arrow::ExtensionTypeGuard> extension_guard_; +}; + +TEST_P(TestArrowReadWriteFileType, FileExtensionRoundtrip) { + const auto& test_case = GetParam(); + const auto& storage_type = test_case.storage_type; + auto file_type = ::arrow::extension::file(storage_type); + auto storage_array = ::arrow::ArrayFromJSON( + storage_type, + R"([{"uri":"u","offset":null,"inline":null},{"uri":null,"offset":null,"inline":"x"},null])"); + auto file_array = ::arrow::ExtensionType::WrapArray(file_type, storage_array); + auto input = + ::arrow::Table::Make(::arrow::schema({::arrow::field("file", file_type)}), + {std::make_shared<::arrow::ChunkedArray>(file_array)}); + + auto arrow_writer_properties = ArrowWriterProperties::Builder().store_schema()->build(); + auto sink = CreateOutputStream(); + ASSERT_OK(WriteTable(*input, ::arrow::default_memory_pool(), sink, input->num_rows(), + default_writer_properties(), arrow_writer_properties)); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + + ASSERT_OK_AND_ASSIGN(auto reader, OpenReader(buffer, file_type)); + + ASSERT_OK_AND_ASSIGN(auto full, reader->ReadTable()); + ::arrow::AssertTypeEqual(file_type, full->schema()->field(0)->type()); + ASSERT_TRUE(full->Equals(*input)); + + ASSERT_OK_AND_ASSIGN(auto partial, reader->ReadTable(std::vector{0})); + ASSERT_EQ(partial->schema()->field(0)->type()->id(), ::arrow::Type::STRUCT); + ASSERT_EQ(partial->schema()->field(0)->type()->num_fields(), 1); + ASSERT_EQ(partial->schema()->field(0)->type()->field(0)->name(), "uri"); + auto storage_struct = + ::arrow::internal::checked_pointer_cast(storage_array); + ASSERT_OK_AND_ASSIGN(auto expected_partial, + ::arrow::StructArray::Make({storage_struct->field(0)}, {"uri"}, + storage_struct->null_bitmap(), + storage_struct->null_count())); + ASSERT_TRUE(partial->column(0)->Equals( + std::make_shared<::arrow::ChunkedArray>(std::move(expected_partial)))); +} + +INSTANTIATE_TEST_SUITE_P( + FileTypes, TestArrowReadWriteFileType, + ::testing::Values( + FileRoundtripTestCase{ + ::arrow::struct_({::arrow::field("uri", ::arrow::utf8()), + ::arrow::field("offset", ::arrow::int64()), + ::arrow::field("inline", ::arrow::binary())})}, + FileRoundtripTestCase{ + ::arrow::struct_({::arrow::field("uri", ::arrow::utf8_view()), + ::arrow::field("offset", ::arrow::int64()), + ::arrow::field("inline", ::arrow::large_binary())})})); + } // namespace arrow } // namespace parquet diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index 27c302fe0d4..8a0abbc0a6d 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include "gmock/gmock-matchers.h" @@ -33,6 +34,7 @@ #include "arrow/array.h" #include "arrow/extension/json.h" +#include "arrow/extension/parquet_file.h" #include "arrow/extension/parquet_variant.h" #include "arrow/extension/uuid.h" #include "arrow/ipc/writer.h" @@ -1061,6 +1063,52 @@ TEST_F(TestConvertParquetSchema, ParquetVariant) { } } +TEST_F(TestConvertParquetSchema, ParquetFile) { + std::vector parquet_fields; + parquet_fields.push_back(PrimitiveNode::Make( + "uri", Repetition::OPTIONAL, LogicalType::String(), ParquetType::BYTE_ARRAY)); + parquet_fields.push_back(PrimitiveNode::Make("offset", Repetition::OPTIONAL, + LogicalType::None(), ParquetType::INT64)); + parquet_fields.push_back(PrimitiveNode::Make( + "inline", Repetition::OPTIONAL, LogicalType::None(), ParquetType::BYTE_ARRAY)); + auto file = + GroupNode::Make("file", Repetition::OPTIONAL, parquet_fields, LogicalType::File()); + + auto storage = ::arrow::struct_({::arrow::field("uri", ::arrow::utf8()), + ::arrow::field("offset", ::arrow::int64()), + ::arrow::field("inline", ::arrow::binary())}); + auto file_extension = ::arrow::extension::file(storage); + + std::shared_ptr<::arrow::KeyValueMetadata> metadata; + auto stored_schema = ::arrow::schema({::arrow::field("file", file_extension)}); + ASSERT_OK(ArrowSchemaToParquetMetadata(stored_schema, metadata)); + + auto check_file_schema = + [&]( + bool enable_extensions, const std::shared_ptr& metadata, + const std::shared_ptr<::arrow::DataType>& expected_type, bool check_metadata) { + std::optional<::arrow::ExtensionTypeGuard> guard; + if constexpr (RegisterExtension) { + guard.emplace(::arrow::DataTypeVector{file_extension}); + } + + ArrowReaderProperties props; + props.set_arrow_extensions_enabled(enable_extensions); + ASSERT_OK(ConvertSchema({file}, metadata, props)); + CheckFlatSchema(::arrow::schema({::arrow::field("file", expected_type)}), + check_metadata); + }; // NOLINT(readability/braces) + + check_file_schema.operator()(true, std::shared_ptr{}, + file_extension, true); + check_file_schema.operator()(false, std::shared_ptr{}, + storage, true); + check_file_schema.operator()(true, std::shared_ptr{}, + storage, true); + check_file_schema.operator()(false, metadata, file_extension, true); + check_file_schema.operator()(true, metadata, storage, false); +} + TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) { std::vector parquet_fields; parquet_fields.push_back(PrimitiveNode::Make( @@ -1838,6 +1886,29 @@ TEST_F(TestConvertArrowSchema, ParquetFlatDecimals) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); } +TEST_F(TestConvertArrowSchema, ParquetFile) { + auto storage = ::arrow::struct_({::arrow::field("inline", ::arrow::binary()), + ::arrow::field("uri", ::arrow::utf8()), + ::arrow::field("size", ::arrow::int64())}); + auto file_type = ::arrow::extension::file(storage); + auto arrow_fields = ::arrow::FieldVector{::arrow::field("file", file_type)}; + + auto expected = GroupNode::Make( + "file", Repetition::OPTIONAL, + {PrimitiveNode::Make("inline", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY), + PrimitiveNode::Make("uri", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY, + ConvertedType::UTF8), + PrimitiveNode::Make("size", Repetition::OPTIONAL, ParquetType::INT64)}, + LogicalType::File()); + + ASSERT_OK(ConvertSchema(arrow_fields)); + ASSERT_EQ(result_schema_->group_node()->field_count(), 1); + ASSERT_TRUE(result_schema_->group_node()->field(0)->Equals(expected.get())); + + ASSERT_OK(ConvertSchema({::arrow::field("file", storage)})); + ASSERT_FALSE(result_schema_->group_node()->field(0)->logical_type()->is_file()); +} + TEST_F(TestConvertArrowSchema, ParquetTimeAdjustedToUTC) { // Verify Parquet Time types have the appropriate isAdjustedToUTC value, depending // on the return value of ArrowWriterProperties::write_time_adjusted_to_utc() diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index eca83e8576d..6989ed85f40 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -29,6 +29,7 @@ #include "arrow/array.h" // IWYU pragma: keep #include "arrow/array/concatenate.h" #include "arrow/buffer.h" +#include "arrow/extension/parquet_file.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" #include "arrow/memory_pool.h" @@ -73,6 +74,7 @@ using arrow::Status; using arrow::StructArray; using arrow::Table; using arrow::TimestampArray; +using arrow::extension::kFileExtensionName; using arrow::internal::checked_cast; using arrow::internal::Iota; @@ -942,12 +944,15 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& arrow_f auto type_id = arrow_field->type()->id(); if (type_id == ::arrow::Type::EXTENSION) { - auto storage_field = arrow_field->WithType( - checked_cast(*arrow_field->type()).storage_type()); + const auto& extension_type = checked_cast(*arrow_field->type()); + auto storage_field = arrow_field->WithType(extension_type.storage_type()); RETURN_NOT_OK(GetReader(field, storage_field, ctx, out)); if (*out) { auto storage_type = (*out)->field()->type(); if (!storage_type->Equals(storage_field->type())) { + if (extension_type.extension_name() == kFileExtensionName) { + return Status::OK(); + } return Status::Invalid( "Due to column pruning only part of an extension's storage type was loaded. " "An extension type cannot be created without all of its fields"); diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index bc4de6c39b5..465142a70f1 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -23,6 +23,7 @@ #include #include "arrow/extension/json.h" +#include "arrow/extension/parquet_file.h" #include "arrow/extension/parquet_variant.h" #include "arrow/extension/uuid.h" #include "arrow/extension_type.h" @@ -50,6 +51,7 @@ using arrow::Field; using arrow::FieldVector; using arrow::KeyValueMetadata; using arrow::Status; +using arrow::extension::kFileExtensionName; using arrow::internal::checked_cast; using arrow::internal::ToChars; @@ -151,6 +153,25 @@ Status VariantToNode( return Status::OK(); } +Status FileToNode(const std::shared_ptr<::arrow::extension::FileExtensionType>& type, + const std::string& name, bool nullable, int field_id, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, NodePtr* out) { + std::vector children; + children.reserve(type->storage_type()->num_fields()); + for (const auto& child : type->storage_type()->fields()) { + NodePtr child_node; + RETURN_NOT_OK( + FieldToNode(child->name(), child, properties, arrow_properties, &child_node)); + children.push_back(std::move(child_node)); + } + + auto file_node = GroupNode::Make(name, RepetitionFromNullable(nullable), children, + LogicalType::File(), field_id); + *out = std::move(file_node); + return Status::OK(); +} + Status StructToNode(const std::shared_ptr<::arrow::StructType>& type, const std::string& name, bool nullable, int field_id, const WriterProperties& properties, @@ -497,6 +518,11 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, return VariantToNode(variant_type, name, field->nullable(), field_id, properties, arrow_properties, out); + } else if (ext_type->extension_name() == kFileExtensionName) { + auto file_type = std::static_pointer_cast<::arrow::extension::FileExtensionType>( + field->type()); + return FileToNode(file_type, name, field->nullable(), field_id, properties, + arrow_properties, out); } std::shared_ptr<::arrow::Field> storage_field = ::arrow::field( @@ -602,13 +628,21 @@ Status GroupToStruct(const GroupNode& node, LevelInfo current_levels, arrow_fields.push_back(out->children[i].field); } auto struct_type = ::arrow::struct_(arrow_fields); - if (ctx->properties.get_arrow_extensions_enabled() && - node.logical_type()->is_variant()) { - auto extension_type = ::arrow::GetExtensionType("arrow.parquet.variant"); - if (extension_type) { - ARROW_ASSIGN_OR_RAISE( - struct_type, - extension_type->Deserialize(std::move(struct_type), /*serialized_data=*/"")); + if (ctx->properties.get_arrow_extensions_enabled()) { + if (node.logical_type()->is_variant()) { + auto extension_type = ::arrow::GetExtensionType("arrow.parquet.variant"); + if (extension_type) { + ARROW_ASSIGN_OR_RAISE( + struct_type, + extension_type->Deserialize(std::move(struct_type), /*serialized_data=*/"")); + } + } else if (node.logical_type()->is_file()) { + auto extension_type = ::arrow::GetExtensionType(std::string(kFileExtensionName)); + if (extension_type) { + ARROW_ASSIGN_OR_RAISE( + struct_type, + extension_type->Deserialize(std::move(struct_type), /*serialized_data=*/"")); + } } } out->field = ::arrow::field(node.name(), struct_type, node.is_optional(), @@ -1042,7 +1076,8 @@ std::function(FieldVector)> GetNestedFactory( } Result ApplyOriginalStorageMetadata(const Field& origin_field, - SchemaField* inferred) { + SchemaField* inferred, + bool match_children_by_name = false) { bool modified = false; auto& origin_type = origin_field.type(); @@ -1059,9 +1094,17 @@ Result ApplyOriginalStorageMetadata(const Field& origin_field, // Apply original metadata recursively to children for (int i = 0; i < inferred_type->num_fields(); ++i) { + std::shared_ptr<::arrow::Field> origin_child; + if (match_children_by_name) { + origin_child = checked_cast(*origin_type) + .GetFieldByName(inferred_type->field(i)->name()); + DCHECK(origin_child != nullptr); + } else { + origin_child = origin_type->field(i); + } ARROW_ASSIGN_OR_RAISE( const bool child_modified, - ApplyOriginalMetadata(*origin_type->field(i), &inferred->children[i])); + ApplyOriginalMetadata(*origin_child, &inferred->children[i])); modified |= child_modified; } if (modified) { @@ -1150,6 +1193,21 @@ Result ApplyOriginalStorageMetadata(const Field& origin_field, return modified; } +bool FileStorageTypesCompatible(const std::shared_ptr<::arrow::DataType>& origin_type, + const std::shared_ptr<::arrow::DataType>& inferred_type) { + if (origin_type->num_fields() != inferred_type->num_fields()) { + return false; + } + const auto& inferred_struct_type = + checked_cast(*inferred_type); + for (const auto& origin_field : origin_type->fields()) { + if (inferred_struct_type.GetFieldByName(origin_field->name()) == nullptr) { + return false; + } + } + return true; +} + Result ApplyOriginalMetadata(const Field& origin_field, SchemaField* inferred) { bool modified = false; @@ -1161,28 +1219,42 @@ Result ApplyOriginalMetadata(const Field& origin_field, SchemaField* infer if (origin_type->id() == ::arrow::Type::EXTENSION) { const auto& origin_extension_type = checked_cast(*origin_type); + std::string origin_extension_name = origin_extension_type.extension_name(); + + // Whether or not the inferred type is also an extension type. This can occur when + // arrow_extensions_enabled is true in the ArrowReaderProperties. Extension types + // are not currently inferred for any other reason. + bool arrow_extension_inferred = + inferred->field->type()->id() == ::arrow::Type::EXTENSION; + + bool restore_file_extension = false; + if (origin_extension_name == kFileExtensionName && arrow_extension_inferred) { + const auto& inferred_extension_type = + checked_cast(*inferred->field->type()); + if (FileStorageTypesCompatible(origin_extension_type.storage_type(), + inferred_extension_type.storage_type())) { + inferred->field = + inferred->field->WithType(inferred_extension_type.storage_type()); + restore_file_extension = true; + } + } // (Recursively) Apply the original storage metadata from the original storage field // This applies extension types to child elements, if any. auto origin_storage_field = origin_field.WithType(origin_extension_type.storage_type()); - RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred)); + RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred, + restore_file_extension)); // Use the inferred type after child updates for below checks to see if // we can restore an extension type on the output. const auto& inferred_type = inferred->field->type(); - // Whether or not the inferred type is also an extension type. This can occur when - // arrow_extensions_enabled is true in the ArrowReaderProperties. Extension types - // are not currently inferred for any other reason. - bool arrow_extension_inferred = inferred_type->id() == ::arrow::Type::EXTENSION; - // Check if the inferred storage type is compatible with the extension type // we're hoping to apply. We assume that if an extension type was inferred // that it was constructed with a valid storage type. Otherwise, we check with // extension types that we know about for valid storage, falling back to // storage type equality for extension types that we don't know about. - std::string origin_extension_name = origin_extension_type.extension_name(); bool extension_supports_inferred_storage; if (origin_extension_name == "arrow.json") { @@ -1198,6 +1270,8 @@ Result ApplyOriginalMetadata(const Field& origin_field, SchemaField* infer extension_supports_inferred_storage = arrow_extension_inferred || ::arrow::extension::VariantExtensionType::IsSupportedStorageType(inferred_type); + } else if (origin_extension_name == kFileExtensionName && restore_file_extension) { + extension_supports_inferred_storage = true; } else { extension_supports_inferred_storage = origin_extension_type.storage_type()->Equals(*inferred_type); @@ -1207,7 +1281,14 @@ Result ApplyOriginalMetadata(const Field& origin_field, SchemaField* infer // the Arrow storage type we would otherwise return, we restore the extension // type to the output. if (extension_supports_inferred_storage) { - inferred->field = inferred->field->WithType(origin_type); + if (restore_file_extension) { + ARROW_ASSIGN_OR_RAISE(auto restored_type, + origin_extension_type.Deserialize( + inferred_type, origin_extension_type.Serialize())); + inferred->field = inferred->field->WithType(std::move(restored_type)); + } else { + inferred->field = inferred->field->WithType(origin_type); + } } modified = true; diff --git a/cpp/src/parquet/schema_test.cc b/cpp/src/parquet/schema_test.cc index 6c8e6366adf..e4a7ff11934 100644 --- a/cpp/src/parquet/schema_test.cc +++ b/cpp/src/parquet/schema_test.cc @@ -1224,6 +1224,9 @@ TEST(TestLogicalTypeConstruction, NewTypeIncompatibility) { auto check_is_variant = [](const std::shared_ptr& logical_type) { return logical_type->is_variant(); }; + auto check_is_file = [](const std::shared_ptr& logical_type) { + return logical_type->is_file(); + }; auto check_is_null = [](const std::shared_ptr& logical_type) { return logical_type->is_null(); }; @@ -1238,6 +1241,7 @@ TEST(TestLogicalTypeConstruction, NewTypeIncompatibility) { {LogicalType::UUID(), check_is_UUID}, {LogicalType::Float16(), check_is_float16}, {LogicalType::Variant(), check_is_variant}, + {LogicalType::File(), check_is_file}, {LogicalType::Null(), check_is_null}, {LogicalType::Time(false, LogicalType::TimeUnit::MILLIS), check_is_time}, {LogicalType::Time(false, LogicalType::TimeUnit::MICROS), check_is_time}, @@ -1323,6 +1327,7 @@ TEST(TestLogicalTypeOperation, LogicalTypeProperties) { {UUIDLogicalType::Make(), false, true, true}, {Float16LogicalType::Make(), false, true, true}, {VariantLogicalType::Make(), true, true, true}, + {FileLogicalType::Make(), true, true, true}, {NoLogicalType::Make(), false, false, true}, }; @@ -1653,6 +1658,7 @@ TEST(TestLogicalTypeOperation, LogicalTypeRepresentation) { R"({"Type":"Geography","crs":"srid:1234","algorithm":"karney"})"}, {LogicalType::Variant(), "Variant(1)", R"({"Type":"Variant","SpecVersion":1})"}, {LogicalType::Variant(2), "Variant(2)", R"({"Type":"Variant","SpecVersion":2})"}, + {LogicalType::File(), "File", R"({"Type":"File"})"}, {LogicalType::None(), "None", R"({"Type":"None"})"}, }; @@ -1706,6 +1712,7 @@ TEST(TestLogicalTypeOperation, LogicalTypeSortOrder) { {LogicalType::Geometry(), SortOrder::UNKNOWN}, {LogicalType::Geography(), SortOrder::UNKNOWN}, {LogicalType::Variant(), SortOrder::UNKNOWN}, + {LogicalType::File(), SortOrder::UNKNOWN}, {LogicalType::None(), SortOrder::UNKNOWN}}; for (const ExpectedSortOrder& c : cases) { @@ -1873,6 +1880,10 @@ TEST(TestSchemaNodeCreation, FactoryExceptions) { VariantLogicalType::Make(), Type::FIXED_LEN_BYTE_ARRAY, 2)); + // Incompatible primitive type ... + ASSERT_ANY_THROW(PrimitiveNode::Make("file", Repetition::REQUIRED, + FileLogicalType::Make(), Type::DOUBLE)); + // Non-positive length argument for fixed length binary ... ASSERT_ANY_THROW(PrimitiveNode::Make("negative_length", Repetition::REQUIRED, NoLogicalType::Make(), Type::FIXED_LEN_BYTE_ARRAY, @@ -2426,6 +2437,7 @@ TEST(TestLogicalTypeSerialization, Roundtrips) { ConfirmGroupNodeRoundtrip("map", LogicalType::Map()); ConfirmGroupNodeRoundtrip("list", LogicalType::List()); ConfirmGroupNodeRoundtrip("variant", LogicalType::Variant()); + ConfirmGroupNodeRoundtrip("file", LogicalType::File()); } TEST(TestLogicalTypeSerialization, VariantSpecificationVersion) { diff --git a/cpp/src/parquet/types.cc b/cpp/src/parquet/types.cc index fda5e319e05..9ac9b17ed40 100644 --- a/cpp/src/parquet/types.cc +++ b/cpp/src/parquet/types.cc @@ -611,6 +611,8 @@ std::shared_ptr LogicalType::FromThrift( } return VariantLogicalType::Make(spec_version); + } else if (type.__isset.FILE) { + return FileLogicalType::Make(); } else { // Sentinel type for one we do not recognize return UndefinedLogicalType::Make(); @@ -682,6 +684,8 @@ std::shared_ptr LogicalType::Variant(int8_t spec_version) { return VariantLogicalType::Make(spec_version); } +std::shared_ptr LogicalType::File() { return FileLogicalType::Make(); } + std::shared_ptr LogicalType::None() { return NoLogicalType::Make(); } /* @@ -770,6 +774,7 @@ class LogicalType::Impl { class Geometry; class Geography; class Variant; + class File; class No; class Undefined; @@ -851,6 +856,7 @@ bool LogicalType::is_geography() const { bool LogicalType::is_variant() const { return impl_->type() == LogicalType::Type::VARIANT; } +bool LogicalType::is_file() const { return impl_->type() == LogicalType::Type::FILE; } bool LogicalType::is_none() const { return impl_->type() == LogicalType::Type::NONE; } bool LogicalType::is_valid() const { return impl_->type() != LogicalType::Type::UNDEFINED; @@ -859,7 +865,8 @@ bool LogicalType::is_invalid() const { return !is_valid(); } bool LogicalType::is_nested() const { return impl_->type() == LogicalType::Type::LIST || impl_->type() == LogicalType::Type::MAP || - impl_->type() == LogicalType::Type::VARIANT; + impl_->type() == LogicalType::Type::VARIANT || + impl_->type() == LogicalType::Type::FILE; } bool LogicalType::is_nonnested() const { return !is_nested(); } bool LogicalType::is_serialized() const { return impl_->is_serialized(); } @@ -2031,6 +2038,20 @@ std::shared_ptr VariantLogicalType::Make(const int8_t spec_ve return logical_type; } +class LogicalType::Impl::File final : public LogicalType::Impl::Incompatible, + public LogicalType::Impl::Inapplicable { + public: + friend class FileLogicalType; + + OVERRIDE_TOSTRING(File) + OVERRIDE_TOTHRIFT(FileType, FILE) + + private: + File() : LogicalType::Impl(LogicalType::Type::FILE, SortOrder::UNKNOWN) {} +}; + +GENERATE_MAKE(File) + class LogicalType::Impl::No final : public LogicalType::Impl::SimpleCompatible, public LogicalType::Impl::UniversalApplicable { public: diff --git a/cpp/src/parquet/types.h b/cpp/src/parquet/types.h index 687353aa9bc..affbc4b9224 100644 --- a/cpp/src/parquet/types.h +++ b/cpp/src/parquet/types.h @@ -162,6 +162,7 @@ class PARQUET_EXPORT LogicalType { GEOMETRY, GEOGRAPHY, VARIANT, + FILE, NONE // Not a real logical type; should always be last element }; }; @@ -230,6 +231,7 @@ class PARQUET_EXPORT LogicalType { static std::shared_ptr Float16(); static std::shared_ptr Variant( int8_t specVersion = kVariantSpecVersion); + static std::shared_ptr File(); static std::shared_ptr Geometry(std::string crs = ""); @@ -293,6 +295,7 @@ class PARQUET_EXPORT LogicalType { bool is_geometry() const; bool is_geography() const; bool is_variant() const; + bool is_file() const; bool is_none() const; /// \brief Return true if this logical type is of a known type. bool is_valid() const; @@ -509,6 +512,15 @@ class PARQUET_EXPORT VariantLogicalType : public LogicalType { VariantLogicalType() = default; }; +/// \brief Allowed for group nodes only. +class PARQUET_EXPORT FileLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + FileLogicalType() = default; +}; + /// \brief Allowed for any physical type. class PARQUET_EXPORT NoLogicalType : public LogicalType { public: