Skip to content
Open
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
53 changes: 38 additions & 15 deletions src/iceberg/arrow_c_data_util.cc
Original file line number Diff line number Diff line change
Expand Up @@ -232,9 +232,7 @@ Status AppendValue(const ArrowSchema& input_schema, const ArrowArray& input_arra
} // namespace

ProjectionContext::ProjectionContext(ProjectionContext&& other) noexcept
: input_schema_(std::exchange(other.input_schema_, nullptr)),
output_schema_(std::exchange(other.output_schema_, nullptr)),
selected_field_indices_(std::move(other.selected_field_indices_)),
: selected_field_indices_(std::move(other.selected_field_indices_)),
input_arrow_schema_(other.input_arrow_schema_),
output_arrow_schema_(other.output_arrow_schema_),
project_batch_function_(std::exchange(other.project_batch_function_, nullptr)),
Expand All @@ -256,8 +254,6 @@ ProjectionContext& ProjectionContext::operator=(ProjectionContext&& other) noexc
output_arrow_schema_.release(&output_arrow_schema_);
}

input_schema_ = std::exchange(other.input_schema_, nullptr);
output_schema_ = std::exchange(other.output_schema_, nullptr);
selected_field_indices_ = std::move(other.selected_field_indices_);
input_arrow_schema_ = other.input_arrow_schema_;
other.input_arrow_schema_.release = nullptr;
Expand All @@ -280,27 +276,54 @@ ProjectionContext::~ProjectionContext() {

Result<ProjectionContext> ProjectionContext::Make(
const Schema& input_schema, const Schema& output_schema,
const ArrowSchema& physical_input_schema,
ProjectionContext::ProjectBatchFunction project_batch_function, bool validate_batch) {
ArrowSchemaView schema_view;
ArrowError error;
ArrowErrorInit(&error);
ICEBERG_NANOARROW_RETURN_UNEXPECTED_WITH_ERROR(
ArrowSchemaViewInit(&schema_view, &physical_input_schema, &error), error);
ICEBERG_PRECHECK(schema_view.type == NANOARROW_TYPE_STRUCT,
"Physical input schema must be a struct, but is {}",
ArrowTypeString(schema_view.type));
ICEBERG_PRECHECK(
std::cmp_equal(physical_input_schema.n_children, input_schema.fields().size()),
"Physical input schema must have {} fields to match the logical schema, but has {}",
input_schema.fields().size(), physical_input_schema.n_children);

ICEBERG_ASSIGN_OR_RAISE(
auto selected_field_indices,
BuildSelectedFieldIndices(input_schema.fields(), output_schema.fields()));

ProjectionContext context;
context.input_schema_ = &input_schema;
context.output_schema_ = &output_schema;
context.selected_field_indices_ = std::move(selected_field_indices);
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(input_schema, &context.input_arrow_schema_));
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_));
context.project_batch_function_ = project_batch_function;
context.validate_batch_ = validate_batch;
ICEBERG_NANOARROW_RETURN_UNEXPECTED(
ArrowSchemaDeepCopy(&physical_input_schema, &context.input_arrow_schema_));
ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(output_schema, &context.output_arrow_schema_));

// Replace each logical output field with the physical layout of its selected input
// field. The deep copy also carries the physical metadata (e.g. field ids and
// extension annotations); only the top-level name and nullability come from
// output_schema.
const auto& output_fields = output_schema.fields();
for (size_t i = 0; i < output_fields.size(); ++i) {
auto* output_field = context.output_arrow_schema_.children[i];
output_field->release(output_field);
ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowSchemaDeepCopy(
physical_input_schema.children[selected_field_indices[i]], output_field));
ICEBERG_NANOARROW_RETURN_UNEXPECTED(
ArrowSchemaSetName(output_field, std::string(output_fields[i].name()).c_str()));
// Preserve type-specific flags such as sorted map keys when setting nullability.
if (output_fields[i].optional()) {
output_field->flags |= ARROW_FLAG_NULLABLE;
} else {
output_field->flags &= ~ARROW_FLAG_NULLABLE;
}
}
context.selected_field_indices_ = std::move(selected_field_indices);
return context;
}

const Schema& ProjectionContext::input_schema() const { return *input_schema_; }

const Schema& ProjectionContext::output_schema() const { return *output_schema_; }

const ArrowSchema& ProjectionContext::input_arrow_schema() const {
return input_arrow_schema_;
}
Expand Down
25 changes: 12 additions & 13 deletions src/iceberg/arrow_c_data_util_internal.h
Original file line number Diff line number Diff line change
Expand Up @@ -61,21 +61,28 @@ class ICEBERG_EXPORT ProjectionContext {
/// \brief Resolve the registered ProjectBatch implementation.
static auto ResolveProjectBatchFunction() -> ProjectBatchFunction;

/// \brief Build reusable projection state for a validated schema pair.
/// \brief Build reusable projection state using the reader's physical Arrow layout.
///
/// \param input_schema Schema that describes every input batch.
/// \param output_schema Final schema and column order requested by the caller.
/// \param physical_input_schema Physical fields corresponding to input_schema in
/// the same order. Each selected output field becomes a deep copy of its physical
/// counterpart, preserving the physical layout, nested list offset widths,
/// metadata, and type-specific flags; only the top-level name and nullability come
/// from output_schema.
/// \param project_batch_function Optional implementation returned by
/// ProjectionContext::ResolveProjectBatchFunction, or nullptr to use the nanoarrow
/// path.
/// \param validate_batch Whether to perform full content validation. Disable only for
/// trusted input.
/// \note It validates that output_schema selects or reorders complete top-level fields
/// by field id. Nested pruning and type changes are rejected. The input_schema and
/// output_schema passed to Make must outlive the context. ProjectBatch may lazily
/// initialize backend cache; do not share one context across concurrent calls.
/// by field id. Nested pruning and type changes are rejected. Schema inputs are only
/// used during construction; the context owns its physical Arrow schemas.
/// ProjectBatch may lazily initialize backend cache; do not share one context across
/// concurrent calls.
static Result<ProjectionContext> Make(const Schema& input_schema,
const Schema& output_schema,
const ArrowSchema& physical_input_schema,
ProjectBatchFunction project_batch_function,
bool validate_batch = true);

Expand All @@ -86,10 +93,6 @@ class ICEBERG_EXPORT ProjectionContext {
ProjectionContext(const ProjectionContext&) = delete;
ProjectionContext& operator=(const ProjectionContext&) = delete;

const Schema& input_schema() const;

const Schema& output_schema() const;

const ArrowSchema& input_arrow_schema() const;

const ArrowSchema& output_arrow_schema() const;
Expand All @@ -105,10 +108,6 @@ class ICEBERG_EXPORT ProjectionContext {
private:
ProjectionContext() = default;

// Raw schema pointers are borrowed from caller-owned schemas. FileScanTaskReader
// keeps those schema objects alive in the same stream source that owns this context.
const Schema* input_schema_ = nullptr;
const Schema* output_schema_ = nullptr;
std::vector<int32_t> selected_field_indices_;
ArrowSchema input_arrow_schema_{};
ArrowSchema output_arrow_schema_{};
Expand Down Expand Up @@ -244,7 +243,7 @@ ICEBERG_EXPORT Result<ArrowArrayStream> MakeArrowArrayStream(
/// \param row_indices Zero-based row positions to copy from `input_batch`.
/// \param projection Reusable schema/projection state created by
/// ProjectionContext::Make.
/// \return A newly owned ArrowArray matching `projection.output_schema()`.
/// \return A newly owned ArrowArray matching `projection.output_arrow_schema()`.
ICEBERG_EXPORT Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
std::span<const int32_t> row_indices,
ProjectionContext& projection);
Expand Down
Loading
Loading