diff --git a/crates/iceberg/src/arrow/reader/pipeline.rs b/crates/iceberg/src/arrow/reader/pipeline.rs index 779a6cd3a4..ccc7679ed0 100644 --- a/crates/iceberg/src/arrow/reader/pipeline.rs +++ b/crates/iceberg/src/arrow/reader/pipeline.rs @@ -20,12 +20,14 @@ //! predicates, row-group / row selection, and delete handling into a stream //! of transformed Arrow `RecordBatch`es. +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::AtomicU64; +use arrow_schema::{DataType, Field}; use futures::{StreamExt, TryStreamExt}; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; -use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder}; +use parquet::arrow::{PARQUET_FIELD_ID_META_KEY, ParquetRecordBatchStreamBuilder, RowNumber}; use parquet::encryption::decrypt::FileDecryptionProperties; use super::{ @@ -40,7 +42,8 @@ use crate::encryption::StandardKeyMetadata; use crate::error::Result; use crate::io::{FileIO, FileMetadata, FileRead}; use crate::metadata_columns::{ - RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_SPEC_ID, is_metadata_field, + RESERVED_COL_NAME_POS, RESERVED_FIELD_ID_FILE, RESERVED_FIELD_ID_POS, + RESERVED_FIELD_ID_SPEC_ID, is_metadata_field, }; use crate::scan::{ArrowRecordBatchStream, FileScanTask, FileScanTaskStream}; use crate::spec::Datum; @@ -212,6 +215,35 @@ impl FileScanTaskReader { arrow_metadata }; + let project_pos = task.project_field_ids().contains(&RESERVED_FIELD_ID_POS); + + let arrow_metadata = if project_pos { + let row_number_field = Arc::new( + Field::new(RESERVED_COL_NAME_POS, DataType::Int64, false) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_POS.to_string(), + )])) + .with_extension_type(RowNumber), + ); + + let options = ArrowReaderOptions::new() + .with_schema(Arc::clone(arrow_metadata.schema())) + .with_virtual_columns(vec![row_number_field])?; + + ArrowReaderMetadata::try_new(Arc::clone(arrow_metadata.metadata()), options).map_err( + |e| { + Error::new( + ErrorKind::Unexpected, + "Failed to create ArrowReaderMetadata with the 'row_number' virtual_column", + ) + .with_source(e) + }, + )? + } else { + arrow_metadata + }; + // Build the stream reader, reusing the already-opened file reader let mut record_batch_stream_builder = ParquetRecordBatchStreamBuilder::new_with_metadata(parquet_file_reader, arrow_metadata); @@ -274,6 +306,11 @@ impl FileScanTaskReader { record_batch_transformer_builder.with_partition(partition_spec, partition_data)?; } + if project_pos { + record_batch_transformer_builder = + record_batch_transformer_builder.with_virtual_field(RESERVED_FIELD_ID_POS); + } + let mut record_batch_transformer = record_batch_transformer_builder.build(); if let Some(batch_size) = self.batch_size { diff --git a/crates/iceberg/src/arrow/record_batch_transformer.rs b/crates/iceberg/src/arrow/record_batch_transformer.rs index b7b8d609a7..12f384ff52 100644 --- a/crates/iceberg/src/arrow/record_batch_transformer.rs +++ b/crates/iceberg/src/arrow/record_batch_transformer.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow_array::{ @@ -28,7 +28,7 @@ use arrow_schema::{ use parquet::arrow::PARQUET_FIELD_ID_META_KEY; use crate::arrow::value::{create_primitive_array_repeated, create_primitive_array_single_element}; -use crate::arrow::{datum_to_arrow_type_with_ree, schema_to_arrow_schema}; +use crate::arrow::{datum_to_arrow_type_with_ree, schema_to_arrow_schema, type_to_arrow_type}; use crate::metadata_columns::get_metadata_field; use crate::spec::{ Datum, Literal, PartitionSpec, PrimitiveLiteral, Schema as IcebergSchema, Struct, Transform, @@ -193,6 +193,7 @@ pub(crate) struct RecordBatchTransformerBuilder { snapshot_schema: Arc, projected_iceberg_field_ids: Vec, constant_fields: HashMap, + virtual_fields: HashSet, } impl RecordBatchTransformerBuilder { @@ -204,6 +205,7 @@ impl RecordBatchTransformerBuilder { snapshot_schema, projected_iceberg_field_ids: projected_iceberg_field_ids.to_vec(), constant_fields: HashMap::new(), + virtual_fields: HashSet::new(), } } @@ -240,11 +242,19 @@ impl RecordBatchTransformerBuilder { Ok(self) } + /// Set virtual fields such as '_pos'. Caller must have registered a corresponding + /// ArrowReaderOptions::with_virtual_columns call on the reader + pub(crate) fn with_virtual_field(mut self, field_id: i32) -> Self { + self.virtual_fields.insert(field_id); + self + } + pub(crate) fn build(self) -> RecordBatchTransformer { RecordBatchTransformer { snapshot_schema: self.snapshot_schema, projected_iceberg_field_ids: self.projected_iceberg_field_ids, constant_fields: self.constant_fields, + virtual_fields: self.virtual_fields, batch_transform: None, } } @@ -289,6 +299,11 @@ pub(crate) struct RecordBatchTransformer { // Datum holds both the Iceberg type and the value constant_fields: HashMap, + // Field IDs whose data is delivered by the source batch as an arrow-rs virtual column + // (e.g. _pos via RowNumber). These fields bypass the snapshot-schema lookup and the + // Iceberg projection rules (name mapping / initial-default / null) + virtual_fields: HashSet, + // BatchTransform gets lazily constructed based on the schema of // the first RecordBatch we receive from the file batch_transform: Option, @@ -330,6 +345,7 @@ impl RecordBatchTransformer { self.snapshot_schema.as_ref(), &self.projected_iceberg_field_ids, &self.constant_fields, + &self.virtual_fields, )?); self.process_record_batch(record_batch)? @@ -349,6 +365,7 @@ impl RecordBatchTransformer { snapshot_schema: &IcebergSchema, projected_iceberg_field_ids: &[i32], constant_fields: &HashMap, + virtual_fields: &HashSet, ) -> Result { let mapped_unprojected_arrow_schema = Arc::new(schema_to_arrow_schema(snapshot_schema)?); let field_id_to_mapped_schema_map = @@ -372,16 +389,26 @@ impl RecordBatchTransformer { PARQUET_FIELD_ID_META_KEY.to_string(), iceberg_field.id.to_string(), )])); - return Ok(Arc::new(arrow_field)); + Ok(Arc::new(arrow_field)) + } else if virtual_fields.contains(field_id) { + let virtual_field = get_metadata_field(*field_id)?; + let arrow_type = type_to_arrow_type(&virtual_field.field_type)?; + let arrow_field = + Field::new(&virtual_field.name, arrow_type, !virtual_field.required) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + virtual_field.id.to_string(), + )])); + Ok(Arc::new(arrow_field)) + } else { + // Regular fields and identity-partitioned constant fields both exist in the + // table schema, so use the mapped Arrow field as-is. + Ok(field_id_to_mapped_schema_map + .get(field_id) + .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? + .0 + .clone()) } - - // Regular fields and identity-partitioned constant fields both exist in the - // table schema, so use the mapped Arrow field as-is. - Ok(field_id_to_mapped_schema_map - .get(field_id) - .ok_or(Error::new(ErrorKind::Unexpected, "field not found"))? - .0 - .clone()) }) .collect(); @@ -397,24 +424,27 @@ impl RecordBatchTransformer { projected_iceberg_field_ids, field_id_to_mapped_schema_map, constant_fields, + virtual_fields, )?, target_schema, }), } } - /// Compares the source and target schemas - /// Determines if they have changed in any meaningful way: - /// * If they have different numbers of fields, then we need to modify - /// the incoming RecordBatch schema AND columns - /// * If they have the same number of fields, but some of them differ in - /// either data type or nullability, then we need to modify the - /// incoming RecordBatch schema AND columns - /// * If the schemas differ only in the column names, then we need - /// to modify the RecordBatch schema BUT we can keep the - /// original column data unmodified - /// * If the schemas are identical (or differ only in inconsequential - /// ways) then we can pass through the original RecordBatch unmodified + /// Compares the source and target schemas to decide how much work the + /// transform must do. The checks are applied in this order: + /// 1. If the schemas have different numbers of fields, they are + /// `Different`: we must rebuild both the schema AND the columns. + /// 2. If every field matches positionally in name, data type, and + /// nullability, they are `Equivalent`: the original RecordBatch can be + /// passed through unmodified. + /// 3. If they contain the same set of fields but in a different order + /// (detected via set-equality, since field names are unique), they are + /// `Different`: the columns must be rebuilt in the target order. + /// 4. Otherwise, the fields align positionally: if any differs in data + /// type or nullability the schemas are `Different` (rebuild schema AND + /// columns); if they differ only in names, it is `NameChangesOnly` and + /// we can relabel the schema BUT keep the original column data. fn compare_schemas( source_schema: &ArrowSchemaRef, target_schema: &ArrowSchemaRef, @@ -423,28 +453,37 @@ impl RecordBatchTransformer { return SchemaComparison::Different; } - let mut names_changed = false; + let field_eq = |a: &FieldRef, b: &FieldRef| { + a.name() == b.name() + && a.data_type() == b.data_type() + && a.is_nullable() == b.is_nullable() + }; + + let positional = || source_schema.fields().iter().zip(target_schema.fields()); - for (source_field, target_field) in source_schema + // Field-by-field identical: nothing to do. + if positional().all(|(s, t)| field_eq(s, t)) { + return SchemaComparison::Equivalent; + } + + // Same set of fields but reordered: rebuild in target (selection) order. + // Lengths are equal and field names are unique, so this is set-equality. + let reordered = target_schema .fields() .iter() - .zip(target_schema.fields().iter()) - { - if source_field.data_type() != target_field.data_type() - || source_field.is_nullable() != target_field.is_nullable() - { - return SchemaComparison::Different; - } - - if source_field.name() != target_field.name() { - names_changed = true; - } + .all(|t| source_schema.fields().iter().any(|s| field_eq(s, t))); + if reordered { + return SchemaComparison::Different; } - if names_changed { - SchemaComparison::NameChangesOnly + // Not a reordering: any type/nullability difference needs a full + // rebuild; otherwise only names differ and we can just relabel. + if positional() + .any(|(s, t)| s.data_type() != t.data_type() || s.is_nullable() != t.is_nullable()) + { + SchemaComparison::Different } else { - SchemaComparison::Equivalent + SchemaComparison::NameChangesOnly } } @@ -454,6 +493,7 @@ impl RecordBatchTransformer { projected_iceberg_field_ids: &[i32], field_id_to_mapped_schema_map: HashMap, constant_fields: &HashMap, + virtual_fields: &HashSet, ) -> Result> { let field_id_to_source_schema_map = Self::build_field_id_to_arrow_schema_map(source_schema)?; @@ -496,12 +536,22 @@ impl RecordBatchTransformer { } } + if virtual_fields.contains(field_id) { + let (_, source_index) = field_id_to_source_schema_map + .get(field_id) + .ok_or_else(|| Error::new( + ErrorKind::Unexpected, + format!("virtual column for field id {field_id} is not present in source batch — ensure ArrowReaderOptions::with_virtual_columns was applied"), + ))?; + return Ok(ColumnSource::PassThrough { source_index: *source_index }); + } + let (target_field, _) = field_id_to_mapped_schema_map .get(field_id) .ok_or(Error::new( ErrorKind::Unexpected, - "could not find field in schema", + format!("could not find field {field_id} in schema"), ))?; let target_type = target_field.data_type(); @@ -1733,4 +1783,152 @@ mod test { assert!(data_col.is_null(1)); assert!(data_col.is_null(2)); } + + #[test] + fn pos_column() { + use crate::metadata_columns::RESERVED_FIELD_ID_POS; + + let snapshot_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + // Simulate what arrow-rs's virtual-column reader produces: file columns + // followed by the _pos Int64 column carrying absolute file row indices. + let parquet_schema = Arc::new(ArrowSchema::new(vec![ + simple_field("id", DataType::Int32, false, "1"), + simple_field("name", DataType::Utf8, true, "2"), + simple_field( + "_pos", + DataType::Int64, + false, + &RESERVED_FIELD_ID_POS.to_string(), + ), + ])); + + let projected_field_ids = [1, 2, RESERVED_FIELD_ID_POS]; + + let mut transformer = + RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) + .with_virtual_field(RESERVED_FIELD_ID_POS) + .build(); + + let parquet_batch = RecordBatch::try_new(parquet_schema, vec![ + Arc::new(Int32Array::from(vec![100, 200, 300])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + Arc::new(Int64Array::from(vec![0, 1, 2])), + ]) + .unwrap(); + + let result = transformer.process_record_batch(parquet_batch).unwrap(); + + assert_eq!(result.num_columns(), 3); + assert_eq!(result.num_rows(), 3); + + // id column from file + let id_col = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(id_col.values(), &[100, 200, 300]); + + // name column from file + let name_col = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(name_col.value(0), "a"); + + // _pos column + let pos_col = result + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(pos_col.len(), 3); + assert_eq!(pos_col.values(), &[0, 1, 2]); + } + + /// field 1 and RESERVED_FIELD_ID_POS are of identical type + /// swapping their projection order should not lead to [`SchemaComparison::NameChangesOnly`] + #[test] + fn reorder_pos_column_with_columns_of_identical_type() { + use crate::metadata_columns::RESERVED_FIELD_ID_POS; + + let snapshot_schema = Arc::new( + Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "id_long", Type::Primitive(PrimitiveType::Long)) + .into(), + NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(), + ]) + .build() + .unwrap(), + ); + + let parquet_schema = Arc::new(ArrowSchema::new(vec![ + simple_field("id_long", DataType::Int64, false, "1"), + simple_field("name", DataType::Utf8, true, "2"), + simple_field( + "_pos", + DataType::Int64, + false, + &RESERVED_FIELD_ID_POS.to_string(), + ), + ])); + + let projected_field_ids = [RESERVED_FIELD_ID_POS, 2, 1]; + + let mut transformer = + RecordBatchTransformerBuilder::new(snapshot_schema, &projected_field_ids) + .with_virtual_field(RESERVED_FIELD_ID_POS) + .build(); + + let parquet_batch = RecordBatch::try_new(parquet_schema, vec![ + Arc::new(Int64Array::from(vec![100, 200, 300])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + Arc::new(Int64Array::from(vec![0, 1, 2])), + ]) + .unwrap(); + + let result = transformer.process_record_batch(parquet_batch).unwrap(); + + assert_eq!(result.num_columns(), 3); + assert_eq!(result.num_rows(), 3); + + let pos_col = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(pos_col.values(), &[0, 1, 2]); + + let name_col = result + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(name_col.value(0), "a"); + assert_eq!(name_col.value(1), "b"); + assert_eq!(name_col.value(2), "c"); + + let id_col = result + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(id_col.len(), 3); + assert_eq!(id_col.values(), &[100, 200, 300]); + } } diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index b377493729..3fd9dfacf2 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -642,7 +642,12 @@ pub mod tests { use crate::arrow::ArrowReaderBuilder; use crate::expr::{BoundPredicate, Reference}; use crate::io::{FileIO, OutputFile}; - use crate::metadata_columns::{RESERVED_COL_NAME_FILE, RESERVED_COL_NAME_SPEC_ID}; + use crate::metadata_columns::{ + RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, + RESERVED_COL_NAME_FILE, RESERVED_COL_NAME_POS, RESERVED_COL_NAME_SPEC_ID, + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, + RESERVED_FIELD_ID_POS, + }; use crate::scan::FileScanTask; use crate::spec::{ DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum, @@ -1416,6 +1421,200 @@ pub mod tests { .unwrap(); manifest_list_write.close().await.unwrap(); } + + /// Sets up a single data file `mrg.parquet` with three 100-row row groups + /// (column `x` = 1000..1300, so row position `p` carries `x = 1000 + p`) and + /// registers it in the current snapshot. When `delete_positions` is non-empty, + /// also writes a positional delete file targeting those file-absolute positions + /// and registers it in a delete manifest. + /// + /// Used to exercise the `_pos` metadata column through the real `TableScan` + /// planning path across row-group boundaries and (optionally) positional deletes. + pub async fn setup_multi_row_group_manifest(&mut self, delete_positions: &[i64]) { + let current_snapshot = self.table.metadata().current_snapshot().unwrap(); + let current_schema = current_snapshot.schema(self.table.metadata()).unwrap(); + let current_partition_spec = self.table.metadata().default_partition_spec(); + + // The table's spec 0 is identity on `x`, so give the data and delete files a + // fixed partition value. Filter tests deliberately filter on `y` (a + // non-partition column) so pruning is driven by Parquet row-group statistics + // rather than partition values. + let partition = Struct::from_iter([Some(Literal::long(1000))]); + + let (data_file_path, data_file_size) = self.write_multi_row_group_data_file(); + + let mut data_writer = ManifestWriterBuilder::new( + self.next_manifest_file(), + Some(current_snapshot.snapshot_id()), + current_schema.clone(), + current_partition_spec.as_ref().clone(), + ) + .build_v2_data(); + data_writer + .add_entry( + ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file( + DataFileBuilder::default() + .partition_spec_id(0) + .content(DataContentType::Data) + .file_path(data_file_path.clone()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(data_file_size) + .record_count(300) + .partition(partition.clone()) + .key_metadata(None) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + let data_manifest = data_writer.write_manifest_file().await.unwrap(); + + let mut manifests = vec![data_manifest]; + + if !delete_positions.is_empty() { + let (del_path, del_size) = + self.write_positional_delete_file(&data_file_path, delete_positions); + + let mut delete_writer = ManifestWriterBuilder::new( + self.next_manifest_file(), + Some(current_snapshot.snapshot_id()), + current_schema.clone(), + current_partition_spec.as_ref().clone(), + ) + .build_v2_deletes(); + delete_writer + .add_entry( + ManifestEntry::builder() + .status(ManifestStatus::Added) + .data_file( + DataFileBuilder::default() + .partition_spec_id(0) + .content(DataContentType::PositionDeletes) + .file_path(del_path) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(del_size) + .record_count(delete_positions.len() as u64) + .partition(partition.clone()) + .build() + .unwrap(), + ) + .build(), + ) + .unwrap(); + manifests.push(delete_writer.write_manifest_file().await.unwrap()); + } + + let manifest_list_writer = self + .table + .file_io() + .new_output(current_snapshot.manifest_list()) + .unwrap() + .writer() + .await + .unwrap(); + let mut manifest_list_write = ManifestListWriter::v2( + manifest_list_writer, + current_snapshot.snapshot_id(), + current_snapshot.parent_snapshot_id(), + current_snapshot.sequence_number(), + ); + manifest_list_write + .add_manifests(manifests.into_iter()) + .unwrap(); + manifest_list_write.close().await.unwrap(); + } + + /// Writes `mrg.parquet` with three 100-row row groups. Columns `x` (field + /// id `1`) and `y` (field id `2`) both run 1000..1300, so row position `p` + /// carries `x = y = 1000 + p`. Returns `(path, file_size_in_bytes)`. + fn write_multi_row_group_data_file(&self) -> (String, u64) { + fs::create_dir_all(&self.table_location).unwrap(); + + let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new("x", arrow_schema::DataType::Int64, false).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]), + ), + arrow_schema::Field::new("y", arrow_schema::DataType::Int64, false).with_metadata( + HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())]), + ), + ])); + + let path = format!("{}/mrg.parquet", &self.table_location); + let max_row_group_row_count = 100; + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .set_max_row_group_row_count(Some(max_row_group_row_count)) + .build(); + + let file = File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, arrow_schema.clone(), Some(props)).unwrap(); + for group in 0..3i64 { + let base = 1000 + group * max_row_group_row_count as i64; + let col = Arc::new(Int64Array::from_iter_values( + base..base + max_row_group_row_count as i64, + )) as ArrayRef; + let batch = + RecordBatch::try_new(arrow_schema.clone(), vec![col.clone(), col]).unwrap(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + let size = fs::metadata(&path).unwrap().len(); + (path, size) + } + + /// Writes a positional delete file targeting `positions` in `data_path`. + /// Returns `(path, file_size_in_bytes)`. + fn write_positional_delete_file( + &self, + data_path: &str, + positions: &[i64], + ) -> (String, u64) { + let del_schema = Arc::new(arrow_schema::Schema::new(vec![ + arrow_schema::Field::new( + RESERVED_COL_NAME_DELETE_FILE_PATH, + arrow_schema::DataType::Utf8, + false, + ) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), // 2147483546 + )])), + arrow_schema::Field::new( + RESERVED_COL_NAME_DELETE_FILE_POS, + arrow_schema::DataType::Int64, + false, + ) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), // 2147483545 + )])), + ])); + + let batch = RecordBatch::try_new(del_schema.clone(), vec![ + Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + data_path.to_string(), + positions.len(), + ))) as ArrayRef, + Arc::new(Int64Array::from_iter_values(positions.iter().copied())) as ArrayRef, + ]) + .unwrap(); + + let path = format!("{}/pos-del.parquet", &self.table_location); + let props = WriterProperties::builder() + .set_compression(Compression::SNAPPY) + .build(); + let file = File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new(file, del_schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let size = fs::metadata(&path).unwrap().len(); + (path, size) + } } #[tokio::test] @@ -2680,4 +2879,350 @@ pub mod tests { let spec_id = int_values.value(0); assert_eq!(spec_id, 2, "_spec_id should be 2, got: {spec_id}"); } + + #[tokio::test] + async fn test_select_with_pos_and_file_columns() { + use arrow_array::cast::AsArray; + + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + // Select regular columns plus the _pos column + let table_scan = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_POS, RESERVED_COL_NAME_FILE]) + .with_row_selection_enabled(true) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 2); + + // Examine batches are 1.paruqet and 3.parquet, 2.parquet is deleted. + for batch in batches.iter() { + // Verify we have 3 columns: x, _pos and _file + assert_eq!(batch.num_columns(), 3); + + // Verify the x column exists and has correct data + let x_col = batch.column_by_name("x").unwrap(); + let x_arr = x_col.as_primitive::(); + assert_eq!(x_arr.value(0), 1); + + // The _pos column exists and verify it is Int64Array with the expected values + let pos_col = batch.column(1); + let pos_array: &Int64Array = pos_col + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array"); + assert_eq!(*pos_array, Int64Array::from_iter_values(0i64..1024)); + + // Verify the _file column exists + let file_col = batch.column_by_name(RESERVED_COL_NAME_FILE); + assert!( + file_col.is_some(), + "_file column should be present in the batch" + ); + } + } + + #[tokio::test] + async fn test_pos_column_at_start_with_filters() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + // y is in [4, 5) + let predicate = Reference::new("y") + .greater_than(Datum::long(4i64)) + .and(Reference::new("y").less_than_or_equal_to(Datum::long(5i64))); + // Select _pos at the start + let table_scan = fixture + .table + .scan() + .select([RESERVED_COL_NAME_POS, "x", "y"]) + .with_filter(predicate) + .with_row_selection_enabled(true) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 2); + + // Examine batches are 1.paruqet and 3.parquet, 2.parquet is deleted. + for batch in batches.iter() { + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.num_rows(), 12); + + // Verify _pos is at position 0 + let schema = batch.schema(); + assert_eq!(schema.field(0).name(), RESERVED_COL_NAME_POS); + assert_eq!(schema.field(1).name(), "x"); + assert_eq!(schema.field(2).name(), "y"); + + let pos_col = batch.column(0); + let pos_array: &Int64Array = pos_col + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array"); + assert_eq!(*pos_array, Int64Array::from_iter_values(1012i64..1024)); + } + } + + #[tokio::test] + async fn test_repeated_pos_column_with_filter() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + + // a NOT STARTSWITH "Apa" + let predicate = Reference::new("a").not_starts_with(Datum::string("Apa")); + // Select '_pos' columns twice + let table_scan = fixture + .table + .scan() + .select([RESERVED_COL_NAME_POS, "a", RESERVED_COL_NAME_POS, "x"]) + .with_row_selection_enabled(true) + .with_filter(predicate) + .build() + .unwrap(); + + let batch_stream = table_scan.to_arrow().await.unwrap(); + + let batches: Vec<_> = batch_stream.try_collect().await.unwrap(); + assert_eq!(batches.len(), 2); + + // Examine batches are 1.paruqet and 3.parquet, 2.parquet is deleted. + for batch in batches.iter() { + assert_eq!(batch.num_rows(), 512); + + // fetch the 1st _pos column by name and verify it is Int64Array with the expected values + let pos_col = batch + .column_by_name("_pos") + .expect("_pos column should be present in the batch"); + let pos_array: &Int64Array = pos_col + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array"); + assert_eq!(*pos_array, Int64Array::from_iter_values(512i64..1024)); + + // fetch the 2nd _pos column by index and verify it is Int64Array with the expected values + let pos_col = batch.column(2); + let pos_array: &Int64Array = pos_col + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array"); + assert_eq!(*pos_array, Int64Array::from_iter_values(512i64..1024)); + } + } + + /// End-to-end through `TableScan`: a data file with three row groups planned + /// as a single whole-file `FileScanTask` must yield contiguous, file-absolute + /// `_pos` values (0..300) across the row-group boundaries. + #[tokio::test] + async fn test_pos_across_row_groups_via_table_scan() { + let mut fixture = TableTestFixture::new(); + fixture.setup_multi_row_group_manifest(&[]).await; + + // Planning must produce exactly one whole-file task with _pos projected and + // no delete files, confirming TableScan does not sub-split the file. + let tasks: Vec<_> = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_POS]) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1, "expected a single FileScanTask"); + let task = &tasks[0]; + assert!( + task.project_field_ids.contains(&RESERVED_FIELD_ID_POS), + "_pos field id must be projected into the FileScanTask" + ); + assert_eq!(task.start, 0, "TableScan should plan whole-file tasks"); + assert_eq!(task.length, task.file_size_in_bytes); + assert!(task.deletes.is_empty()); + + // Reading that task yields absolute _pos 0..300 in order. + let batches: Vec<_> = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_POS]) + .build() + .unwrap() + .to_arrow() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let pos: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name(RESERVED_COL_NAME_POS) + .expect("_pos column should be present") + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array") + .values() + .to_vec() + }) + .collect(); + assert_eq!(pos, (0..300).collect::>()); + + // Sanity: x == 1000 + _pos, proving _pos aligns with the actual rows read. + let x: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("x") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(x, (1000..1300).collect::>()); + } + + /// A positional delete file registered in the manifest must be attached to the planned + /// `FileScanTask` and applied on read, while surviving `_pos` values stay file-absolute. + #[tokio::test] + async fn test_pos_with_positional_deletes_via_table_scan() { + let mut fixture = TableTestFixture::new(); + // Delete file-absolute positions 150 (middle row group) and 299 (last row). + fixture.setup_multi_row_group_manifest(&[150, 299]).await; + + // Planning must attach the positional delete file to the task. + let tasks: Vec<_> = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_POS]) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!( + tasks[0].deletes.len(), + 1, + "positional delete file should be planned into the task" + ); + assert_eq!( + tasks[0].deletes[0].file_type, + DataContentType::PositionDeletes + ); + + // Reading applies the deletes; _pos must skip 150 and 299 and stay absolute. + let batches: Vec<_> = fixture + .table + .scan() + .select(["x", RESERVED_COL_NAME_POS]) + .build() + .unwrap() + .to_arrow() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let pos: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name(RESERVED_COL_NAME_POS) + .expect("_pos column should be present") + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array") + .values() + .to_vec() + }) + .collect(); + + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 298, + "two rows should be removed by positional deletes" + ); + assert!(!pos.contains(&150) && !pos.contains(&299), "got {pos:?}"); + let expected: Vec = (0..150).chain(151..299).collect(); + assert_eq!(pos, expected); + } + + /// A filter that only matches the middle row group (`y` in [1100, 1200)) and + /// prunes the other two row groups by statistics, so only the middle row group + /// is read. `_pos` must report the file-absolute positions 100..200 for those + /// rows, not values reset to 0..100. + /// + /// `y` is a non-partition column, so pruning here is driven purely by Parquet + /// row-group statistics (the TableTestFixture's partition column is `x`). + #[tokio::test] + async fn test_pos_reads_only_middle_row_group_via_filter() { + let mut fixture = TableTestFixture::new(); + fixture.setup_multi_row_group_manifest(&[]).await; + + // Middle row group holds y = 1100..1200 at file positions 100..200. + let predicate = Reference::new("y") + .greater_than_or_equal_to(Datum::long(1100)) + .and(Reference::new("y").less_than(Datum::long(1200))); + + let batches: Vec<_> = fixture + .table + .scan() + .select(["y", RESERVED_COL_NAME_POS]) + .with_filter(predicate) + .with_row_group_filtering_enabled(true) + .build() + .unwrap() + .to_arrow() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 100, "only the middle row group should be read"); + + let pos: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name(RESERVED_COL_NAME_POS) + .expect("_pos column should be present") + .as_any() + .downcast_ref::() + .expect("_pos column should be a Int64Array") + .values() + .to_vec() + }) + .collect(); + assert_eq!( + pos, + (100..200).collect::>(), + "_pos must be file-absolute for the middle row group" + ); + + // Cross-check: y == 1000 + _pos for every surviving row. + let y: Vec = batches + .iter() + .flat_map(|b| { + b.column_by_name("y") + .unwrap() + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(y, (1100..1200).collect::>()); + } }