laskoviymishka commented on code in PR #2986: URL: https://github.com/apache/iceberg-rust/pull/2986#discussion_r3980265573
########## crates/iceberg/src/writer/base_writer/position_delete_writer.rs: ########## @@ -0,0 +1,914 @@ +// 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. + +//! This module provides `PositionDeleteFileWriter`. +//! +//! A position delete file has two required columns: `file_path` (`string`, field id +//! [`RESERVED_FIELD_ID_DELETE_FILE_PATH`]) and `pos` (`long`, field id +//! [`RESERVED_FIELD_ID_DELETE_FILE_POS`]). The writer takes batches already shaped as +//! those two columns (see [`position_delete_schema`]) and sets +//! [`DataContentType::PositionDeletes`] on the output. It does not sort its input; see +//! [`PositionDeleteFileWriter::write`]. +//! +//! Position delete files are a v2 construct. v3 replaces them with deletion vectors and +//! forbids adding new position delete files, so callers must not route v3 writes here. +//! This base writer has no format-version gate by design; that gating belongs at the +//! transaction/commit layer. + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::{DataType, Field}; +use once_cell::sync::Lazy; +use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + +use crate::metadata_columns::{ + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, delete_file_path_field, + delete_file_pos_field, +}; +use crate::spec::{DataContentType, DataFile, PartitionKey, Schema, SchemaRef}; +use crate::writer::file_writer::FileWriterBuilder; +use crate::writer::file_writer::location_generator::{FileNameGenerator, LocationGenerator}; +use crate::writer::file_writer::rolling_writer::{RollingFileWriter, RollingFileWriterBuilder}; +use crate::writer::{IcebergWriter, IcebergWriterBuilder}; +use crate::{Error, ErrorKind, Result}; + +/// The canonical Iceberg schema of a position delete file: the required `file_path` +/// (`string`) and `pos` (`long`) columns with their reserved field ids. +static POSITION_DELETE_SCHEMA: Lazy<SchemaRef> = Lazy::new(|| { + Arc::new( + Schema::builder() + .with_fields(vec![ + delete_file_path_field().clone(), + delete_file_pos_field().clone(), + ]) + .build() + .expect("position delete schema is statically valid"), + ) +}); + +/// [`POSITION_DELETE_SCHEMA`] converted to Arrow, keeping the reserved field ids in each +/// field's Parquet field-id metadata. Test-only for now: callers configure the writer with +/// the Iceberg [`position_delete_schema`], so no non-test code needs the Arrow form yet. +#[cfg(test)] +static POSITION_DELETE_ARROW_SCHEMA: Lazy<arrow_schema::SchemaRef> = Lazy::new(|| { + Arc::new( + crate::arrow::schema_to_arrow_schema(&POSITION_DELETE_SCHEMA) + .expect("position delete arrow schema is statically valid"), + ) +}); + +/// Returns the canonical Iceberg schema of a position delete file. +/// +/// Use this to build the [`ParquetWriterBuilder`](crate::writer::file_writer::ParquetWriterBuilder) +/// that backs a [`PositionDeleteFileWriter`], so the written file matches the +/// spec exactly. +pub fn position_delete_schema() -> SchemaRef { + POSITION_DELETE_SCHEMA.clone() +} + +/// Returns the canonical Arrow schema of a position delete file. +#[cfg(test)] +fn position_delete_arrow_schema() -> arrow_schema::SchemaRef { + POSITION_DELETE_ARROW_SCHEMA.clone() +} + +/// Reads a field's Iceberg field id from its Parquet field-id metadata. +fn field_id(field: &Field) -> Result<i32> { + field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Position delete column `{}` is missing its Iceberg field id metadata.", + field.name() + ), + ) + })? + .parse::<i32>() + .map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Position delete column `{}` has an invalid field id: {e}", + field.name() + ), + ) + }) +} + +/// Validates that a batch is a position delete file: the `file_path` (`Utf8`) and +/// `pos` (`Int64`) columns, in order, with the two reserved field ids. Checking it +/// here gives a clear error before the batch reaches the Parquet writer. +fn validate_position_delete_batch(batch: &RecordBatch) -> Result<()> { + let fields = batch.schema_ref().fields(); + if fields.len() != 2 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "This writer supports only the two required position delete columns (`file_path`, `pos`); \ + batches with a different column count (e.g. including the optional `row` column) are not supported. Got {} columns.", + fields.len() + ), + )); + } + + let path = &fields[0]; + let path_id = field_id(path)?; + if path_id != RESERVED_FIELD_ID_DELETE_FILE_PATH { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The first position delete column must be `file_path` (field id {RESERVED_FIELD_ID_DELETE_FILE_PATH}), but got field id {path_id}." + ), + )); + } + // The canonical schema maps Iceberg `string` to `Utf8` and the file writer is + // configured with it, so a `LargeUtf8` column has to be cast to `Utf8` first. + if path.data_type() != &DataType::Utf8 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The position delete `file_path` column must be Utf8 (cast it first); got {:?}.", + path.data_type() + ), + )); + } + // Required column: a nullable field could write nulls under a required schema. + if path.is_nullable() { + return Err(Error::new( + ErrorKind::DataInvalid, + "The position delete `file_path` column must be required (non-nullable).", + )); + } + + let pos = &fields[1]; + let pos_id = field_id(pos)?; + if pos_id != RESERVED_FIELD_ID_DELETE_FILE_POS { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The second position delete column must be `pos` (field id {RESERVED_FIELD_ID_DELETE_FILE_POS}), but got field id {pos_id}." + ), + )); + } + if pos.data_type() != &DataType::Int64 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The position delete `pos` column must be Int64, but got {:?}.", + pos.data_type() + ), + )); + } + if pos.is_nullable() { + return Err(Error::new( + ErrorKind::DataInvalid, + "The position delete `pos` column must be required (non-nullable).", + )); + } + + Ok(()) +} + +/// Builder for [`PositionDeleteFileWriter`]. +#[derive(Debug)] +pub struct PositionDeleteFileWriterBuilder< + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +> { + inner: RollingFileWriterBuilder<B, L, F>, +} + +impl<B, L, F> PositionDeleteFileWriterBuilder<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + /// Create a new `PositionDeleteFileWriterBuilder` using a `RollingFileWriterBuilder`. + /// + /// The `RollingFileWriterBuilder` must be backed by a file writer configured + /// with the [`position_delete_schema`]; the per-batch validation in + /// [`PositionDeleteFileWriter::write`] guards against a mismatched batch, but + /// the caller is responsible for wiring the same schema into the file writer. + pub fn new(inner: RollingFileWriterBuilder<B, L, F>) -> Self { + Self { inner } + } +} + +#[async_trait::async_trait] +impl<B, L, F> IcebergWriterBuilder for PositionDeleteFileWriterBuilder<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + type R = PositionDeleteFileWriter<B, L, F>; + + async fn build(&self, partition_key: Option<PartitionKey>) -> Result<Self::R> { + Ok(PositionDeleteFileWriter { + inner: Some(self.inner.build()), + partition_key, + }) + } +} + +/// Writer used to write position delete files within one spec/partition. +#[derive(Debug)] +pub struct PositionDeleteFileWriter< + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +> { + inner: Option<RollingFileWriter<B, L, F>>, + partition_key: Option<PartitionKey>, +} + +#[async_trait::async_trait] +impl<B, L, F> IcebergWriter for PositionDeleteFileWriter<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + /// Writes a batch of `(file_path, pos)` records; the shape is validated on every Review Comment: Fixed! ########## crates/iceberg/src/writer/base_writer/position_delete_writer.rs: ########## @@ -0,0 +1,914 @@ +// 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. + +//! This module provides `PositionDeleteFileWriter`. +//! +//! A position delete file has two required columns: `file_path` (`string`, field id +//! [`RESERVED_FIELD_ID_DELETE_FILE_PATH`]) and `pos` (`long`, field id +//! [`RESERVED_FIELD_ID_DELETE_FILE_POS`]). The writer takes batches already shaped as +//! those two columns (see [`position_delete_schema`]) and sets +//! [`DataContentType::PositionDeletes`] on the output. It does not sort its input; see +//! [`PositionDeleteFileWriter::write`]. +//! +//! Position delete files are a v2 construct. v3 replaces them with deletion vectors and +//! forbids adding new position delete files, so callers must not route v3 writes here. +//! This base writer has no format-version gate by design; that gating belongs at the +//! transaction/commit layer. + +use std::sync::Arc; + +use arrow_array::RecordBatch; +use arrow_schema::{DataType, Field}; +use once_cell::sync::Lazy; +use parquet::arrow::PARQUET_FIELD_ID_META_KEY; + +use crate::metadata_columns::{ + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, delete_file_path_field, + delete_file_pos_field, +}; +use crate::spec::{DataContentType, DataFile, PartitionKey, Schema, SchemaRef}; +use crate::writer::file_writer::FileWriterBuilder; +use crate::writer::file_writer::location_generator::{FileNameGenerator, LocationGenerator}; +use crate::writer::file_writer::rolling_writer::{RollingFileWriter, RollingFileWriterBuilder}; +use crate::writer::{IcebergWriter, IcebergWriterBuilder}; +use crate::{Error, ErrorKind, Result}; + +/// The canonical Iceberg schema of a position delete file: the required `file_path` +/// (`string`) and `pos` (`long`) columns with their reserved field ids. +static POSITION_DELETE_SCHEMA: Lazy<SchemaRef> = Lazy::new(|| { + Arc::new( + Schema::builder() + .with_fields(vec![ + delete_file_path_field().clone(), + delete_file_pos_field().clone(), + ]) + .build() + .expect("position delete schema is statically valid"), + ) +}); + +/// [`POSITION_DELETE_SCHEMA`] converted to Arrow, keeping the reserved field ids in each +/// field's Parquet field-id metadata. Test-only for now: callers configure the writer with +/// the Iceberg [`position_delete_schema`], so no non-test code needs the Arrow form yet. +#[cfg(test)] +static POSITION_DELETE_ARROW_SCHEMA: Lazy<arrow_schema::SchemaRef> = Lazy::new(|| { + Arc::new( + crate::arrow::schema_to_arrow_schema(&POSITION_DELETE_SCHEMA) + .expect("position delete arrow schema is statically valid"), + ) +}); + +/// Returns the canonical Iceberg schema of a position delete file. +/// +/// Use this to build the [`ParquetWriterBuilder`](crate::writer::file_writer::ParquetWriterBuilder) +/// that backs a [`PositionDeleteFileWriter`], so the written file matches the +/// spec exactly. +pub fn position_delete_schema() -> SchemaRef { + POSITION_DELETE_SCHEMA.clone() +} + +/// Returns the canonical Arrow schema of a position delete file. +#[cfg(test)] +fn position_delete_arrow_schema() -> arrow_schema::SchemaRef { + POSITION_DELETE_ARROW_SCHEMA.clone() +} + +/// Reads a field's Iceberg field id from its Parquet field-id metadata. +fn field_id(field: &Field) -> Result<i32> { + field + .metadata() + .get(PARQUET_FIELD_ID_META_KEY) + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Position delete column `{}` is missing its Iceberg field id metadata.", + field.name() + ), + ) + })? + .parse::<i32>() + .map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!( + "Position delete column `{}` has an invalid field id: {e}", + field.name() + ), + ) + }) +} + +/// Validates that a batch is a position delete file: the `file_path` (`Utf8`) and +/// `pos` (`Int64`) columns, in order, with the two reserved field ids. Checking it +/// here gives a clear error before the batch reaches the Parquet writer. +fn validate_position_delete_batch(batch: &RecordBatch) -> Result<()> { + let fields = batch.schema_ref().fields(); + if fields.len() != 2 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "This writer supports only the two required position delete columns (`file_path`, `pos`); \ + batches with a different column count (e.g. including the optional `row` column) are not supported. Got {} columns.", + fields.len() + ), + )); + } + + let path = &fields[0]; + let path_id = field_id(path)?; + if path_id != RESERVED_FIELD_ID_DELETE_FILE_PATH { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The first position delete column must be `file_path` (field id {RESERVED_FIELD_ID_DELETE_FILE_PATH}), but got field id {path_id}." + ), + )); + } + // The canonical schema maps Iceberg `string` to `Utf8` and the file writer is + // configured with it, so a `LargeUtf8` column has to be cast to `Utf8` first. + if path.data_type() != &DataType::Utf8 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The position delete `file_path` column must be Utf8 (cast it first); got {:?}.", + path.data_type() + ), + )); + } + // Required column: a nullable field could write nulls under a required schema. + if path.is_nullable() { + return Err(Error::new( + ErrorKind::DataInvalid, + "The position delete `file_path` column must be required (non-nullable).", + )); + } + + let pos = &fields[1]; + let pos_id = field_id(pos)?; + if pos_id != RESERVED_FIELD_ID_DELETE_FILE_POS { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The second position delete column must be `pos` (field id {RESERVED_FIELD_ID_DELETE_FILE_POS}), but got field id {pos_id}." + ), + )); + } + if pos.data_type() != &DataType::Int64 { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "The position delete `pos` column must be Int64, but got {:?}.", + pos.data_type() + ), + )); + } + if pos.is_nullable() { + return Err(Error::new( + ErrorKind::DataInvalid, + "The position delete `pos` column must be required (non-nullable).", + )); + } + + Ok(()) +} + +/// Builder for [`PositionDeleteFileWriter`]. +#[derive(Debug)] +pub struct PositionDeleteFileWriterBuilder< + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +> { + inner: RollingFileWriterBuilder<B, L, F>, +} + +impl<B, L, F> PositionDeleteFileWriterBuilder<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + /// Create a new `PositionDeleteFileWriterBuilder` using a `RollingFileWriterBuilder`. + /// + /// The `RollingFileWriterBuilder` must be backed by a file writer configured + /// with the [`position_delete_schema`]; the per-batch validation in + /// [`PositionDeleteFileWriter::write`] guards against a mismatched batch, but + /// the caller is responsible for wiring the same schema into the file writer. + pub fn new(inner: RollingFileWriterBuilder<B, L, F>) -> Self { + Self { inner } + } +} + +#[async_trait::async_trait] +impl<B, L, F> IcebergWriterBuilder for PositionDeleteFileWriterBuilder<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + type R = PositionDeleteFileWriter<B, L, F>; + + async fn build(&self, partition_key: Option<PartitionKey>) -> Result<Self::R> { + Ok(PositionDeleteFileWriter { + inner: Some(self.inner.build()), + partition_key, + }) + } +} + +/// Writer used to write position delete files within one spec/partition. +#[derive(Debug)] +pub struct PositionDeleteFileWriter< + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +> { + inner: Option<RollingFileWriter<B, L, F>>, + partition_key: Option<PartitionKey>, +} + +#[async_trait::async_trait] +impl<B, L, F> IcebergWriter for PositionDeleteFileWriter<B, L, F> +where + B: FileWriterBuilder, + L: LocationGenerator, + F: FileNameGenerator, +{ + /// Writes a batch of `(file_path, pos)` records; the shape is validated on every Review Comment: Fixed! -- 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]
