feat(shared-shredding): Expose shared-shredding utilities

dalingmeng · dalingmeng · commit dac41bc37a6d · 2026-07-09T18:28:48.000+08:00
diff --git a/include/paimon/data/shredding/map_shared_shredding_schema_utils.h b/include/paimon/data/shredding/map_shared_shredding_schema_utils.h
@@ -0,0 +1,78 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "paimon/arrow/abi.h"
+#include "paimon/result.h"
+#include "paimon/status.h"
+#include "paimon/visibility.h"
+
+namespace paimon {
+
+/// Parsed file-level meta for one shared-shredding MAP column.
+struct PAIMON_EXPORT MapSharedShreddingFieldMeta {
+    /// field_name -> field_id
+    std::map<std::string, int32_t> name_to_id;
+    /// field_id -> set of physical column indices S
+    std::map<int32_t, std::vector<int32_t>> field_to_columns;
+    /// Set of field_ids that ever spilled into __overflow
+    std::set<int32_t> overflow_field_set;
+    /// Number of physical columns K in this file
+    int32_t num_columns = 0;
+    /// Maximum row width observed in this file
+    int32_t max_row_width = 0;
+
+    bool operator==(const MapSharedShreddingFieldMeta& other) const {
+        if (this == &other) {
+            return true;
+        }
+        return name_to_id == other.name_to_id && field_to_columns == other.field_to_columns &&
+               overflow_field_set == other.overflow_field_set && num_columns == other.num_columns &&
+               max_row_width == other.max_row_width;
+    }
+};
+
+class PAIMON_EXPORT MapSharedShreddingSchemaUtils {
+ public:
+    MapSharedShreddingSchemaUtils() = delete;
+    ~MapSharedShreddingSchemaUtils() = delete;
+
+    /// Converts a logical schema to a physical schema by replacing shredding MAP columns
+    /// with their physical Struct representation.
+    /// @param logical_schema The original Arrow C schema with MAP<STRING, T> columns.
+    ///        Ownership of schema resources is transferred to this method.
+    /// @param field_to_num_columns Map from field name to its physical column count K.
+    ///        Each shredding column can have its own width.
+    /// @return The exported Arrow C schema for file writing.
+    static Result<std::unique_ptr<::ArrowSchema>> LogicalToPhysicalSchema(
+        std::unique_ptr<::ArrowSchema> logical_schema,
+        const std::map<std::string, int32_t>& field_to_num_columns);
+
+    /// Attaches shared-shredding metadata to fields in a physical schema.
+    /// @param physical_schema The Arrow C physical schema whose fields should receive metadata.
+    ///        Ownership of schema resources is transferred to this method.
+    /// @param field_name_to_meta Map from physical field name to its shared-shredding metadata.
+    /// @param compression Compression codec name for field_dict serialization.
+    /// @return A new Arrow C schema with shared-shredding metadata attached to matching fields.
+    static Result<std::unique_ptr<::ArrowSchema>> AttachMetadataToSchema(
+        std::unique_ptr<::ArrowSchema> physical_schema,
+        const std::map<std::string, MapSharedShreddingFieldMeta>& field_name_to_meta,
+        const std::string& compression);
+
+    /// Extracts shared-shredding metadata from a named field in a physical schema.
+    /// @param physical_schema The Arrow C physical schema that contains the target field.
+    ///        Ownership of schema resources is transferred to this method.
+    /// @param field_name The physical field name whose metadata should be extracted.
+    /// @param compression Compression codec name for field_dict deserialization.
+    /// @return Parsed shared-shredding metadata for the field.
+    static Result<MapSharedShreddingFieldMeta> ExtractMetadataFromField(
+        std::unique_ptr<::ArrowSchema> physical_schema, const std::string& field_name,
+        const std::string& compression);
+};
+
+}  // namespace paimon
diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt
@@ -139,6 +139,7 @@ set(PAIMON_COMMON_SRCS
     common/utils/crc32c.cpp
     common/utils/decimal_utils.cpp
     common/data/shredding/map_shared_shredding_utils.cpp
+    common/data/shredding/map_shared_shredding_schema_utils.cpp
     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
@@ -547,6 +548,7 @@ if(PAIMON_BUILD_TESTS)
                     common/utils/threadsafe_queue_test.cpp
                     common/utils/generic_lru_cache_test.cpp
                     common/data/shredding/map_shared_shredding_utils_test.cpp
+                    common/data/shredding/map_shared_shredding_schema_utils_test.cpp
                     common/data/shredding/map_shared_shredding_batch_converter_test.cpp
                     common/data/shredding/lru_map_shared_shredding_column_allocator_test.cpp
                     common/data/shredding/plain_map_shared_shredding_column_allocator_test.cpp
diff --git a/src/paimon/common/data/shredding/map_shared_shredding_schema_utils.cpp b/src/paimon/common/data/shredding/map_shared_shredding_schema_utils.cpp
@@ -0,0 +1,96 @@
+/*
+ * Copyright 2026-present Alibaba Inc.
+ *
+ * Licensed 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 "paimon/data/shredding/map_shared_shredding_schema_utils.h"
+
+#include "arrow/c/bridge.h"
+#include "arrow/type.h"
+#include "arrow/util/key_value_metadata.h"
+#include "fmt/format.h"
+#include "paimon/common/data/shredding/map_shared_shredding_utils.h"
+#include "paimon/common/utils/arrow/status_utils.h"
+
+namespace paimon {
+
+Result<std::unique_ptr<::ArrowSchema>> MapSharedShreddingSchemaUtils::LogicalToPhysicalSchema(
+    std::unique_ptr<::ArrowSchema> logical_schema,
+    const std::map<std::string, int32_t>& field_to_num_columns) {
+    if (!logical_schema) {
+        return Status::Invalid("logical schema is null");
+    }
+    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> arrow_logical_schema,
+                                      arrow::ImportSchema(logical_schema.get()));
+    std::shared_ptr<arrow::Schema> physical_schema =
+        MapSharedShreddingUtils::LogicalToPhysicalSchema(arrow_logical_schema,
+                                                         field_to_num_columns);
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*physical_schema, c_schema.get()));
+    return c_schema;
+}
+
+Result<std::unique_ptr<::ArrowSchema>> MapSharedShreddingSchemaUtils::AttachMetadataToSchema(
+    std::unique_ptr<::ArrowSchema> physical_schema,
+    const std::map<std::string, MapSharedShreddingFieldMeta>& field_name_to_meta,
+    const std::string& compression) {
+    if (!physical_schema) {
+        return Status::Invalid("physical schema is null");
+    }
+    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> arrow_physical_schema,
+                                      arrow::ImportSchema(physical_schema.get()));
+
+    arrow::FieldVector updated_fields = arrow_physical_schema->fields();
+    for (const auto& [field_name, field_meta] : field_name_to_meta) {
+        int32_t field_index = arrow_physical_schema->GetFieldIndex(field_name);
+        if (field_index < 0) {
+            return Status::Invalid(fmt::format(
+                "Shared-shredding field '{}' not found in physical schema.", field_name));
+        }
+
+        const auto& field = arrow_physical_schema->field(field_index);
+        auto metadata = field->metadata() ? field->metadata()->Copy()
+                                          : std::make_shared<arrow::KeyValueMetadata>();
+        PAIMON_RETURN_NOT_OK(
+            MapSharedShreddingUtils::SerializeMetadata(field_meta, compression, metadata.get()));
+        updated_fields[field_index] = field->WithMetadata(metadata);
+    }
+
+    auto updated_schema =
+        arrow::schema(std::move(updated_fields), arrow_physical_schema->metadata());
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*updated_schema, c_schema.get()));
+    return c_schema;
+}
+
+Result<MapSharedShreddingFieldMeta> MapSharedShreddingSchemaUtils::ExtractMetadataFromField(
+    std::unique_ptr<::ArrowSchema> physical_schema, const std::string& field_name,
+    const std::string& compression) {
+    if (!physical_schema) {
+        return Status::Invalid("physical schema is null");
+    }
+    PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Schema> arrow_physical_schema,
+                                      arrow::ImportSchema(physical_schema.get()));
+    const auto& field = arrow_physical_schema->GetFieldByName(field_name);
+    if (!field) {
+        return Status::Invalid(
+            fmt::format("Shared-shredding field '{}' not found in physical schema.", field_name));
+    }
+
+    auto metadata =
+        field->metadata() ? field->metadata()->Copy() : std::shared_ptr<arrow::KeyValueMetadata>();
+    return MapSharedShreddingUtils::DeserializeMetadata(metadata, compression);
+}
+
+}  // namespace paimon
diff --git a/src/paimon/common/data/shredding/map_shared_shredding_schema_utils_test.cpp b/src/paimon/common/data/shredding/map_shared_shredding_schema_utils_test.cpp
@@ -0,0 +1,172 @@
+/*
+ * Copyright 2026-present Alibaba Inc.
+ *
+ * Licensed 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 "paimon/data/shredding/map_shared_shredding_schema_utils.h"
+
+#include <string>
+#include <utility>
+
+#include "arrow/c/bridge.h"
+#include "arrow/type.h"
+#include "arrow/util/key_value_metadata.h"
+#include "gtest/gtest.h"
+#include "paimon/common/data/shredding/map_shared_shredding_utils.h"
+#include "paimon/testing/utils/testharness.h"
+
+namespace paimon::test {
+
+TEST(MapSharedShreddingSchemaUtilsTest, AttachMetadataToSchemaBasic) {
+    MapSharedShreddingFieldMeta tags_meta;
+    tags_meta.name_to_id = {{"host", 0}, {"region", 1}};
+    tags_meta.field_to_columns = {{0, {0}}, {1, {1}}};
+    tags_meta.num_columns = 2;
+    tags_meta.max_row_width = 2;
+
+    auto id_metadata = std::make_shared<arrow::KeyValueMetadata>();
+    id_metadata->Append("paimon.field.id", "1");
+    auto schema_metadata = std::make_shared<arrow::KeyValueMetadata>();
+    schema_metadata->Append("schema.key", "schema.value");
+    auto schema = arrow::schema(
+        {arrow::field("id", arrow::int32(), true, id_metadata),
+         arrow::field("tags", arrow::struct_({arrow::field("__field_mapping",
+                                                           arrow::list(arrow::int32()), true),
+                                              arrow::field("__col_0", arrow::utf8(), true),
+                                              arrow::field("__col_1", arrow::utf8(), true)}))},
+        schema_metadata);
+
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok());
+    ASSERT_OK_AND_ASSIGN(auto c_updated_schema,
+                         MapSharedShreddingSchemaUtils::AttachMetadataToSchema(
+                             std::move(c_schema), {{"tags", tags_meta}}, "none"));
+    ASSERT_TRUE(c_updated_schema);
+    ASSERT_TRUE(c_updated_schema->release);
+    auto updated_schema = arrow::ImportSchema(c_updated_schema.get()).ValueOrDie();
+
+    ASSERT_TRUE(updated_schema->metadata()->Equals(*schema_metadata));
+    ASSERT_TRUE(updated_schema->field(0)->metadata()->Equals(*id_metadata));
+
+    auto tags_metadata = updated_schema->GetFieldByName("tags")->metadata()->Copy();
+    ASSERT_TRUE(MapSharedShreddingUtils::HasShreddingMetadata(tags_metadata));
+    ASSERT_OK_AND_ASSIGN(auto deserialized,
+                         MapSharedShreddingUtils::DeserializeMetadata(tags_metadata, "none"));
+    ASSERT_EQ(deserialized, tags_meta);
+}
+
+TEST(MapSharedShreddingSchemaUtilsTest, AttachMetadataToSchemaMissingField) {
+    MapSharedShreddingFieldMeta tags_meta;
+    tags_meta.name_to_id = {{"host", 0}};
+    tags_meta.field_to_columns = {{0, {0}}};
+    tags_meta.num_columns = 1;
+    tags_meta.max_row_width = 1;
+
+    auto schema = arrow::schema({arrow::field("id", arrow::int32())});
+
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok());
+    ASSERT_NOK_WITH_MSG(MapSharedShreddingSchemaUtils::AttachMetadataToSchema(
+                            std::move(c_schema), {{"tags", tags_meta}}, "none"),
+                        "Shared-shredding field 'tags' not found in physical schema.");
+}
+
+TEST(MapSharedShreddingSchemaUtilsTest, AttachMetadataToSchemaPreservesExistingFieldMetadata) {
+    MapSharedShreddingFieldMeta tags_meta;
+    tags_meta.name_to_id = {{"host", 0}};
+    tags_meta.field_to_columns = {{0, {0}}};
+    tags_meta.num_columns = 1;
+    tags_meta.max_row_width = 1;
+
+    auto tags_metadata = std::make_shared<arrow::KeyValueMetadata>();
+    tags_metadata->Append("paimon.field.id", "7");
+    tags_metadata->Append("description", "original tags field");
+    auto schema = arrow::schema({arrow::field(
+        "tags",
+        arrow::struct_({arrow::field("__field_mapping", arrow::list(arrow::int32()), true),
+                        arrow::field("__col_0", arrow::utf8(), true)}),
+        true, tags_metadata)});
+
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok());
+    ASSERT_OK_AND_ASSIGN(auto c_updated_schema,
+                         MapSharedShreddingSchemaUtils::AttachMetadataToSchema(
+                             std::move(c_schema), {{"tags", tags_meta}}, "none"));
+    ASSERT_TRUE(c_updated_schema);
+    ASSERT_TRUE(c_updated_schema->release);
+    auto updated_schema = arrow::ImportSchema(c_updated_schema.get()).ValueOrDie();
+
+    auto updated_metadata = updated_schema->field(0)->metadata()->Copy();
+    ASSERT_EQ(updated_metadata->value(updated_metadata->FindKey("paimon.field.id")), "7");
+    ASSERT_EQ(updated_metadata->value(updated_metadata->FindKey("description")),
+              "original tags field");
+    ASSERT_TRUE(MapSharedShreddingUtils::HasShreddingMetadata(updated_metadata));
+}
+
+TEST(MapSharedShreddingSchemaUtilsTest, ExtractMetadataFromField) {
+    MapSharedShreddingFieldMeta tags_meta;
+    tags_meta.name_to_id = {{"host", 0}, {"region", 1}};
+    tags_meta.field_to_columns = {{0, {0}}, {1, {1}}};
+    tags_meta.num_columns = 2;
+    tags_meta.max_row_width = 2;
+
+    auto metadata = std::make_shared<arrow::KeyValueMetadata>();
+    ASSERT_OK(MapSharedShreddingUtils::SerializeMetadata(tags_meta, "none", metadata.get()));
+    auto field = arrow::field(
+        "tags",
+        arrow::struct_({arrow::field("__field_mapping", arrow::list(arrow::int32()), true),
+                        arrow::field("__col_0", arrow::utf8(), true),
+                        arrow::field("__col_1", arrow::utf8(), true)}),
+        true, metadata);
+    auto schema = arrow::schema({arrow::field("id", arrow::int32()), field});
+
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok());
+    ASSERT_OK_AND_ASSIGN(auto parsed_meta, MapSharedShreddingSchemaUtils::ExtractMetadataFromField(
+                                               std::move(c_schema), "tags", "none"));
+    ASSERT_EQ(parsed_meta, tags_meta);
+}
+
+TEST(MapSharedShreddingSchemaUtilsTest, LogicalToPhysicalSchemaNestedListValue) {
+    // MAP<STRING, STRUCT<a:array<int32>, b: array<int32>, c: array<int32>>>
+    auto nested_value = arrow::struct_({arrow::field("a", arrow::list(arrow::int32())),
+                                        arrow::field("b", arrow::list(arrow::int32())),
+                                        arrow::field("c", arrow::list(arrow::int32()))});
+    auto map_type = arrow::map(arrow::utf8(), nested_value);
+    auto schema =
+        arrow::schema({arrow::field("ts", arrow::int64()), arrow::field("data", map_type)});
+
+    std::map<std::string, int32_t> field_to_num_columns = {{"data", 3}};
+    auto c_schema = std::make_unique<::ArrowSchema>();
+    EXPECT_TRUE(arrow::ExportSchema(*schema, c_schema.get()).ok());
+    ASSERT_OK_AND_ASSIGN(auto c_physical_schema,
+                         MapSharedShreddingSchemaUtils::LogicalToPhysicalSchema(
+                             std::move(c_schema), field_to_num_columns));
+    ASSERT_TRUE(c_physical_schema);
+    ASSERT_TRUE(c_physical_schema->release);
+    auto physical_schema = arrow::ImportSchema(c_physical_schema.get()).ValueOrDie();
+
+    auto expected_struct = arrow::struct_({
+        arrow::field("__field_mapping", arrow::list(arrow::int32()), true),
+        arrow::field("__col_0", nested_value, true),
+        arrow::field("__col_1", nested_value, true),
+        arrow::field("__col_2", nested_value, true),
+        arrow::field("__overflow", arrow::map(arrow::int32(), nested_value), true),
+    });
+    auto expected_schema = arrow::schema(
+        {arrow::field("ts", arrow::int64()), arrow::field("data", expected_struct, true)});
+    ASSERT_TRUE(physical_schema->Equals(expected_schema));
+}
+
+}  // namespace paimon::test
diff --git a/src/paimon/common/data/shredding/map_shared_shredding_utils.h b/src/paimon/common/data/shredding/map_shared_shredding_utils.h
diff --git a/src/paimon/common/data/shredding/map_shredding_defs.h b/src/paimon/common/data/shredding/map_shredding_defs.h
