Skip to content

Commit e10f48e

Browse files
author
kiwimeng
committed
feat(parquet): apply name mapping to files without field IDs
1 parent 5107eda commit e10f48e

5 files changed

Lines changed: 361 additions & 8 deletions

File tree

‎src/iceberg/parquet/parquet_reader.cc‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -61,19 +61,23 @@ Result<std::shared_ptr<::arrow::io::RandomAccessFile>> OpenInputStream(
6161
}
6262

6363
Result<SchemaProjection> BuildProjection(::parquet::arrow::FileReader* reader,
64-
const Schema& read_schema) {
64+
const Schema& read_schema,
65+
const NameMapping* name_mapping) {
6566
auto metadata = reader->parquet_reader()->metadata();
6667

67-
if (!HasFieldIds(metadata->schema()->schema_root())) {
68-
// TODO(gangwu): apply name mapping to Parquet schema
69-
return NotImplemented("Applying name mapping to Parquet schema is not implemented");
70-
}
71-
7268
::parquet::arrow::SchemaManifest schema_manifest;
7369
ICEBERG_ARROW_RETURN_NOT_OK(::parquet::arrow::SchemaManifest::Make(
7470
metadata->schema(), metadata->key_value_metadata(), reader->properties(),
7571
&schema_manifest));
7672

73+
if (!HasFieldIds(metadata->schema()->schema_root())) {
74+
if (name_mapping == nullptr) {
75+
return InvalidSchema(
76+
"Parquet file schema has no field IDs and no name mapping provided");
77+
}
78+
ApplyNameMapping(*name_mapping, &schema_manifest);
79+
}
80+
7781
// Leverage SchemaManifest to project the schema
7882
ICEBERG_ASSIGN_OR_RAISE(auto projection, Project(read_schema, schema_manifest));
7983
return projection;
@@ -269,7 +273,6 @@ struct ReadContext {
269273
// 1. Make the memory pool configurable
270274
// 2. Catch ParquetException and convert to Status/Result
271275
// 3. Add utility to convert Arrow Status/Result to Iceberg Status/Result
272-
// 4. Check field ids and apply name mapping if needed
273276
class ParquetReader::Impl {
274277
public:
275278
// Open the Parquet reader with the given options
@@ -301,7 +304,8 @@ class ParquetReader::Impl {
301304
ICEBERG_RETURN_UNEXPECTED(PlanRowGroups(options));
302305

303306
// Project read schema onto the Parquet file schema
304-
ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(), *read_schema_));
307+
ICEBERG_ASSIGN_OR_RAISE(projection_, BuildProjection(reader_.get(), *read_schema_,
308+
options.name_mapping.get()));
305309
metadata_context_ = {.file_path = options.path,
306310
.first_row_id = options.first_row_id,
307311
.data_sequence_number = options.data_sequence_number};

‎src/iceberg/parquet/parquet_schema_util.cc‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
#include "iceberg/constants.h"
2828
#include "iceberg/metadata_columns.h"
29+
#include "iceberg/name_mapping.h"
2930
#include "iceberg/parquet/parquet_schema_util_internal.h"
3031
#include "iceberg/result.h"
3132
#include "iceberg/schema_internal.h"
@@ -716,6 +717,52 @@ void CollectColumnIds(const FieldProjection& field_projection,
716717
}
717718
}
718719

720+
// Forward declaration
721+
void ApplyNameMappingToField(const NameMapping& mapping, std::vector<std::string>& names,
722+
::parquet::arrow::SchemaField* parquet_field);
723+
724+
void ApplyNameMappingToChild(const NameMapping& mapping, std::vector<std::string>& names,
725+
std::string name, ::parquet::arrow::SchemaField* child) {
726+
names.push_back(std::move(name));
727+
ApplyNameMappingToField(mapping, names, child);
728+
names.pop_back();
729+
}
730+
731+
void ApplyNameMappingToField(const NameMapping& mapping, std::vector<std::string>& names,
732+
::parquet::arrow::SchemaField* parquet_field) {
733+
if (auto mapped_field = mapping.Find(names);
734+
mapped_field.has_value() && mapped_field->get().field_id.has_value()) {
735+
parquet_field->field =
736+
parquet_field->field->WithMergedMetadata(::arrow::key_value_metadata(
737+
{std::string(kParquetFieldIdKey)},
738+
{std::to_string(mapped_field->get().field_id.value())}));
739+
}
740+
741+
auto& children = parquet_field->children;
742+
switch (parquet_field->field->type()->id()) {
743+
case ::arrow::Type::STRUCT:
744+
for (auto& child : children) {
745+
ApplyNameMappingToChild(mapping, names, child.field->name(), &child);
746+
}
747+
break;
748+
case ::arrow::Type::LIST:
749+
case ::arrow::Type::LARGE_LIST:
750+
if (children.size() == 1) {
751+
ApplyNameMappingToChild(mapping, names, "element", &children[0]);
752+
}
753+
break;
754+
case ::arrow::Type::MAP:
755+
// The key and value sit under the entries struct, which the name path skips.
756+
if (children.size() == 1 && children[0].children.size() == 2) {
757+
ApplyNameMappingToChild(mapping, names, "key", &children[0].children[0]);
758+
ApplyNameMappingToChild(mapping, names, "value", &children[0].children[1]);
759+
}
760+
break;
761+
default:
762+
break;
763+
}
764+
}
765+
719766
} // namespace
720767

721768
Result<SchemaProjection> Project(const Schema& expected_schema,
@@ -753,6 +800,14 @@ bool HasFieldIds(const ::parquet::schema::NodePtr& node) {
753800
return false;
754801
}
755802

803+
void ApplyNameMapping(const NameMapping& mapping,
804+
::parquet::arrow::SchemaManifest* parquet_schema) {
805+
std::vector<std::string> names;
806+
for (auto& parquet_field : parquet_schema->schema_fields) {
807+
ApplyNameMappingToChild(mapping, names, parquet_field.field->name(), &parquet_field);
808+
}
809+
}
810+
756811
Result<std::shared_ptr<::parquet::SchemaDescriptor>> ToParquetSchema(
757812
const Schema& iceberg_schema) {
758813
::parquet::schema::NodeVector fields;

‎src/iceberg/parquet/parquet_schema_util_internal.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626

2727
#include "iceberg/schema.h"
2828
#include "iceberg/schema_util.h"
29+
#include "iceberg/type_fwd.h"
2930

3031
namespace iceberg::parquet {
3132

@@ -63,6 +64,16 @@ std::vector<int32_t> SelectedColumnIndices(const SchemaProjection& projection);
6364
/// \return True if the Parquet schema has field IDs, false otherwise.
6465
bool HasFieldIds(const ::parquet::schema::NodePtr& root_node);
6566

67+
/// \brief Assign field IDs from a name mapping to a Parquet schema without field IDs.
68+
///
69+
/// List elements use "element"; map keys and values use "key" and "value".
70+
/// Fields without a mapped ID remain unassigned.
71+
///
72+
/// \param mapping The name mapping to apply.
73+
/// \param parquet_schema The Parquet schema to update in place.
74+
void ApplyNameMapping(const NameMapping& mapping,
75+
::parquet::arrow::SchemaManifest* parquet_schema);
76+
6677
/// \brief Convert an Iceberg schema to a Parquet schema.
6778
Result<std::shared_ptr<::parquet::SchemaDescriptor>> ToParquetSchema(
6879
const Schema& iceberg_schema);

‎src/iceberg/test/parquet_schema_test.cc‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
#include <parquet/types.h>
3232

3333
#include "iceberg/metadata_columns.h"
34+
#include "iceberg/name_mapping.h"
3435
#include "iceberg/parquet/parquet_schema_util_internal.h"
3536
#include "iceberg/schema.h"
3637
#include "iceberg/test/matchers.h"
@@ -203,6 +204,158 @@ TEST(HasFieldIdsTest, GroupNode) {
203204
}
204205
// NOLINTEND(clang-analyzer-cplusplus.NewDeleteLeaks)
205206

207+
TEST(ParquetNameMappingTest, ApplyNameMapping) {
208+
Schema expected_schema({
209+
SchemaField::MakeOptional(/*field_id=*/1, "id", int64()),
210+
SchemaField::MakeOptional(/*field_id=*/2, "name", string()),
211+
});
212+
213+
auto parquet_schema =
214+
MakeGroupNode("iceberg_schema", {MakeInt64Node("id"), MakeStringNode("name")});
215+
auto schema_manifest = MakeSchemaManifest(parquet_schema);
216+
auto mapping = NameMapping::Make({
217+
MappedField{.names = {"id"}, .field_id = 1},
218+
MappedField{.names = {"name"}, .field_id = 2},
219+
});
220+
221+
ApplyNameMapping(*mapping, &schema_manifest);
222+
EXPECT_EQ(schema_manifest.schema_fields[0].field->metadata()->Get(
223+
std::string(kParquetFieldIdKey)),
224+
"1");
225+
EXPECT_EQ(schema_manifest.schema_fields[1].field->metadata()->Get(
226+
std::string(kParquetFieldIdKey)),
227+
"2");
228+
229+
auto projection_result = Project(expected_schema, schema_manifest);
230+
ASSERT_THAT(projection_result, IsOk());
231+
232+
const auto& projection = *projection_result;
233+
ASSERT_EQ(projection.fields.size(), 2);
234+
ASSERT_PROJECTED_FIELD(projection.fields[0], 0);
235+
ASSERT_PROJECTED_FIELD(projection.fields[1], 1);
236+
ASSERT_EQ(SelectedColumnIndices(projection), std::vector<int32_t>({0, 1}));
237+
}
238+
239+
TEST(ParquetNameMappingTest, ApplyNameMappingWithAlias) {
240+
Schema expected_schema({
241+
SchemaField::MakeOptional(/*field_id=*/2, "full_name", string()),
242+
});
243+
244+
auto parquet_schema = MakeGroupNode("iceberg_schema", {MakeStringNode("user_name")});
245+
auto schema_manifest = MakeSchemaManifest(parquet_schema);
246+
auto mapping = NameMapping::Make({
247+
MappedField{.names = {"full_name", "user_name"}, .field_id = 2},
248+
});
249+
250+
ApplyNameMapping(*mapping, &schema_manifest);
251+
auto projection_result = Project(expected_schema, schema_manifest);
252+
ASSERT_THAT(projection_result, IsOk());
253+
254+
const auto& projection = *projection_result;
255+
ASSERT_EQ(projection.fields.size(), 1);
256+
ASSERT_PROJECTED_FIELD(projection.fields[0], 0);
257+
}
258+
259+
TEST(ParquetNameMappingTest, ApplyNameMappingToNestedFields) {
260+
Schema expected_schema({
261+
SchemaField::MakeOptional(
262+
/*field_id=*/1, "profile",
263+
std::make_shared<StructType>(std::vector<SchemaField>{
264+
SchemaField::MakeOptional(/*field_id=*/101, "name", string()),
265+
})),
266+
SchemaField::MakeOptional(
267+
/*field_id=*/2, "numbers",
268+
std::make_shared<ListType>(SchemaField::MakeOptional(
269+
/*field_id=*/201, "element", int32()))),
270+
SchemaField::MakeOptional(
271+
/*field_id=*/3, "counts",
272+
std::make_shared<MapType>(
273+
SchemaField::MakeRequired(/*field_id=*/301, "key", string()),
274+
SchemaField::MakeOptional(/*field_id=*/302, "value", int32()))),
275+
});
276+
277+
auto parquet_schema = MakeGroupNode(
278+
"iceberg_schema",
279+
{
280+
MakeGroupNode("profile", {MakeStringNode("name")}),
281+
MakeListNode("numbers", MakeInt32Node("item")),
282+
MakeMapNode("counts",
283+
MakeStringNode("key", /*field_id=*/-1, /*optional=*/false),
284+
MakeInt32Node("value")),
285+
});
286+
auto mapping = NameMapping::Make({
287+
MappedField{.names = {"profile"},
288+
.field_id = 1,
289+
.nested_mapping = MappedFields::Make({
290+
MappedField{.names = {"name"}, .field_id = 101},
291+
})},
292+
MappedField{.names = {"numbers"},
293+
.field_id = 2,
294+
.nested_mapping = MappedFields::Make({
295+
MappedField{.names = {"element"}, .field_id = 201},
296+
})},
297+
MappedField{.names = {"counts"},
298+
.field_id = 3,
299+
.nested_mapping = MappedFields::Make({
300+
MappedField{.names = {"key"}, .field_id = 301},
301+
MappedField{.names = {"value"}, .field_id = 302},
302+
})},
303+
});
304+
305+
auto schema_manifest = MakeSchemaManifest(parquet_schema);
306+
ApplyNameMapping(*mapping, &schema_manifest);
307+
auto projection_result = Project(expected_schema, schema_manifest);
308+
ASSERT_THAT(projection_result, IsOk());
309+
310+
const auto& projection = *projection_result;
311+
ASSERT_EQ(projection.fields.size(), 3);
312+
ASSERT_PROJECTED_FIELD(projection.fields[0], 0);
313+
ASSERT_EQ(projection.fields[0].children.size(), 1);
314+
ASSERT_PROJECTED_FIELD(projection.fields[0].children[0], 0);
315+
ASSERT_PROJECTED_FIELD(projection.fields[1], 1);
316+
ASSERT_EQ(projection.fields[1].children.size(), 1);
317+
ASSERT_PROJECTED_FIELD(projection.fields[1].children[0], 0);
318+
ASSERT_PROJECTED_FIELD(projection.fields[2], 2);
319+
ASSERT_EQ(projection.fields[2].children.size(), 2);
320+
ASSERT_PROJECTED_FIELD(projection.fields[2].children[0], 0);
321+
ASSERT_PROJECTED_FIELD(projection.fields[2].children[1], 1);
322+
ASSERT_EQ(SelectedColumnIndices(projection), std::vector<int32_t>({0, 1, 2, 3}));
323+
}
324+
325+
TEST(ParquetNameMappingTest, ApplyNameMappingWithUnmappedFields) {
326+
Schema expected_schema({
327+
SchemaField::MakeOptional(/*field_id=*/1, "id", int64()),
328+
SchemaField::MakeOptional(/*field_id=*/2, "extra", string()),
329+
SchemaField::MakeOptional(/*field_id=*/3, "unassigned", string()),
330+
});
331+
332+
auto parquet_schema = MakeGroupNode(
333+
"iceberg_schema",
334+
{MakeInt64Node("id"), MakeStringNode("extra"), MakeStringNode("unassigned")});
335+
auto schema_manifest = MakeSchemaManifest(parquet_schema);
336+
auto mapping = NameMapping::Make({
337+
MappedField{.names = {"id"}, .field_id = 1},
338+
MappedField{.names = {"unassigned"}, .field_id = std::nullopt},
339+
});
340+
341+
ApplyNameMapping(*mapping, &schema_manifest);
342+
for (size_t i = 1; i < schema_manifest.schema_fields.size(); ++i) {
343+
const auto& metadata = schema_manifest.schema_fields[i].field->metadata();
344+
EXPECT_TRUE(metadata == nullptr ||
345+
metadata->FindKey(std::string(kParquetFieldIdKey)) < 0);
346+
}
347+
348+
auto projection_result = Project(expected_schema, schema_manifest);
349+
ASSERT_THAT(projection_result, IsOk());
350+
351+
const auto& projection = *projection_result;
352+
ASSERT_EQ(projection.fields.size(), 3);
353+
ASSERT_PROJECTED_FIELD(projection.fields[0], 0);
354+
ASSERT_PROJECTED_NULL_FIELD(projection.fields[1]);
355+
ASSERT_PROJECTED_NULL_FIELD(projection.fields[2]);
356+
ASSERT_EQ(SelectedColumnIndices(projection), std::vector<int32_t>({0}));
357+
}
358+
206359
TEST(ParquetSchemaProjectionTest, ProjectIdenticalSchemas) {
207360
Schema expected_schema({
208361
SchemaField::MakeRequired(/*field_id=*/1, "id", iceberg::int64()),

0 commit comments

Comments
 (0)