feat(shared_shredding): add read inte test for shared_shredding (alibaba#382)

Author: lszskye

src/paimon/CMakeLists.txt

@@ -142,7 +142,7 @@ set(PAIMON_COMMON_SRCS
     common/data/shredding/map_shared_shredding_context.cpp
     common/data/shredding/map_shared_shredding_batch_converter.cpp
     common/data/shredding/map_shared_shredding_column_allocator.cpp
-    common/data/shredding/shared_shredding_file_reader.cpp
+    common/data/shredding/map_shared_shredding_file_reader.cpp
     common/utils/delta_varint_compressor.cpp
     common/utils/fields_comparator.cpp
     common/utils/path_util.cpp
@@ -550,7 +550,7 @@ if(PAIMON_BUILD_TESTS)
                     common/data/shredding/map_shared_shredding_column_allocator_test.cpp
                     common/data/shredding/map_shared_shredding_field_dict_test.cpp
                     common/data/shredding/map_shared_shredding_context_test.cpp
-                    common/data/shredding/shared_shredding_file_reader_test.cpp
+                    common/data/shredding/map_shared_shredding_file_reader_test.cpp
                     STATIC_LINK_LIBS
                     paimon_shared
                     test_utils_static

…redding/shared_shredding_file_reader.cpp …ing/map_shared_shredding_file_reader.cppsrc/paimon/common/data/shredding/shared_shredding_file_reader.cpp renamed to src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp src/paimon/common/data/shredding/shared_shredding_file_reader.cpp renamed to src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp

@@ -18,7 +18,9 @@
 
 #include <map>
 #include <memory>
+#include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "arrow/api.h"
@@ -28,10 +30,22 @@
 
 namespace paimon {
 
-class SharedShreddingFileReader : public FileBatchReader {
+class MapSharedShreddingFileReader : public FileBatchReader {
  public:
-    static Result<std::unique_ptr<SharedShreddingFileReader>> Create(
-        std::unique_ptr<FileBatchReader>&& reader, const std::shared_ptr<MemoryPool>& pool);
+    struct SharedShreddingContext {
+        SharedShreddingContext(const MapSharedShreddingFieldMeta& _meta,
+                               const std::vector<std::string>& _selected_keys,
+                               const std::shared_ptr<arrow::MapType>& _map_type)
+            : meta(_meta), selected_keys(_selected_keys), map_type(_map_type) {}
+        MapSharedShreddingFieldMeta meta;
+        std::vector<std::string> selected_keys;
+        std::shared_ptr<arrow::MapType> map_type;
+    };
+
+    MapSharedShreddingFileReader(
+        std::unique_ptr<FileBatchReader>&& reader,
+        std::map<std::string, SharedShreddingContext>&& shared_shredding_name_to_context,
+        const std::shared_ptr<MemoryPool>& pool);
 
     Result<std::unique_ptr<::ArrowSchema>> GetFileSchema() const override;
 
@@ -53,11 +67,6 @@ class SharedShreddingFileReader : public FileBatchReader {
     bool SupportPreciseBitmapSelection() const override;
 
  private:
-    SharedShreddingFileReader(
-        std::unique_ptr<FileBatchReader>&& reader,
-        const std::map<std::string, MapSharedShreddingFieldMeta>& shared_shredding_name_to_meta,
-        const std::shared_ptr<MemoryPool>& pool);
-
     Result<std::shared_ptr<arrow::Array>> RebuildLogicalMapArray(
         const std::shared_ptr<arrow::Field>& physical_field,
         const std::shared_ptr<arrow::StructArray>& physical_struct_array) const;
@@ -76,9 +85,7 @@ class SharedShreddingFileReader : public FileBatchReader {
  private:
     std::shared_ptr<arrow::MemoryPool> arrow_pool_;
     std::unique_ptr<FileBatchReader> reader_;
-    std::map<std::string, MapSharedShreddingFieldMeta> shared_shredding_name_to_meta_;
-    std::map<std::string, std::vector<std::string>> shared_shredding_name_to_selected_keys_;
-    std::map<std::string, std::shared_ptr<arrow::MapType>> shared_shredding_name_to_map_type_;
+    std::map<std::string, SharedShreddingContext> shared_shredding_name_to_context_;
 };
 
 }  // namespace paimon

…shredding/shared_shredding_file_reader.h …dding/map_shared_shredding_file_reader.hsrc/paimon/common/data/shredding/shared_shredding_file_reader.h renamed to src/paimon/common/data/shredding/map_shared_shredding_file_reader.h src/paimon/common/data/shredding/shared_shredding_file_reader.h renamed to src/paimon/common/data/shredding/map_shared_shredding_file_reader.h

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "paimon/common/data/shredding/shared_shredding_file_reader.h"
+#include "paimon/common/data/shredding/map_shared_shredding_file_reader.h"
 
 #include <map>
 #include <memory>
@@ -47,7 +47,7 @@
 #include "paimon/testing/utils/testharness.h"
 
 namespace paimon::test {
-class SharedShreddingFileReaderTest : public ::testing::Test {
+class MapSharedShreddingFileReaderTest : public ::testing::Test {
  public:
     void SetUp() override {
         pool_ = GetDefaultPool();
@@ -94,9 +94,56 @@ class SharedShreddingFileReaderTest : public ::testing::Test {
             .ValueOrDie();
     }
 
-    std::unique_ptr<SharedShreddingFileReader> CreateReader(
+    std::unique_ptr<MapSharedShreddingFileReader> WrapReader(
+        std::unique_ptr<FileBatchReader>&& reader,
+        const std::optional<std::string>& selected_keys_str = std::nullopt) const {
+        EXPECT_OK_AND_ASSIGN(auto c_file_schema, reader->GetFileSchema());
+        auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie();
+        std::map<std::string, MapSharedShreddingFileReader::SharedShreddingContext>
+            shared_shredding_name_to_context;
+        for (const auto& field : file_schema->fields()) {
+            auto metadata = std::const_pointer_cast<arrow::KeyValueMetadata>(field->metadata());
+            if (!MapSharedShreddingUtils::HasShreddingMetadata(metadata)) {
+                continue;
+            }
+            EXPECT_OK_AND_ASSIGN(auto meta,
+                                 MapSharedShreddingUtils::DeserializeMetadata(
+                                     metadata, MapSharedShreddingDefine::kDefaultDictCompression));
+            auto physical_type =
+                arrow::internal::checked_pointer_cast<arrow::StructType>(field->type());
+            std::shared_ptr<arrow::Field> item_field;
+            for (const auto& child : physical_type->fields()) {
+                if (child->name() != MapSharedShreddingDefine::kFieldMapping &&
+                    child->name() != MapSharedShreddingDefine::kOverflow) {
+                    item_field = child;
+                    break;
+                }
+            }
+            EXPECT_TRUE(item_field);
+            auto map_type = arrow::internal::checked_pointer_cast<arrow::MapType>(arrow::map(
+                arrow::utf8(), arrow::field("value", item_field->type(), item_field->nullable())));
+            std::vector<std::string> selected_keys;
+            if (selected_keys_str.has_value()) {
+                selected_keys = StringUtils::Split(selected_keys_str.value(), ",",
+                                                   /*ignore_empty=*/false);
+            } else {
+                selected_keys.reserve(meta.name_to_id.size());
+                for (const auto& [key_name, _] : meta.name_to_id) {
+                    selected_keys.push_back(key_name);
+                }
+            }
+            shared_shredding_name_to_context.emplace(
+                field->name(), MapSharedShreddingFileReader::SharedShreddingContext(
+                                   meta, selected_keys, map_type));
+        }
+        return std::make_unique<MapSharedShreddingFileReader>(
+            std::move(reader), std::move(shared_shredding_name_to_context), pool_);
+    }
+
+    std::unique_ptr<MapSharedShreddingFileReader> CreateReader(
         std::shared_ptr<arrow::Array> physical_array = nullptr,
-        std::shared_ptr<arrow::Schema> physical_schema = nullptr) const {
+        std::shared_ptr<arrow::Schema> physical_schema = nullptr,
+        const std::optional<std::string>& selected_keys = std::nullopt) const {
         if (!physical_schema) {
             physical_schema = PhysicalSchemaWithMetadata();
         }
@@ -106,9 +153,7 @@ class SharedShreddingFileReaderTest : public ::testing::Test {
         auto mock_reader = std::make_unique<MockFileBatchReader>(
             physical_array, arrow::struct_(physical_schema->fields()), /*read_batch_size=*/10);
         mock_reader->EnableRandomizeBatchSize(false);
-        EXPECT_OK_AND_ASSIGN(auto shared_shredding_reader,
-                             SharedShreddingFileReader::Create(std::move(mock_reader), pool_));
-        return shared_shredding_reader;
+        return WrapReader(std::move(mock_reader), selected_keys);
     }
 
     std::shared_ptr<arrow::Schema> ReadSchema(
@@ -196,7 +241,7 @@ class SharedShreddingFileReaderTest : public ::testing::Test {
     };
 };
 
-TEST_F(SharedShreddingFileReaderTest, TestGetFileSchemaReturnsLogicalMapSchema) {
+TEST_F(MapSharedShreddingFileReaderTest, TestGetFileSchemaReturnsLogicalMapSchema) {
     auto reader = CreateReader();
 
     ASSERT_OK_AND_ASSIGN(auto c_schema, reader->GetFileSchema());
@@ -209,8 +254,9 @@ TEST_F(SharedShreddingFileReaderTest, TestGetFileSchemaReturnsLogicalMapSchema)
     ASSERT_FALSE(schema->field(1)->HasMetadata());
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestAllExistSelectedKeysWithoutOverflow) {
-    auto reader = CreateReader();
+TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithoutOverflow) {
+    auto reader = CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr,
+                               /*selected_keys=*/"b");
     auto read_schema = ExportSchema(ReadSchema("b"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
@@ -229,8 +275,9 @@ TEST_F(SharedShreddingFileReaderTest, TestAllExistSelectedKeysWithoutOverflow) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) {
-    auto reader = CreateReader();
+TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) {
+    auto reader = CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr,
+                               /*selected_keys=*/"a,c");
     auto read_schema = ExportSchema(ReadSchema("a,c"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
@@ -249,8 +296,9 @@ TEST_F(SharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestPartialExistSelectedKeys) {
-    auto reader = CreateReader();
+TEST_F(MapSharedShreddingFileReaderTest, TestPartialExistSelectedKeys) {
+    auto reader = CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr,
+                               /*selected_keys=*/"a,c,missing");
     auto read_schema = ExportSchema(ReadSchema("a,c,missing"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
@@ -270,15 +318,7 @@ TEST_F(SharedShreddingFileReaderTest, TestPartialExistSelectedKeys) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestDuplicatedSelectedKeys) {
-    auto reader = CreateReader();
-    auto read_schema = ExportSchema(ReadSchema("a,c,a"));
-    ASSERT_NOK_WITH_MSG(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
-                                              /*selection_bitmap=*/std::nullopt),
-                        "duplicate key [a] in paimon.map.selected-keys for field tags");
-}
-
-TEST_F(SharedShreddingFileReaderTest, TestMissingSelectedKeysReadsWholeMap) {
+TEST_F(MapSharedShreddingFileReaderTest, TestMissingSelectedKeysReadsWholeMap) {
     auto reader = CreateReader();
     auto read_schema = ExportSchema(ReadSchema(std::nullopt));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
@@ -298,7 +338,7 @@ TEST_F(SharedShreddingFileReaderTest, TestMissingSelectedKeysReadsWholeMap) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestSpecialSelectedKeys) {
+TEST_F(MapSharedShreddingFileReaderTest, TestSpecialSelectedKeys) {
     MapSharedShreddingFieldMeta meta;
     meta.name_to_id = {{"", 0}, {"     ", 1}, {".", 2}, {"a", 3}};
     meta.field_to_columns = {{0, {0}}, {1, {1}}, {2, {0}}, {3, {1}}};
@@ -315,7 +355,7 @@ TEST_F(SharedShreddingFileReaderTest, TestSpecialSelectedKeys) {
             .ValueOrDie();
 
     auto assert_read = [&](const std::string& selected_keys, const std::string& expected_json) {
-        auto reader = CreateReader(physical_array, physical_schema);
+        auto reader = CreateReader(physical_array, physical_schema, selected_keys);
         auto read_schema = ExportSchema(ReadSchema(selected_keys));
         ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                         /*selection_bitmap=*/std::nullopt));
@@ -350,18 +390,9 @@ TEST_F(SharedShreddingFileReaderTest, TestSpecialSelectedKeys) {
     ])");
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestSpecialSelectedKeysWithDuplicatedEmptyKey) {
-    for (const auto& selected_keys : {",", ",,"}) {
-        auto reader = CreateReader();
-        auto read_schema = ExportSchema(ReadSchema(selected_keys));
-        ASSERT_NOK_WITH_MSG(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
-                                                  /*selection_bitmap=*/std::nullopt),
-                            "duplicate key [] in paimon.map.selected-keys for field tags");
-    }
-}
-
-TEST_F(SharedShreddingFileReaderTest, TestUnknownSelectedKeyReturnsEmptyMap) {
-    auto reader = CreateReader();
+TEST_F(MapSharedShreddingFileReaderTest, TestUnknownSelectedKeyReturnsEmptyMap) {
+    auto reader = CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr,
+                               /*selected_keys=*/"missing");
     auto read_schema = ExportSchema(ReadSchema("missing"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
@@ -381,39 +412,39 @@ TEST_F(SharedShreddingFileReaderTest, TestUnknownSelectedKeyReturnsEmptyMap) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestInvalidNullFieldMappingField) {
+TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingField) {
     auto physical_schema = PhysicalSchemaWithMetadata();
     std::string json = R"([
         [1, [null, 10, null, null]]
     ])";
     auto physical_array =
         arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(physical_schema->fields()), json)
             .ValueOrDie();
-    auto reader = CreateReader(physical_array, physical_schema);
+    auto reader = CreateReader(physical_array, physical_schema, /*selected_keys=*/"a");
     auto read_schema = ExportSchema(ReadSchema("a"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
     ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(reader.get()),
                         "__field_mapping cannot be null");
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestInvalidNullFieldMappingFieldElement) {
+TEST_F(MapSharedShreddingFileReaderTest, TestInvalidNullFieldMappingFieldElement) {
     auto physical_schema = PhysicalSchemaWithMetadata();
     std::string json = R"([
         [1, [[0, null], 10, null, null]]
     ])";
     auto physical_array =
         arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(physical_schema->fields()), json)
             .ValueOrDie();
-    auto reader = CreateReader(physical_array, physical_schema);
+    auto reader = CreateReader(physical_array, physical_schema, /*selected_keys=*/"b");
     auto read_schema = ExportSchema(ReadSchema("b"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
                                     /*selection_bitmap=*/std::nullopt));
     ASSERT_NOK_WITH_MSG(ReadResultCollector::CollectResult(reader.get()),
                         "__field_mapping element cannot be null");
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestListValue) {
+TEST_F(MapSharedShreddingFileReaderTest, TestListValue) {
     std::shared_ptr<arrow::Schema> logical_schema = arrow::schema({
         arrow::field("id", arrow::int32()),
         arrow::field("tags", arrow::map(arrow::utf8(), arrow::list(arrow::int32()))),
@@ -443,7 +474,8 @@ TEST_F(SharedShreddingFileReaderTest, TestListValue) {
         [4, [[1, 0], [8], [9, 10], [[2, [null]]]]]
     ])")
             .ValueOrDie();
-    auto reader = CreateReader(physical_array, physical_schema);
+    auto reader = CreateReader(physical_array, physical_schema,
+                               /*selected_keys=*/"a,c");  // NOLINT(whitespace/comma)
 
     auto read_metadata = std::make_shared<arrow::KeyValueMetadata>();
     read_metadata->Append("paimon.map.selected-keys", "a,c");
@@ -467,7 +499,7 @@ TEST_F(SharedShreddingFileReaderTest, TestListValue) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) {
+TEST_F(MapSharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) {
     std::shared_ptr<arrow::Schema> logical_schema = arrow::schema({
         arrow::field("id", arrow::int32()),
         arrow::field("tags", arrow::map(arrow::utf8(), arrow::utf8())),
@@ -503,9 +535,8 @@ TEST_F(SharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) {
     std::string data_file_path =
         path_factory->ToPath(inc.GetNewFilesIncrement().NewFiles()[0]->file_name);
     std::map<std::string, std::string> reader_options = {{"orc.read.enable-lazy-decoding", "true"}};
-    ASSERT_OK_AND_ASSIGN(auto reader,
-                         SharedShreddingFileReader::Create(
-                             OpenFormatReader(data_file_path, format, reader_options), pool_));
+    auto reader = WrapReader(OpenFormatReader(data_file_path, format, reader_options),
+                             /*selected_keys_str=*/"a,c");
 
     auto read_metadata = std::make_shared<arrow::KeyValueMetadata>();
     read_metadata->Append("paimon.map.selected-keys", "a,c");
@@ -528,7 +559,7 @@ TEST_F(SharedShreddingFileReaderTest, TestOrcDictionaryEncodedStringValue) {
     AssertChunkedArrayEquals(expected, actual);
 }
 
-TEST_F(SharedShreddingFileReaderTest, TestReadsRealFormatFile) {
+TEST_F(MapSharedShreddingFileReaderTest, TestReadsRealFormatFile) {
     // TODO(lisizhuo.lsz): support other format
     auto options = options_;
     std::string format = "orc";
@@ -559,8 +590,8 @@ TEST_F(SharedShreddingFileReaderTest, TestReadsRealFormatFile) {
 
     std::string data_file_path =
         path_factory->ToPath(inc.GetNewFilesIncrement().NewFiles()[0]->file_name);
-    ASSERT_OK_AND_ASSIGN(auto reader, SharedShreddingFileReader::Create(
-                                          OpenFormatReader(data_file_path, format), pool_));
+    auto reader = WrapReader(OpenFormatReader(data_file_path, format),
+                             /*selected_keys_str=*/"a,c");
 
     auto read_schema = ExportSchema(ReadSchema("a,c"));
     ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,