anoopj commented on code in PR #3231: URL: https://github.com/apache/iceberg-rust/pull/3231#discussion_r4031374338
########## crates/iceberg/src/writer/delta_writer/record_ops.rs: ########## @@ -0,0 +1,602 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Splitting a change-type [`RecordBatch`] into insert and delete batches. +//! +//! A delta writer receives a single stream of row-level changes where each row +//! carries a `_change_type` indicator describing what kind of change it is. This +//! is the repo's [`_change_type`](RESERVED_COL_NAME_CHANGE_TYPE) convention: a +//! non-nullable [`Utf8`] column whose value is one of [`CHANGE_TYPE_INSERT`], +//! [`CHANGE_TYPE_DELETE`], [`CHANGE_TYPE_UPDATE_BEFORE`], or +//! [`CHANGE_TYPE_UPDATE_AFTER`]. +//! +//! [`split_by_change_type`] partitions such a batch into the insert and delete +//! payloads that the underlying data and delete writers consume, collapsing the +//! four change types onto two sides the way Java's `BaseDeltaTaskWriter` does: +//! `INSERT` and `UPDATE_AFTER` become inserts, `DELETE` and `UPDATE_BEFORE` +//! become deletes. The `_change_type` column is stripped from both outputs. +//! +//! This is a building block for the `DeltaWriter` epic (see +//! <https://github.com/apache/iceberg-rust/issues/2218>). +//! +//! [`RecordBatch`]: arrow_array::RecordBatch +//! [`Utf8`]: arrow_schema::DataType::Utf8 + +use std::sync::Arc; + +use arrow_arith::boolean::not; +use arrow_array::{Array, BooleanArray, RecordBatch, StringArray}; +use arrow_buffer::BooleanBufferBuilder; +use arrow_schema::{DataType, Schema}; +use arrow_select::filter::filter_record_batch; + +use crate::metadata_columns::RESERVED_COL_NAME_CHANGE_TYPE; +use crate::{Error, ErrorKind, Result}; + +/// `_change_type` value marking a row as an insert. +pub(crate) const CHANGE_TYPE_INSERT: &str = "INSERT"; +/// `_change_type` value marking a row as a delete. +pub(crate) const CHANGE_TYPE_DELETE: &str = "DELETE"; +/// `_change_type` value marking the pre-image of an updated row. +pub(crate) const CHANGE_TYPE_UPDATE_BEFORE: &str = "UPDATE_BEFORE"; +/// `_change_type` value marking the post-image of an updated row. +pub(crate) const CHANGE_TYPE_UPDATE_AFTER: &str = "UPDATE_AFTER"; + +/// The result of splitting a change-type batch with [`split_by_change_type`]. +/// +/// Both batches share the same schema: the schema of the input batch with the +/// `_change_type` column removed. Field metadata (including Parquet field ids) +/// and the schema-level metadata are preserved. +#[derive(Debug)] +pub struct SplitBatches { + /// Rows whose `_change_type` was [`CHANGE_TYPE_INSERT`] or + /// [`CHANGE_TYPE_UPDATE_AFTER`], with the `_change_type` column removed. + pub inserts: RecordBatch, + /// Rows whose `_change_type` was [`CHANGE_TYPE_DELETE`] or + /// [`CHANGE_TYPE_UPDATE_BEFORE`], with the `_change_type` column removed. + pub deletes: RecordBatch, +} + +/// Split a change-type [`RecordBatch`] into separate insert and delete batches. +/// +/// The batch must contain a [`_change_type`](RESERVED_COL_NAME_CHANGE_TYPE) +/// column, located **by name** wherever it sits in the schema (not by position). +/// It must be a non-nullable [`Utf8`] column whose values are all one of the four +/// spec change types. The rows are collapsed onto two sides the way Java's +/// `BaseDeltaTaskWriter` does: +/// - [`CHANGE_TYPE_INSERT`] and [`CHANGE_TYPE_UPDATE_AFTER`] go to `inserts`; +/// - [`CHANGE_TYPE_DELETE`] and [`CHANGE_TYPE_UPDATE_BEFORE`] go to `deletes`. +/// +/// Every other column is treated as payload. The returned [`SplitBatches`] carry +/// the payload columns (and their field-id metadata) but not the `_change_type` +/// column. +/// +/// # Returns +/// +/// A [`SplitBatches`] whose `inserts` and `deletes` preserve the input row order +/// within each side. Either side may be empty when all rows share a side (e.g. an +/// all-insert batch yields an empty `deletes` that still carries the payload +/// schema). This splitter neither reorders rows nor enforces any +/// `UPDATE_BEFORE`/`UPDATE_AFTER` pairing or ordering; that is the `DeltaWriter`'s +/// concern. +/// +/// # Errors +/// +/// Returns [`ErrorKind::DataInvalid`] when: +/// - the batch has no `_change_type` column; +/// - removing the `_change_type` column would leave no payload columns (a batch +/// consisting of only the `_change_type` column has nothing to write); +/// - the `_change_type` column is not of type [`Utf8`]; +/// - the `_change_type` column is declared nullable or contains null values; +/// - the `_change_type` column holds a value other than one of the four spec +/// change types. +/// +/// Returns [`ErrorKind::Unexpected`] if the `_change_type` column passes the +/// [`Utf8`] type check but cannot be downcast to a [`StringArray`], which would +/// be an internal invariant violation. +/// +/// [`RecordBatch`]: arrow_array::RecordBatch +/// [`Utf8`]: arrow_schema::DataType::Utf8 +pub fn split_by_change_type(batch: RecordBatch) -> Result<SplitBatches> { + let schema = batch.schema(); + + // Locate the change-type column by name, wherever it sits in the schema. + let (change_type_idx, change_type_field) = schema + .column_with_name(RESERVED_COL_NAME_CHANGE_TYPE) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "the batch has no '{RESERVED_COL_NAME_CHANGE_TYPE}' column required to split delta changes" + ), + ) + })?; + + // There must be at least one payload column besides the change-type column. + // A batch consisting of only the change-type column carries no payload to + // write, which is a caller error. + if batch.num_columns() < 2 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "removing the '{RESERVED_COL_NAME_CHANGE_TYPE}' column leaves no payload columns to write" + ), + )); + } + + if change_type_field.data_type() != &DataType::Utf8 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "the '{RESERVED_COL_NAME_CHANGE_TYPE}' column must be of type Utf8, but found {}", + change_type_field.data_type() + ), + )); + } + + let change_type = batch + .column(change_type_idx) + .as_any() + .downcast_ref::<StringArray>() + .ok_or_else(|| { + Error::new( + ErrorKind::Unexpected, + format!( + "internal: the '{RESERVED_COL_NAME_CHANGE_TYPE}' column passed the Utf8 type check but could not be downcast to a StringArray" + ), + ) + })?; + + // The contract is a non-nullable column with no nulls. Reject any actual + // nulls first (so `value(i)` below is safe), then reject a nullable-declared + // column even when it happens to carry no nulls. + if change_type.null_count() > 0 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("the '{RESERVED_COL_NAME_CHANGE_TYPE}' column must not contain null values"), + )); + } + if change_type_field.is_nullable() { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("the '{RESERVED_COL_NAME_CHANGE_TYPE}' column must be declared non-nullable"), + )); + } + + // Build the insert mask in a single pass, validating the change-type domain + // as we go. `change_type` has no nulls, so `value(i)` is safe. + let mut insert_builder = BooleanBufferBuilder::new(change_type.len()); + for i in 0..change_type.len() { + let value = change_type.value(i); + let is_insert = match value { + CHANGE_TYPE_INSERT | CHANGE_TYPE_UPDATE_AFTER => true, + CHANGE_TYPE_DELETE | CHANGE_TYPE_UPDATE_BEFORE => false, + other => { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "unexpected '{RESERVED_COL_NAME_CHANGE_TYPE}' value {other:?}: expected one of {CHANGE_TYPE_INSERT}, {CHANGE_TYPE_DELETE}, {CHANGE_TYPE_UPDATE_BEFORE}, {CHANGE_TYPE_UPDATE_AFTER}" + ), + )); + } + }; + insert_builder.append(is_insert); + } + let insert_mask = BooleanArray::new(insert_builder.finish(), None); + // After validation every row is exactly one side, so the delete side is the + // complement of the insert side. + let delete_mask = not(&insert_mask)?; + + // Rebuild the payload schema from the original fields minus the change-type + // column, preserving each field's metadata (field ids) and the schema-level + // metadata. + let payload_fields = schema + .fields() + .iter() + .enumerate() + .filter(|(idx, _)| *idx != change_type_idx) + .map(|(_, field)| field.clone()) + .collect::<Vec<_>>(); + let payload_schema = Arc::new(Schema::new_with_metadata( + payload_fields, + schema.metadata().clone(), + )); + let payload_columns = batch + .columns() + .iter() + .enumerate() + .filter(|(idx, _)| *idx != change_type_idx) + .map(|(_, column)| column.clone()) + .collect::<Vec<_>>(); Review Comment: This reimplements `RecordBatch::project`. You could just do something like: ``` let payload_indices: Vec<usize> = (0..batch.num_columns()) .filter(|idx| *idx != change_type_idx) .collect(); let payload = batch.project(&payload_indices)?; ``` -- 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]
