Skip to content

Commit c3efa47

Browse files
lszskyelszskye
committed
feat(shredding): support selected-key pushdown by MapSharedShreddingAccessBuilder
1 parent 3663237 commit c3efa47

14 files changed

Lines changed: 1054 additions & 168 deletions

‎include/paimon/data/shredding/map_shared_shredding_schema_utils.h‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,48 @@ struct PAIMON_EXPORT MapSharedShreddingFieldMeta {
5656
}
5757
};
5858

59+
/// Builds a selected-key projection field for a top-level shared-shredding MAP column.
60+
///
61+
/// The built field is a STRUCT which replaces the MAP field in the read schema. Each child
62+
/// corresponds to one selected key and contains that key's MAP value, or NULL when the key is
63+
/// absent. Children are named by their ordinal ("0", "1", ...), and preserve insertion order.
64+
///
65+
/// Example: read keys "age" and "score" from MAP column `attributes`:
66+
///
67+
/// auto builder = MapSharedShreddingAccessBuilder::Create(attributes_field);
68+
/// builder->AddKey("age");
69+
/// builder->AddKey("score");
70+
/// auto field = builder->Build();
71+
///
72+
/// Use the returned field in `ReadContextBuilder::SetReadSchema`.
73+
class PAIMON_EXPORT MapSharedShreddingAccessBuilder {
74+
public:
75+
/// Creates a builder bound to the original MAP field.
76+
///
77+
/// The field must be a MAP with STRING keys. Its name, nullability, and value type are
78+
/// retained for the selected-key projection. Ownership of the Arrow C schema resources is
79+
/// transferred to this method.
80+
static Result<std::unique_ptr<MapSharedShreddingAccessBuilder>> Create(
81+
struct ArrowSchema* map_field);
82+
83+
~MapSharedShreddingAccessBuilder();
84+
85+
/// Adds a selected MAP key.
86+
///
87+
/// @param key The string MAP key. Keys are returned in insertion order.
88+
Status AddKey(const std::string& key);
89+
90+
/// Builds a STRUCT projection field which retains the original MAP field's name and
91+
/// nullability. Every selected-key child uses the complete MAP value type and is nullable.
92+
Result<std::unique_ptr<struct ArrowSchema>> Build() const;
93+
94+
private:
95+
class Impl;
96+
explicit MapSharedShreddingAccessBuilder(std::unique_ptr<Impl>&& impl);
97+
98+
std::unique_ptr<Impl> impl_;
99+
};
100+
59101
class PAIMON_EXPORT MapSharedShreddingSchemaUtils {
60102
public:
61103
MapSharedShreddingSchemaUtils() = delete;

‎src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp‎

Lines changed: 469 additions & 114 deletions
Large diffs are not rendered by default.

‎src/paimon/common/data/shredding/map_shared_shredding_file_reader.h‎

Lines changed: 43 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -33,21 +33,52 @@
3333

3434
namespace paimon {
3535

36-
class MapSharedShreddingFileReader : public FileBatchReader {
36+
class MapFieldReadPlan {
3737
public:
38-
struct SharedShreddingContext {
39-
SharedShreddingContext(const MapSharedShreddingFieldMeta& _meta,
40-
const std::vector<std::string>& _selected_keys,
41-
const std::shared_ptr<arrow::MapType>& _map_type)
42-
: meta(_meta), selected_keys(_selected_keys), map_type(_map_type) {}
43-
MapSharedShreddingFieldMeta meta;
44-
std::vector<std::string> selected_keys;
45-
std::shared_ptr<arrow::MapType> map_type;
46-
};
38+
virtual ~MapFieldReadPlan() = default;
39+
40+
MapFieldReadPlan(const std::shared_ptr<arrow::Field>& logical_field,
41+
const std::shared_ptr<arrow::Field>& physical_read_field)
42+
: logical_field_(logical_field), physical_read_field_(physical_read_field) {}
43+
44+
const std::shared_ptr<arrow::Field>& LogicalField() const {
45+
return logical_field_;
46+
}
47+
48+
const std::shared_ptr<arrow::Field>& PhysicalReadField() const {
49+
return physical_read_field_;
50+
}
51+
52+
virtual Result<std::shared_ptr<arrow::Array>> Materialize(
53+
const std::shared_ptr<arrow::Array>& physical_array,
54+
arrow::MemoryPool* arrow_pool) const = 0;
55+
56+
private:
57+
std::shared_ptr<arrow::Field> logical_field_;
58+
std::shared_ptr<arrow::Field> physical_read_field_;
59+
};
60+
61+
class MapFieldReadPlanFactory {
62+
public:
63+
static Result<std::unique_ptr<MapFieldReadPlan>> CreateFullMapReadPlan(
64+
const std::shared_ptr<arrow::Field>& logical_map_field,
65+
const MapSharedShreddingFieldMeta& meta, const std::vector<std::string>& selected_keys);
4766

67+
static Result<std::unique_ptr<MapFieldReadPlan>> CreateSharedSelectedKeysReadPlan(
68+
const std::shared_ptr<arrow::Field>& selected_keys_field,
69+
const MapSharedShreddingFieldMeta& meta, const std::vector<std::string>& selected_keys);
70+
71+
static Result<std::unique_ptr<MapFieldReadPlan>> CreateDefaultSelectedKeysReadPlan(
72+
const std::shared_ptr<arrow::Field>& file_map_field,
73+
const std::shared_ptr<arrow::Field>& selected_keys_field,
74+
const std::vector<std::string>& selected_keys);
75+
};
76+
77+
class MapSharedShreddingFileReader : public FileBatchReader {
78+
public:
4879
MapSharedShreddingFileReader(
4980
std::unique_ptr<FileBatchReader>&& reader,
50-
std::map<std::string, SharedShreddingContext>&& shared_shredding_name_to_context,
81+
std::map<std::string, std::unique_ptr<MapFieldReadPlan>>&& field_read_plans,
5182
const std::shared_ptr<MemoryPool>& pool);
5283

5384
Result<std::unique_ptr<::ArrowSchema>> GetFileSchema() const override;
@@ -70,25 +101,13 @@ class MapSharedShreddingFileReader : public FileBatchReader {
70101
bool SupportPreciseBitmapSelection() const override;
71102

72103
private:
73-
Result<std::shared_ptr<arrow::Array>> RebuildLogicalMapArray(
74-
const std::shared_ptr<arrow::Field>& physical_field,
75-
const std::shared_ptr<arrow::StructArray>& physical_struct_array) const;
76-
77-
static std::vector<std::pair<std::string, int32_t>> ResolveSelectedKeyIds(
78-
const MapSharedShreddingFieldMeta& meta, const std::vector<std::string>& selected_keys);
79-
80-
static void CollectPhysicalColumns(
81-
const std::shared_ptr<arrow::StructArray>& physical_struct_array,
82-
std::map<std::string, std::shared_ptr<arrow::Array>>* physical_column_name_to_array,
83-
std::shared_ptr<arrow::MapArray>* overflow_array);
84-
85104
static Result<std::shared_ptr<arrow::Field>> ToLogicalMapField(
86105
const std::shared_ptr<arrow::Field>& physical_field);
87106

88107
private:
89108
std::shared_ptr<arrow::MemoryPool> arrow_pool_;
90109
std::unique_ptr<FileBatchReader> reader_;
91-
std::map<std::string, SharedShreddingContext> shared_shredding_name_to_context_;
110+
std::map<std::string, std::unique_ptr<MapFieldReadPlan>> field_read_plans_;
92111
};
93112

94113
} // namespace paimon

‎src/paimon/common/data/shredding/map_shared_shredding_file_reader_test.cpp‎

Lines changed: 96 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,7 @@ class MapSharedShreddingFileReaderTest : public ::testing::Test {
104104
const std::optional<std::string>& selected_keys_str = std::nullopt) const {
105105
EXPECT_OK_AND_ASSIGN(auto c_file_schema, reader->GetFileSchema());
106106
auto file_schema = arrow::ImportSchema(c_file_schema.get()).ValueOrDie();
107-
std::map<std::string, MapSharedShreddingFileReader::SharedShreddingContext>
108-
shared_shredding_name_to_context;
107+
std::map<std::string, std::unique_ptr<MapFieldReadPlan>> field_read_plans;
109108
for (const auto& field : file_schema->fields()) {
110109
auto metadata = std::const_pointer_cast<arrow::KeyValueMetadata>(field->metadata());
111110
if (!MapSharedShreddingUtils::HasShreddingMetadata(metadata)) {
@@ -135,12 +134,14 @@ class MapSharedShreddingFileReaderTest : public ::testing::Test {
135134
selected_keys.push_back(key_name);
136135
}
137136
}
138-
shared_shredding_name_to_context.emplace(
139-
field->name(), MapSharedShreddingFileReader::SharedShreddingContext(
140-
meta, selected_keys, map_type));
137+
auto logical_map_field = field->WithType(map_type);
138+
EXPECT_OK_AND_ASSIGN(auto field_read_plan,
139+
MapFieldReadPlanFactory::CreateFullMapReadPlan(
140+
logical_map_field, meta, selected_keys));
141+
field_read_plans.emplace(field->name(), std::move(field_read_plan));
141142
}
142-
return std::make_unique<MapSharedShreddingFileReader>(
143-
std::move(reader), std::move(shared_shredding_name_to_context), pool_);
143+
return std::make_unique<MapSharedShreddingFileReader>(std::move(reader),
144+
std::move(field_read_plans), pool_);
144145
}
145146

146147
Result<std::unique_ptr<MapSharedShreddingFileReader>> CreateReader(
@@ -299,6 +300,94 @@ TEST_F(MapSharedShreddingFileReaderTest, TestAllExistSelectedKeysWithOverflow) {
299300
AssertChunkedArrayEquals(expected, actual);
300301
}
301302

303+
TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjection) {
304+
ASSERT_OK_AND_ASSIGN(auto physical_schema, PhysicalSchemaWithMetadata());
305+
ASSERT_OK_AND_ASSIGN(auto physical_array, PhysicalArray());
306+
auto mock_reader = std::make_unique<MockFileBatchReader>(
307+
physical_array, arrow::struct_(physical_schema->fields()), /*read_batch_size=*/10);
308+
mock_reader->EnableRandomizeBatchSize(false);
309+
310+
auto selected_type =
311+
arrow::struct_({arrow::field("0", arrow::int64()), arrow::field("1", arrow::int64()),
312+
arrow::field("2", arrow::int64())});
313+
auto selected_field = arrow::field(
314+
"tags", selected_type, /*nullable=*/true,
315+
arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,c,missing"}));
316+
ASSERT_OK_AND_ASSIGN(auto field_read_plan,
317+
MapFieldReadPlanFactory::CreateSharedSelectedKeysReadPlan(
318+
selected_field, TagsMeta(), {"a", "c", "missing"}));
319+
std::map<std::string, std::unique_ptr<MapFieldReadPlan>> contexts;
320+
contexts.emplace("tags", std::move(field_read_plan));
321+
auto reader = std::make_unique<MapSharedShreddingFileReader>(std::move(mock_reader),
322+
std::move(contexts), pool_);
323+
324+
auto read_schema =
325+
ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field}));
326+
ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
327+
/*selection_bitmap=*/std::nullopt));
328+
ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get()));
329+
330+
auto expected_type = arrow::struct_({arrow::field("id", arrow::int32()), selected_field});
331+
std::shared_ptr<arrow::ChunkedArray> expected;
332+
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(expected_type, {R"([
333+
[1, [10, null, null]],
334+
[2, [40, 30, null]],
335+
[3, null],
336+
[4, [80, null, null]]
337+
])"},
338+
&expected)
339+
.ok());
340+
AssertChunkedArrayEquals(expected, actual);
341+
}
342+
343+
TEST_F(MapSharedShreddingFileReaderTest, TestSelectedKeysStructProjectionFromLegacyMap) {
344+
auto map_type = arrow::internal::checked_pointer_cast<arrow::MapType>(
345+
arrow::map(arrow::utf8(), arrow::field("value", arrow::int64())));
346+
auto file_schema =
347+
arrow::schema({arrow::field("id", arrow::int32()), arrow::field("tags", map_type)});
348+
auto file_array =
349+
arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_(file_schema->fields()),
350+
R"([
351+
[1, [["a", 10], ["c", null]]],
352+
[2, [["b", 20]]],
353+
[3, null]
354+
])")
355+
.ValueOrDie();
356+
auto mock_reader = std::make_unique<MockFileBatchReader>(
357+
file_array, arrow::struct_(file_schema->fields()), /*read_batch_size=*/10);
358+
mock_reader->EnableRandomizeBatchSize(false);
359+
360+
auto selected_type =
361+
arrow::struct_({arrow::field("0", arrow::int64()), arrow::field("1", arrow::int64())});
362+
auto selected_field =
363+
arrow::field("tags", selected_type, /*nullable=*/true,
364+
arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {"a,missing"}));
365+
ASSERT_OK_AND_ASSIGN(auto field_read_plan,
366+
MapFieldReadPlanFactory::CreateDefaultSelectedKeysReadPlan(
367+
file_schema->field(1), selected_field, {"a", "missing"}));
368+
std::map<std::string, std::unique_ptr<MapFieldReadPlan>> contexts;
369+
contexts.emplace("tags", std::move(field_read_plan));
370+
auto reader = std::make_unique<MapSharedShreddingFileReader>(std::move(mock_reader),
371+
std::move(contexts), pool_);
372+
373+
auto read_schema =
374+
ExportSchema(arrow::schema({arrow::field("id", arrow::int32()), selected_field}));
375+
ASSERT_OK(reader->SetReadSchema(read_schema.get(), /*predicate=*/nullptr,
376+
/*selection_bitmap=*/std::nullopt));
377+
ASSERT_OK_AND_ASSIGN(auto actual, ReadResultCollector::CollectResult(reader.get()));
378+
379+
auto expected_type = arrow::struct_({arrow::field("id", arrow::int32()), selected_field});
380+
std::shared_ptr<arrow::ChunkedArray> expected;
381+
ASSERT_TRUE(arrow::ipc::internal::json::ChunkedArrayFromJSON(expected_type, {R"([
382+
[1, [10, null]],
383+
[2, [null, null]],
384+
[3, null]
385+
])"},
386+
&expected)
387+
.ok());
388+
AssertChunkedArrayEquals(expected, actual);
389+
}
390+
302391
TEST_F(MapSharedShreddingFileReaderTest, TestPartialExistSelectedKeys) {
303392
ASSERT_OK_AND_ASSIGN(auto reader,
304393
CreateReader(/*physical_array=*/nullptr, /*physical_schema=*/nullptr,

‎src/paimon/common/data/shredding/map_shared_shredding_schema_utils.cpp‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,15 +19,92 @@
1919

2020
#include "paimon/data/shredding/map_shared_shredding_schema_utils.h"
2121

22+
#include <unordered_set>
23+
#include <utility>
24+
#include <vector>
25+
2226
#include "arrow/c/bridge.h"
2327
#include "arrow/type.h"
2428
#include "arrow/util/key_value_metadata.h"
2529
#include "fmt/format.h"
2630
#include "paimon/common/data/shredding/map_shared_shredding_utils.h"
31+
#include "paimon/common/types/data_field.h"
2732
#include "paimon/common/utils/arrow/status_utils.h"
2833

2934
namespace paimon {
3035

36+
class MapSharedShreddingAccessBuilder::Impl {
37+
public:
38+
Impl(const std::shared_ptr<arrow::Field>& _map_field,
39+
const std::shared_ptr<arrow::MapType>& _map_type)
40+
: map_field(_map_field), map_type(_map_type) {}
41+
42+
std::shared_ptr<arrow::Field> map_field;
43+
std::shared_ptr<arrow::MapType> map_type;
44+
std::vector<std::string> keys;
45+
std::unordered_set<std::string> unique_keys;
46+
};
47+
48+
MapSharedShreddingAccessBuilder::~MapSharedShreddingAccessBuilder() = default;
49+
50+
MapSharedShreddingAccessBuilder::MapSharedShreddingAccessBuilder(std::unique_ptr<Impl>&& impl)
51+
: impl_(std::move(impl)) {}
52+
53+
Result<std::unique_ptr<MapSharedShreddingAccessBuilder>> MapSharedShreddingAccessBuilder::Create(
54+
struct ArrowSchema* map_field) {
55+
if (!map_field) {
56+
return Status::Invalid("MAP field is null");
57+
}
58+
PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Field> field,
59+
arrow::ImportField(map_field));
60+
if (field->type()->id() != arrow::Type::MAP) {
61+
return Status::Invalid(
62+
fmt::format("MapSharedShreddingAccessBuilder requires MAP field, got {}",
63+
field->type()->ToString()));
64+
}
65+
auto map_type = arrow::internal::checked_pointer_cast<arrow::MapType>(field->type());
66+
if (map_type->key_type()->id() != arrow::Type::STRING) {
67+
return Status::Invalid(fmt::format(
68+
"MapSharedShreddingAccessBuilder only supports MAP with STRING keys, got {}",
69+
map_type->key_type()->ToString()));
70+
}
71+
auto impl = std::make_unique<Impl>(field, map_type);
72+
return std::unique_ptr<MapSharedShreddingAccessBuilder>(
73+
new MapSharedShreddingAccessBuilder(std::move(impl)));
74+
}
75+
76+
Status MapSharedShreddingAccessBuilder::AddKey(const std::string& key) {
77+
if (!impl_->unique_keys.insert(key).second) {
78+
return Status::Invalid(fmt::format("selected MAP key must not be duplicated: {}", key));
79+
}
80+
impl_->keys.push_back(key);
81+
return Status::OK();
82+
}
83+
84+
Result<std::unique_ptr<struct ArrowSchema>> MapSharedShreddingAccessBuilder::Build() const {
85+
if (impl_->keys.empty()) {
86+
return Status::Invalid(
87+
"shared shredding MAP selected-key projection needs at least one key");
88+
}
89+
arrow::FieldVector fields;
90+
fields.reserve(impl_->keys.size());
91+
std::string encoded_keys;
92+
for (size_t i = 0; i < impl_->keys.size(); ++i) {
93+
if (i != 0) {
94+
encoded_keys.push_back(',');
95+
}
96+
encoded_keys.append(impl_->keys[i]);
97+
fields.push_back(arrow::field(std::to_string(i), impl_->map_type->item_type(),
98+
/*nullable=*/true));
99+
}
100+
auto metadata = arrow::KeyValueMetadata::Make({DataField::MAP_SELECTED_KEYS}, {encoded_keys});
101+
auto access_field = impl_->map_field->WithType(arrow::struct_(std::move(fields)))
102+
->WithMetadata(std::move(metadata));
103+
auto field = std::make_unique<struct ArrowSchema>();
104+
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportField(*access_field, field.get()));
105+
return field;
106+
}
107+
31108
Result<std::unique_ptr<::ArrowSchema>> MapSharedShreddingSchemaUtils::LogicalToPhysicalSchema(
32109
std::unique_ptr<::ArrowSchema> logical_schema,
33110
const std::map<std::string, int32_t>& field_to_num_columns) {

0 commit comments

Comments
 (0)