Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 61 additions & 12 deletions src/iceberg/schema.cc
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,28 @@

namespace iceberg {

Schema::Schema(std::vector<SchemaField> fields, std::optional<int32_t> schema_id)
: StructType(std::move(fields)), schema_id_(schema_id) {}
Schema::Schema(std::vector<SchemaField> fields, std::optional<int32_t> schema_id,
std::vector<int32_t> identifier_field_ids)
: StructType(std::move(fields)),
schema_id_(schema_id),
identifier_field_ids_(std::move(identifier_field_ids)) {}

Result<std::unique_ptr<Schema>> Schema::Make(
std::vector<SchemaField> fields, std::optional<int32_t> schema_id,
const std::vector<std::string>& identifier_field_names) {
auto schema = std::make_unique<Schema>(std::move(fields), schema_id);

std::vector<int32_t> fresh_identifier_ids;
for (const auto& name : identifier_field_names) {
ICEBERG_ASSIGN_OR_RAISE(auto field, schema->FindFieldByName(name));
if (!field) {
return InvalidSchema("Cannot find identifier field: {}", name);
}
fresh_identifier_ids.push_back(field.value().get().field_id());
}
schema->identifier_field_ids_ = std::move(fresh_identifier_ids);
return schema;
}

std::optional<int32_t> Schema::schema_id() const { return schema_id_; }

Expand All @@ -48,15 +68,16 @@ std::string Schema::ToString() const {
}

bool Schema::Equals(const Schema& other) const {
return schema_id_ == other.schema_id_ && fields_ == other.fields_;
return schema_id_ == other.schema_id_ && fields_ == other.fields_ &&
identifier_field_ids_ == other.identifier_field_ids_;
}

Result<std::optional<std::reference_wrapper<const SchemaField>>> Schema::FindFieldByName(
std::string_view name, bool case_sensitive) const {
if (case_sensitive) {
ICEBERG_ASSIGN_OR_RAISE(auto name_to_id, name_to_id_.Get(*this));
auto it = name_to_id.get().find(name);
if (it == name_to_id.get().end()) {
ICEBERG_ASSIGN_OR_RAISE(auto name_id_map, name_id_map_.Get(*this));
auto it = name_id_map.get().name_to_id.find(name);
if (it == name_id_map.get().name_to_id.end()) {
return std::nullopt;
};
return FindFieldById(it->second);
Expand All @@ -77,21 +98,22 @@ Schema::InitIdToFieldMap(const Schema& self) {
return id_to_field;
}

Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
Schema::InitNameToIdMap(const Schema& self) {
std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>> name_to_id;
NameToIdVisitor visitor(name_to_id, /*case_sensitive=*/true);
Result<Schema::NameIdMap> Schema::InitNameIdMap(const Schema& self) {
NameIdMap name_id_map;
NameToIdVisitor visitor(name_id_map.name_to_id, &name_id_map.id_to_name,
/*case_sensitive=*/true);
ICEBERG_RETURN_UNEXPECTED(
VisitTypeInline(self, &visitor, /*path=*/"", /*short_path=*/""));
visitor.Finish();
return name_to_id;
return name_id_map;
}

Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
Schema::InitLowerCaseNameToIdMap(const Schema& self) {
std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>
lowercase_name_to_id;
NameToIdVisitor visitor(lowercase_name_to_id, /*case_sensitive=*/false);
NameToIdVisitor visitor(lowercase_name_to_id, /*id_to_name=*/nullptr,
/*case_sensitive=*/false);
ICEBERG_RETURN_UNEXPECTED(
VisitTypeInline(self, &visitor, /*path=*/"", /*short_path=*/""));
visitor.Finish();
Expand All @@ -108,6 +130,16 @@ Result<std::optional<std::reference_wrapper<const SchemaField>>> Schema::FindFie
return it->second;
}

Result<std::optional<std::string_view>> Schema::FindColumnNameById(
int32_t field_id) const {
ICEBERG_ASSIGN_OR_RAISE(auto name_id_map, name_id_map_.Get(*this));
auto it = name_id_map.get().id_to_name.find(field_id);
if (it == name_id_map.get().id_to_name.end()) {
return std::nullopt;
}
return it->second;
}

Result<std::unordered_map<int32_t, std::vector<size_t>>> Schema::InitIdToPositionPath(
const Schema& self) {
PositionPathVisitor visitor;
Expand Down Expand Up @@ -179,4 +211,21 @@ Result<std::unique_ptr<Schema>> Schema::Project(
std::nullopt);
}

const std::vector<int32_t>& Schema::IdentifierFieldIds() const {
return identifier_field_ids_;
}

Result<std::vector<std::string>> Schema::IdentifierFieldNames() const {
std::vector<std::string> names;
names.reserve(identifier_field_ids_.size());
for (auto id : identifier_field_ids_) {
ICEBERG_ASSIGN_OR_RAISE(auto name, FindColumnNameById(id));
if (!name.has_value()) {
return InvalidSchema("Cannot find the field of the specified field id: {}", id);
}
names.emplace_back(name.value());
}
return names;
}

} // namespace iceberg
48 changes: 44 additions & 4 deletions src/iceberg/schema.h
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,18 @@ class ICEBERG_EXPORT Schema : public StructType {
static constexpr int32_t kInvalidColumnId = -1;

explicit Schema(std::vector<SchemaField> fields,
std::optional<int32_t> schema_id = std::nullopt);
std::optional<int32_t> schema_id = std::nullopt,
std::vector<int32_t> identifier_field_ids = {});

/// \brief Create a schema.
///
/// \param fields The fields that make up the schema.
/// \param schema_id The unique identifier for this schema (default: kInitialSchemaId).
/// \param identifier_field_names Canonical names of fields that uniquely identify rows
/// in the table (default: empty). \return A new Schema instance or Status if failed.
static Result<std::unique_ptr<Schema>> Make(
std::vector<SchemaField> fields, std::optional<int32_t> schema_id = std::nullopt,
const std::vector<std::string>& identifier_field_names = {});

/// \brief Get the schema ID.
///
Expand Down Expand Up @@ -78,6 +89,13 @@ class ICEBERG_EXPORT Schema : public StructType {
Result<std::optional<std::reference_wrapper<const SchemaField>>> FindFieldById(
int32_t field_id) const;

/// \brief Returns the canonical field name for the given id.
///
/// \param field_id The id of the field to get the canonical name for.
/// \return The canocinal column name of the field with the given id, or std::nullopt if
/// not found.
Result<std::optional<std::string_view>> FindColumnNameById(int32_t field_id) const;

/// \brief Get the accessor to access the field by field id.
///
/// \param field_id The id of the field to get the accessor for.
Expand All @@ -103,26 +121,48 @@ class ICEBERG_EXPORT Schema : public StructType {
Result<std::unique_ptr<Schema>> Project(
const std::unordered_set<int32_t>& field_ids) const;

/// \brief Return the field IDs of the identifier fields.
const std::vector<int32_t>& IdentifierFieldIds() const;

/// \brief Return the canonical field names of the identifier fields.
Result<std::vector<std::string>> IdentifierFieldNames() const;

friend bool operator==(const Schema& lhs, const Schema& rhs) { return lhs.Equals(rhs); }

private:
/// \brief Compare two schemas for equality.
bool Equals(const Schema& other) const;

struct NameIdMap {
/// \brief Mapping from canonical field name to ID
///
/// \note Short names for maps and lists are included for any name that does not
/// conflict with a canonical name. For example, a list, 'l', of structs with field
/// 'x' will produce short name 'l.x' in addition to canonical name 'l.element.x'.
std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>> name_to_id;

/// \brief Mapping from field ID to canonical name
///
/// \note Canonical names, but not short names are set, for example
/// 'list.element.field' instead of 'list.field'.
std::unordered_map<int32_t, std::string> id_to_name;
};

static Result<std::unordered_map<int32_t, std::reference_wrapper<const SchemaField>>>
InitIdToFieldMap(const Schema&);
static Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
InitNameToIdMap(const Schema&);
static Result<NameIdMap> InitNameIdMap(const Schema&);
static Result<std::unordered_map<std::string, int32_t, StringHash, std::equal_to<>>>
InitLowerCaseNameToIdMap(const Schema&);
static Result<std::unordered_map<int32_t, std::vector<size_t>>> InitIdToPositionPath(
const Schema&);

const std::optional<int32_t> schema_id_;
/// Field IDs that uniquely identify rows in the table.
std::vector<int32_t> identifier_field_ids_;
/// Mapping from field id to field.
Lazy<InitIdToFieldMap> id_to_field_;
/// Mapping from field name to field id.
Lazy<InitNameToIdMap> name_to_id_;
Lazy<InitNameIdMap> name_id_map_;
/// Mapping from lowercased field name to field id
Lazy<InitLowerCaseNameToIdMap> lowercase_name_to_id_;
/// Mapping from field id to (nested) position path to access the field.
Expand Down
1 change: 1 addition & 0 deletions src/iceberg/test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ endfunction()

add_iceberg_test(schema_test
SOURCES
assign_id_visitor_test.cc
name_mapping_test.cc
partition_field_test.cc
partition_spec_test.cc
Expand Down
Loading
Loading