wgtmac commented on code in PR #942:
URL: https://github.com/apache/iceberg-cpp/pull/942#discussion_r4322821474
##########
src/iceberg/arrow/arrow_c_data_util.cc:
##########
@@ -128,4 +142,154 @@ void RegisterArrowProjectBatch() {
});
}
+namespace arrow {
+namespace {
+
+Result<std::shared_ptr<::arrow::Schema>> ImportIcebergSchema(const Schema&
schema) {
+ ArrowSchema c_schema;
+ ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(schema, &c_schema));
+ internal::ArrowSchemaGuard c_schema_guard(&c_schema);
+ ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema,
::arrow::ImportSchema(&c_schema));
+ return arrow_schema;
+}
+
+// Writers must not silently drop data at any nesting level.
+Status RejectUnknownInputFields(const NestedType& input_type,
+ const NestedType& write_type) {
+ for (const auto& input_field : input_type.fields()) {
+ ICEBERG_ASSIGN_OR_RAISE(auto write_field,
+ write_type.GetFieldById(input_field.field_id()));
+ if (!write_field.has_value()) {
+ return InvalidSchema("Input field id {} is not present in write schema",
+ input_field.field_id());
+ }
+ if (input_field.type()->is_nested() &&
+ input_field.type()->type_id() == write_field->get().type()->type_id())
{
+ ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(
+ internal::checked_cast<const NestedType&>(*input_field.type()),
+ internal::checked_cast<const
NestedType&>(*write_field->get().type())));
+ }
+ }
+ return {};
+}
+
+Result<std::vector<std::shared_ptr<::arrow::Array>>> MaterializeProjection(
+ std::span<const FieldProjection> projections,
+ const std::vector<std::shared_ptr<::arrow::Array>>& input_arrays,
+ const std::vector<std::shared_ptr<::arrow::Field>>& write_arrow_fields,
+ int64_t length, ::arrow::MemoryPool* pool) {
+ std::vector<std::shared_ptr<::arrow::Array>> result;
+ result.reserve(projections.size());
+ for (size_t index = 0; index < projections.size(); ++index) {
+ const auto& projection = projections[index];
+ const auto& write_arrow_type = write_arrow_fields[index]->type();
+ switch (projection.kind) {
+ case FieldProjection::Kind::kProjected: {
+ const auto& input_array =
input_arrays[std::get<size_t>(projection.from)];
+ if (projection.children.empty()) {
+ result.emplace_back(input_array);
+ break;
+ }
+ auto data = input_array->data()->Copy();
+ data->type = write_arrow_type;
+ auto child_data = data;
+ auto child_type = write_arrow_type;
+ int64_t child_length = data->offset + data->length;
+ switch (write_arrow_type->id()) {
+ case ::arrow::Type::STRUCT:
+ break;
+ case ::arrow::Type::LIST:
+ child_length = data->child_data[0]->length;
+ break;
+ case ::arrow::Type::MAP:
+ // Arrow maps have an entries struct between the map and its
key/value.
+ child_data = data->child_data[0]->Copy();
+ child_type = write_arrow_type->field(0)->type();
+ child_length = child_data->offset + child_data->length;
+ break;
+ default:
+ return InvalidArrowData("Expected nested array for projected field
{}",
+ index);
+ }
+ std::vector<std::shared_ptr<::arrow::Array>> input_children;
+ input_children.reserve(child_data->child_data.size());
+ for (const auto& child : child_data->child_data) {
+ input_children.emplace_back(::arrow::MakeArray(child));
+ }
+ // Raw child data retains the parent's physical row positions.
+ // StructArray::fields() would apply the parent slice, causing its
offset to be
+ // applied twice.
+ ICEBERG_ASSIGN_OR_RAISE(
+ auto children,
+ MaterializeProjection(projection.children, input_children,
+ child_type->fields(), child_length, pool));
+ child_data->type = child_type;
+ child_data->child_data.clear();
+ for (const auto& child : children) {
+ child_data->child_data.emplace_back(child->data());
+ }
+ if (write_arrow_type->id() == ::arrow::Type::MAP) {
+ data->child_data = {std::move(child_data)};
+ }
+ result.emplace_back(::arrow::MakeArray(std::move(data)));
+ break;
+ }
+ case FieldProjection::Kind::kDefault: {
+ ICEBERG_ASSIGN_OR_RAISE(auto default_array,
+
MakeDefaultArray(std::get<Literal>(projection.from),
+ write_arrow_type, length,
pool));
+ result.emplace_back(std::move(default_array));
+ break;
+ }
+ case FieldProjection::Kind::kNull: {
+ ICEBERG_ARROW_ASSIGN_OR_RETURN(
+ auto null_array, ::arrow::MakeArrayOfNull(write_arrow_type,
length, pool));
+ result.emplace_back(std::move(null_array));
+ break;
+ }
+ default:
+ return NotSupported("Projection kind {} is not supported for write
alignment",
+ ToString(projection.kind));
+ }
+ }
+ return result;
+}
+
+} // namespace
+
+Result<ArrowArray> AlignBatchForWrite(ArrowArray* input_batch, const Schema&
input_schema,
+ const Schema& write_schema) {
+ ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null");
+ internal::ArrowArrayGuard input_guard(input_batch);
+ if (input_schema == write_schema) {
+ return std::exchange(*input_batch, ArrowArray{});
+ }
+
+ ICEBERG_ASSIGN_OR_RAISE(auto input_arrow_schema,
ImportIcebergSchema(input_schema));
+ ICEBERG_ARROW_ASSIGN_OR_RETURN(
+ auto input_record_batch,
+ ::arrow::ImportRecordBatch(input_batch, input_arrow_schema));
+
+ ICEBERG_ASSIGN_OR_RAISE(auto write_arrow_schema,
ImportIcebergSchema(write_schema));
+
+ ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(input_schema,
write_schema));
+ ICEBERG_ASSIGN_OR_RAISE(auto projection,
Review Comment:
In write mode, Project still turns a concrete input into null when the
target type is unknown, before reaching the type check. An optional int
containing 123 is therefore silently written as null. Reject
concrete-to-unknown writes and add the reverse of the unknown-input test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
