From 8e32c6897d5474287e0ad666e4928d27db300e09 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 9 Oct 2026 17:20:03 -0700 Subject: [PATCH 1/7] feat(arrow): convert manifest entries to an Arrow RecordBatch --- crates/iceberg/public-api.txt | 1 + crates/iceberg/src/arrow/manifest.rs | 1162 +++++++++++++++++++++ crates/iceberg/src/arrow/mod.rs | 2 + crates/iceberg/src/spec/manifest/entry.rs | 18 + 4 files changed, 1183 insertions(+) create mode 100644 crates/iceberg/src/arrow/manifest.rs diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 8a34bbf7b5..90e899cb13 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -130,6 +130,7 @@ pub fn iceberg::arrow::arrow_schema_to_schema_auto_assign_ids(schema: &arrow_sch pub fn iceberg::arrow::arrow_struct_to_literal(struct_array: &arrow_array::array::ArrayRef, ty: &iceberg::spec::StructType) -> iceberg::Result>> pub fn iceberg::arrow::arrow_type_to_type(ty: &arrow_schema::datatype::DataType) -> iceberg::Result pub fn iceberg::arrow::datum_to_arrow_type_with_ree(datum: &iceberg::spec::Datum) -> arrow_schema::datatype::DataType +pub fn iceberg::arrow::manifest_to_record_batch(manifest: &iceberg::spec::Manifest) -> iceberg::Result pub fn iceberg::arrow::schema_to_arrow_schema(schema: &iceberg::spec::Schema) -> iceberg::Result pub fn iceberg::arrow::strip_metadata_from_schema(schema: &arrow_schema::schema::Schema) -> iceberg::Result pub fn iceberg::arrow::type_to_arrow_type(ty: &iceberg::spec::Type) -> iceberg::Result diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs new file mode 100644 index 0000000000..e12d3a150b --- /dev/null +++ b/crates/iceberg/src/arrow/manifest.rs @@ -0,0 +1,1162 @@ +// 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. + +//! Conversion of manifest entries to Arrow. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_array::builder::{ + ArrayBuilder, Int32Builder, Int64Builder, LargeBinaryBuilder, ListBuilder, MapBuilder, + MapFieldNames, StringBuilder, +}; +use arrow_array::{ + ArrayRef, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, Float32Array, + Float64Array, Int32Array, Int64Array, LargeBinaryArray, NullArray, RecordBatch, StringArray, + StructArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, +}; +use arrow_schema::{DataType, Field, FieldRef, Fields}; + +use crate::arrow::{UTC_TIME_ZONE, schema_to_arrow_schema}; +use crate::error::invalid_data; +use crate::spec::{ + DataFileFormat, Datum, Literal, Manifest, ManifestEntryRef, NestedField, PrimitiveLiteral, + PrimitiveType, Schema, StructType, manifest_entry_fields, +}; +use crate::{Error, ErrorKind, Result}; + +/// Converts the entries of `manifest` into a [`RecordBatch`] with one row per +/// entry. +/// +/// The schema is the spec's `manifest_entry` struct for the latest format +/// version, with field IDs in the Arrow field metadata, so manifests of every +/// format version convert to the same schema apart from `data_file.partition`, +/// which follows the manifest's partition spec. Lower and upper bounds use the +/// spec's binary single-value serialization. The metric maps list their +/// entries in no particular order, and a metric map with no entries is empty +/// rather than null. +/// +/// # Example +/// +/// ``` +/// use std::sync::Arc; +/// +/// use iceberg::arrow::manifest_to_record_batch; +/// use iceberg::spec::{ +/// DataContentType, DataFileBuilder, DataFileFormat, FormatVersion, Manifest, +/// ManifestContentType, ManifestEntry, ManifestMetadata, ManifestStatus, NestedField, +/// PartitionSpec, PrimitiveType, Schema, Type, +/// }; +/// +/// let schema = Arc::new( +/// Schema::builder() +/// .with_fields(vec![ +/// NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), +/// ]) +/// .build()?, +/// ); +/// let data_file = DataFileBuilder::default() +/// .content(DataContentType::Data) +/// .file_path("s3://bucket/data/0.parquet".to_string()) +/// .file_format(DataFileFormat::Parquet) +/// .record_count(100) +/// .file_size_in_bytes(4096) +/// .build()?; +/// let manifest = Manifest::new( +/// ManifestMetadata { +/// schema: schema.clone(), +/// schema_id: 0, +/// partition_spec: PartitionSpec::builder(schema).build()?, +/// format_version: FormatVersion::V2, +/// content: ManifestContentType::Data, +/// }, +/// vec![ +/// ManifestEntry::builder() +/// .status(ManifestStatus::Added) +/// .data_file(data_file) +/// .build(), +/// ], +/// ); +/// +/// let batch = manifest_to_record_batch(&manifest)?; +/// assert_eq!(batch.num_rows(), 1); +/// # Ok::<(), Box>(()) +/// ``` +pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { + let metadata = manifest.metadata(); + let partition_type = metadata + .partition_spec() + .partition_type(metadata.schema())?; + let schema = Arc::new(schema_to_arrow_schema( + &Schema::builder() + .with_fields(manifest_entry_fields(&partition_type)) + .build()?, + )?); + let data_file_fields = struct_fields(schema.field_with_name("data_file")?)?; + let child = |name: &str| -> Result<&FieldRef> { + data_file_fields + .find(name) + .map(|(_, field)| field) + .ok_or_else(|| Error::new(ErrorKind::Unexpected, format!("no data_file field {name}"))) + }; + + let entries = manifest.entries(); + let len = entries.len(); + let mut status = Int32Builder::with_capacity(len); + let mut snapshot_id = Int64Builder::with_capacity(len); + let mut sequence_number = Int64Builder::with_capacity(len); + let mut file_sequence_number = Int64Builder::with_capacity(len); + let mut content = Int32Builder::with_capacity(len); + let mut file_path = StringBuilder::new(); + let mut file_format = StringBuilder::new(); + let mut record_count = Int64Builder::with_capacity(len); + let mut file_size_in_bytes = Int64Builder::with_capacity(len); + let mut column_sizes = map_builder(child("column_sizes")?, Int64Builder::new())?; + let mut value_counts = map_builder(child("value_counts")?, Int64Builder::new())?; + let mut null_value_counts = map_builder(child("null_value_counts")?, Int64Builder::new())?; + let mut nan_value_counts = map_builder(child("nan_value_counts")?, Int64Builder::new())?; + let mut lower_bounds = map_builder(child("lower_bounds")?, LargeBinaryBuilder::new())?; + let mut upper_bounds = map_builder(child("upper_bounds")?, LargeBinaryBuilder::new())?; + let mut key_metadata = LargeBinaryBuilder::new(); + let mut split_offsets = + ListBuilder::new(Int64Builder::new()).with_field(list_element(child("split_offsets")?)?); + let mut equality_ids = + ListBuilder::new(Int32Builder::new()).with_field(list_element(child("equality_ids")?)?); + let mut sort_order_id = Int32Builder::with_capacity(len); + let mut first_row_id = Int64Builder::with_capacity(len); + let mut referenced_data_file = StringBuilder::new(); + let mut content_offset = Int64Builder::with_capacity(len); + let mut content_size_in_bytes = Int64Builder::with_capacity(len); + + for entry in entries { + status.append_value(entry.status as i32); + snapshot_id.append_option(entry.snapshot_id); + sequence_number.append_option(entry.sequence_number); + file_sequence_number.append_option(entry.file_sequence_number); + + let data_file = &entry.data_file; + content.append_value(data_file.content as i32); + file_path.append_value(&data_file.file_path); + // Writers store the format name in upper case. + file_format.append_value(match data_file.file_format { + DataFileFormat::Avro => "AVRO", + DataFileFormat::Orc => "ORC", + DataFileFormat::Parquet => "PARQUET", + DataFileFormat::Puffin => "PUFFIN", + }); + record_count.append_value(data_file.record_count.try_into()?); + file_size_in_bytes.append_value(data_file.file_size_in_bytes.try_into()?); + append_counts(&mut column_sizes, &data_file.column_sizes)?; + append_counts(&mut value_counts, &data_file.value_counts)?; + append_counts(&mut null_value_counts, &data_file.null_value_counts)?; + append_counts(&mut nan_value_counts, &data_file.nan_value_counts)?; + append_bounds(&mut lower_bounds, &data_file.lower_bounds)?; + append_bounds(&mut upper_bounds, &data_file.upper_bounds)?; + key_metadata.append_option(data_file.key_metadata.as_deref()); + match &data_file.split_offsets { + Some(offsets) => { + split_offsets.values().append_slice(offsets); + split_offsets.append(true); + } + None => split_offsets.append(false), + } + match &data_file.equality_ids { + Some(ids) => { + equality_ids.values().append_slice(ids); + equality_ids.append(true); + } + None => equality_ids.append(false), + } + sort_order_id.append_option(data_file.sort_order_id); + first_row_id.append_option(data_file.first_row_id); + referenced_data_file.append_option(data_file.referenced_data_file.as_deref()); + content_offset.append_option(data_file.content_offset); + content_size_in_bytes.append_option(data_file.content_size_in_bytes); + } + + let partition = partition_column(entries, &partition_type, child("partition")?)?; + let data_file = StructArray::try_new( + data_file_fields.clone(), + vec![ + Arc::new(content.finish()), + Arc::new(file_path.finish()), + Arc::new(file_format.finish()), + partition, + Arc::new(record_count.finish()), + Arc::new(file_size_in_bytes.finish()), + Arc::new(column_sizes.finish()), + Arc::new(value_counts.finish()), + Arc::new(null_value_counts.finish()), + Arc::new(nan_value_counts.finish()), + Arc::new(lower_bounds.finish()), + Arc::new(upper_bounds.finish()), + Arc::new(key_metadata.finish()), + Arc::new(split_offsets.finish()), + Arc::new(equality_ids.finish()), + Arc::new(sort_order_id.finish()), + Arc::new(first_row_id.finish()), + Arc::new(referenced_data_file.finish()), + Arc::new(content_offset.finish()), + Arc::new(content_size_in_bytes.finish()), + ], + None, + )?; + + Ok(RecordBatch::try_new(schema, vec![ + Arc::new(status.finish()), + Arc::new(snapshot_id.finish()), + Arc::new(sequence_number.finish()), + Arc::new(file_sequence_number.finish()), + Arc::new(data_file), + ])?) +} + +fn struct_fields(field: &Field) -> Result<&Fields> { + match field.data_type() { + DataType::Struct(fields) => Ok(fields), + data_type => Err(Error::new( + ErrorKind::Unexpected, + format!("field {} is {data_type}, not a struct", field.name()), + )), + } +} + +fn list_element(field: &Field) -> Result { + match field.data_type() { + DataType::List(element) => Ok(element.clone()), + data_type => Err(Error::new( + ErrorKind::Unexpected, + format!("field {} is {data_type}, not a list", field.name()), + )), + } +} + +/// Creates a builder for the `map` column `field` that produces the +/// key and value fields of its schema, including their field IDs. +fn map_builder(field: &Field, values: V) -> Result> { + let DataType::Map(entries, _) = field.data_type() else { + return Err(Error::new( + ErrorKind::Unexpected, + format!("field {} is {}, not a map", field.name(), field.data_type()), + )); + }; + let [key, value] = &struct_fields(entries)?[..] else { + return Err(Error::new( + ErrorKind::Unexpected, + format!("map {} doesn't have two entry fields", field.name()), + )); + }; + let names = MapFieldNames { + entry: entries.name().clone(), + key: key.name().clone(), + value: value.name().clone(), + }; + Ok(MapBuilder::new(Some(names), Int32Builder::new(), values) + .with_keys_field(key.clone()) + .with_values_field(value.clone())) +} + +fn append_counts( + builder: &mut MapBuilder, + counts: &HashMap, +) -> Result<()> { + for (&field_id, &count) in counts { + builder.keys().append_value(field_id); + builder.values().append_value(count.try_into()?); + } + Ok(builder.append(true)?) +} + +fn append_bounds( + builder: &mut MapBuilder, + bounds: &HashMap, +) -> Result<()> { + for (&field_id, bound) in bounds { + builder.keys().append_value(field_id); + builder.values().append_value(bound.to_bytes()?); + } + Ok(builder.append(true)?) +} + +/// Builds the `data_file.partition` column, one child array per partition field. +fn partition_column( + entries: &[ManifestEntryRef], + partition_type: &StructType, + field: &Field, +) -> Result { + if let Some(entry) = entries + .iter() + .find(|entry| entry.data_file.partition.fields().len() != partition_type.fields().len()) + { + return Err(invalid_data!( + "partition of {} has {} values, but the partition spec has {} fields", + entry.data_file.file_path, + entry.data_file.partition.fields().len(), + partition_type.fields().len() + )); + } + + let fields = struct_fields(field)?; + if fields.is_empty() { + return Ok(Arc::new(StructArray::new_empty_fields(entries.len(), None))); + } + let columns = partition_type + .fields() + .iter() + .zip(fields) + .enumerate() + .map(|(index, (partition_field, arrow_field))| { + let values: Vec<_> = entries + .iter() + .map(|entry| match &entry.data_file.partition[index] { + None => Ok(None), + Some(Literal::Primitive(value)) => Ok(Some(value)), + Some(value) => Err(invalid_data!( + "partition field {} has non-primitive value {value:?}", + partition_field.name + )), + }) + .collect::>()?; + partition_array(partition_field, arrow_field.data_type(), &values) + }) + .collect::>()?; + Ok(Arc::new(StructArray::try_new( + fields.clone(), + columns, + None, + )?)) +} + +/// Builds the array of one partition field from its values in each entry. +fn partition_array( + field: &NestedField, + data_type: &DataType, + values: &[Option<&PrimitiveLiteral>], +) -> Result { + let primitive_type = field.field_type.as_primitive_type().ok_or_else(|| { + invalid_data!( + "partition field {} has non-primitive type {}", + field.name, + field.field_type + ) + })?; + // Collects the values, returning an error for a value of another type. + fn collect<'a, T, A: FromIterator>>( + field: &NestedField, + values: &[Option<&'a PrimitiveLiteral>], + value: impl Fn(&'a PrimitiveLiteral) -> Option, + ) -> Result { + values + .iter() + .map(|literal| { + literal + .map(|literal| { + value(literal).ok_or_else(|| { + invalid_data!( + "partition field {} of type {} can't hold {literal:?}", + field.name, + field.field_type + ) + }) + }) + .transpose() + }) + .collect() + } + fn int(literal: &PrimitiveLiteral) -> Option { + match literal { + PrimitiveLiteral::Int(v) => Some(*v), + _ => None, + } + } + fn long(literal: &PrimitiveLiteral) -> Option { + match literal { + PrimitiveLiteral::Long(v) => Some(*v), + _ => None, + } + } + fn binary(literal: &PrimitiveLiteral) -> Option<&[u8]> { + match literal { + PrimitiveLiteral::Binary(v) => Some(v), + _ => None, + } + } + let array: ArrayRef = + match primitive_type { + PrimitiveType::Boolean => Arc::new(collect::<_, BooleanArray>( + field, + values, + |literal| match literal { + PrimitiveLiteral::Boolean(v) => Some(*v), + _ => None, + }, + )?), + PrimitiveType::Int => Arc::new(collect::<_, Int32Array>(field, values, int)?), + PrimitiveType::Date => Arc::new(collect::<_, Date32Array>(field, values, int)?), + PrimitiveType::Long => Arc::new(collect::<_, Int64Array>(field, values, long)?), + PrimitiveType::Time => { + Arc::new(collect::<_, Time64MicrosecondArray>(field, values, long)?) + } + PrimitiveType::Timestamp => Arc::new(collect::<_, TimestampMicrosecondArray>( + field, values, long, + )?), + PrimitiveType::Timestamptz => Arc::new( + collect::<_, TimestampMicrosecondArray>(field, values, long)? + .with_timezone(UTC_TIME_ZONE), + ), + PrimitiveType::TimestampNs => { + Arc::new(collect::<_, TimestampNanosecondArray>(field, values, long)?) + } + PrimitiveType::TimestamptzNs => Arc::new( + collect::<_, TimestampNanosecondArray>(field, values, long)? + .with_timezone(UTC_TIME_ZONE), + ), + PrimitiveType::Float => Arc::new(collect::<_, Float32Array>( + field, + values, + |literal| match literal { + PrimitiveLiteral::Float(v) => Some(v.0), + _ => None, + }, + )?), + PrimitiveType::Double => Arc::new(collect::<_, Float64Array>( + field, + values, + |literal| match literal { + PrimitiveLiteral::Double(v) => Some(v.0), + _ => None, + }, + )?), + PrimitiveType::Decimal { .. } => { + let DataType::Decimal128(precision, scale) = *data_type else { + return Err(Error::new( + ErrorKind::Unexpected, + format!("decimal partition field {} is {data_type}", field.name), + )); + }; + Arc::new( + collect::<_, Decimal128Array>(field, values, |literal| match literal { + PrimitiveLiteral::Int128(v) => Some(*v), + _ => None, + })? + .with_precision_and_scale(precision, scale)?, + ) + } + PrimitiveType::String => Arc::new(collect::<_, StringArray>( + field, + values, + |literal| match literal { + PrimitiveLiteral::String(v) => Some(v.as_str()), + _ => None, + }, + )?), + PrimitiveType::Uuid => { + let uuids: Vec<_> = collect(field, values, |literal| match literal { + PrimitiveLiteral::UInt128(v) => Some(v.to_be_bytes()), + _ => None, + })?; + Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size( + uuids.into_iter(), + 16, + )?) + } + PrimitiveType::Fixed(_) => { + let DataType::FixedSizeBinary(width) = *data_type else { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("fixed partition field {} is {data_type}", field.name), + )); + }; + let bytes: Vec<_> = collect(field, values, binary)?; + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size(bytes.into_iter(), width) + .map_err(|err| { + invalid_data!( + "partition field {} has a value of the wrong width", + field.name + ) + .with_source(err) + })?, + ) + } + PrimitiveType::Binary => { + Arc::new(collect::<_, LargeBinaryArray>(field, values, binary)?) + } + PrimitiveType::Unknown => { + // Only nulls fit, so this returns an error for any value. + collect::<(), Vec<_>>(field, values, |_| None)?; + Arc::new(NullArray::new(values.len())) + } + }; + Ok(array) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + + use arrow_array::cast::AsArray; + use arrow_array::types::{Int32Type, Int64Type}; + use arrow_array::{ + Array, ArrayRef, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, + Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray, NullArray, + StringArray, StructArray, Time64MicrosecondArray, TimestampMicrosecondArray, + TimestampNanosecondArray, + }; + use uuid::Uuid; + + use super::*; + use crate::arrow::arrow_schema_to_schema; + use crate::spec::{ + DataContentType, DataFile, DataFileBuilder, DataFileFormat, Datum, FormatVersion, ListType, + Literal, ManifestContentType, ManifestEntry, ManifestMetadata, ManifestStatus, MapType, + NestedField, PartitionSpec, PrimitiveType, Schema, Struct, StructType, Transform, Type, + }; + + const UUID: &str = "f79c3e09-677c-4bbd-a479-3f349cb785e7"; + + fn table_schema() -> Schema { + let columns = [ + (1, "b", PrimitiveType::Boolean), + (2, "i", PrimitiveType::Int), + (3, "l", PrimitiveType::Long), + (4, "f", PrimitiveType::Float), + (5, "d", PrimitiveType::Double), + (6, "dec", PrimitiveType::Decimal { + precision: 38, + scale: 10, + }), + (7, "date", PrimitiveType::Date), + (8, "time", PrimitiveType::Time), + (9, "ts", PrimitiveType::Timestamp), + (10, "tstz", PrimitiveType::Timestamptz), + (11, "ts_ns", PrimitiveType::TimestampNs), + (12, "tstz_ns", PrimitiveType::TimestamptzNs), + (13, "s", PrimitiveType::String), + (14, "u", PrimitiveType::Uuid), + (15, "fx", PrimitiveType::Fixed(3)), + (16, "bin", PrimitiveType::Binary), + (17, "unk", PrimitiveType::Unknown), + ]; + Schema::builder() + .with_fields( + columns.into_iter().map(|(id, name, ty)| { + NestedField::optional(id, name, Type::Primitive(ty)).into() + }), + ) + .build() + .unwrap() + } + + fn partition_spec(fields: &[(&str, Transform)]) -> PartitionSpec { + fields + .iter() + .fold( + PartitionSpec::builder(table_schema()).with_spec_id(0), + |builder, (source, transform)| { + builder + .add_partition_field(*source, format!("{source}_{transform}"), *transform) + .unwrap() + }, + ) + .build() + .unwrap() + } + + fn manifest(partition_spec: PartitionSpec, entries: Vec) -> Manifest { + Manifest::new( + ManifestMetadata { + schema: Arc::new(table_schema()), + schema_id: 0, + partition_spec, + format_version: FormatVersion::V2, + content: ManifestContentType::Data, + }, + entries, + ) + } + + fn data_file(content: DataContentType, path: &str) -> DataFileBuilder { + let mut builder = DataFileBuilder::default(); + builder + .content(content) + .file_path(path.to_string()) + .file_format(DataFileFormat::Parquet) + .record_count(1) + .file_size_in_bytes(10); + builder + } + + fn entry(status: ManifestStatus, data_file: DataFile) -> ManifestEntry { + ManifestEntry { + status, + snapshot_id: None, + sequence_number: None, + file_sequence_number: None, + data_file, + } + } + + fn data_file_column(batch: &RecordBatch, name: &str) -> ArrayRef { + batch + .column_by_name("data_file") + .unwrap() + .as_struct() + .column_by_name(name) + .unwrap() + .clone() + } + + /// Entries of the map in `row` of a `map` column, sorted by key. + fn long_map(column: &ArrayRef, row: usize) -> Vec<(i32, i64)> { + let entries = column.as_map().value(row); + let keys = entries.column(0).as_primitive::(); + let values = entries.column(1).as_primitive::(); + let mut pairs: Vec<_> = keys + .values() + .iter() + .copied() + .zip(values.values().iter().copied()) + .collect(); + pairs.sort_unstable(); + pairs + } + + /// Entries of the map in `row` of a `map` column, sorted by key. + fn binary_map(column: &ArrayRef, row: usize) -> Vec<(i32, Vec)> { + let entries = column.as_map().value(row); + let keys = entries.column(0).as_primitive::(); + let values = entries.column(1).as_binary::(); + let mut pairs: Vec<_> = keys + .values() + .iter() + .copied() + .zip(values.iter().map(|v| v.unwrap().to_vec())) + .collect(); + pairs.sort_unstable(); + pairs + } + + fn int_map_type(key_id: i32, value_id: i32, value: PrimitiveType) -> Type { + Type::Map(MapType::new( + NestedField::map_key_element(key_id, Type::Primitive(PrimitiveType::Int)).into(), + NestedField::map_value_element(value_id, Type::Primitive(value), true).into(), + )) + } + + fn list_type(element_id: i32, element: PrimitiveType) -> Type { + Type::List(ListType::new( + NestedField::list_element(element_id, Type::Primitive(element), true).into(), + )) + } + + #[test] + fn test_schema_is_spec_manifest_entry() { + let spec = partition_spec(&[("i", Transform::Identity), ("s", Transform::Truncate(2))]); + let batch = manifest_to_record_batch(&manifest(spec, vec![])).unwrap(); + assert_eq!(batch.num_rows(), 0); + + let long = || Type::Primitive(PrimitiveType::Long); + let int = || Type::Primitive(PrimitiveType::Int); + let string = || Type::Primitive(PrimitiveType::String); + let data_file = StructType::new(vec![ + NestedField::required(134, "content", int()).into(), + NestedField::required(100, "file_path", string()).into(), + NestedField::required(101, "file_format", string()).into(), + NestedField::required( + 102, + "partition", + Type::Struct(StructType::new(vec![ + NestedField::optional(1000, "i_identity", int()).into(), + NestedField::optional(1001, "s_truncate[2]", string()).into(), + ])), + ) + .into(), + NestedField::required(103, "record_count", long()).into(), + NestedField::required(104, "file_size_in_bytes", long()).into(), + NestedField::optional( + 108, + "column_sizes", + int_map_type(117, 118, PrimitiveType::Long), + ) + .into(), + NestedField::optional( + 109, + "value_counts", + int_map_type(119, 120, PrimitiveType::Long), + ) + .into(), + NestedField::optional( + 110, + "null_value_counts", + int_map_type(121, 122, PrimitiveType::Long), + ) + .into(), + NestedField::optional( + 137, + "nan_value_counts", + int_map_type(138, 139, PrimitiveType::Long), + ) + .into(), + NestedField::optional( + 125, + "lower_bounds", + int_map_type(126, 127, PrimitiveType::Binary), + ) + .into(), + NestedField::optional( + 128, + "upper_bounds", + int_map_type(129, 130, PrimitiveType::Binary), + ) + .into(), + NestedField::optional(131, "key_metadata", Type::Primitive(PrimitiveType::Binary)) + .into(), + NestedField::optional(132, "split_offsets", list_type(133, PrimitiveType::Long)).into(), + NestedField::optional(135, "equality_ids", list_type(136, PrimitiveType::Int)).into(), + NestedField::optional(140, "sort_order_id", int()).into(), + NestedField::optional(142, "first_row_id", long()).into(), + NestedField::optional(143, "referenced_data_file", string()).into(), + NestedField::optional(144, "content_offset", long()).into(), + NestedField::optional(145, "content_size_in_bytes", long()).into(), + ]); + let expected = Schema::builder() + .with_fields(vec![ + NestedField::required(0, "status", int()).into(), + NestedField::optional(1, "snapshot_id", long()).into(), + NestedField::optional(3, "sequence_number", long()).into(), + NestedField::optional(4, "file_sequence_number", long()).into(), + NestedField::required(2, "data_file", Type::Struct(data_file)).into(), + ]) + .build() + .unwrap(); + assert_eq!( + arrow_schema_to_schema(&batch.schema()).unwrap().as_struct(), + expected.as_struct() + ); + } + + #[test] + fn test_entry_and_data_file_columns() { + let added = DataFile { + column_sizes: HashMap::from([(1, 11)]), + value_counts: HashMap::from([(1, 10)]), + null_value_counts: HashMap::from([(1, 1)]), + nan_value_counts: HashMap::from([(4, 0)]), + lower_bounds: HashMap::from([(2, Datum::int(-1))]), + upper_bounds: HashMap::from([(2, Datum::int(7))]), + key_metadata: Some(vec![1, 2]), + split_offsets: Some(vec![4, 100]), + sort_order_id: Some(0), + first_row_id: Some(1000), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .record_count(10) + .file_size_in_bytes(100) + .build() + .unwrap() + }; + let position_deletes = DataFile { + file_format: DataFileFormat::Puffin, + referenced_data_file: Some("s3://b/data/0.parquet".to_string()), + content_offset: Some(4), + content_size_in_bytes: Some(40), + ..data_file(DataContentType::PositionDeletes, "s3://b/data/1.puffin") + .build() + .unwrap() + }; + let equality_deletes = DataFile { + equality_ids: Some(vec![1, 3]), + split_offsets: Some(vec![]), + ..data_file(DataContentType::EqualityDeletes, "s3://b/data/2.parquet") + .build() + .unwrap() + }; + let entries = vec![ + ManifestEntry { + snapshot_id: Some(1), + sequence_number: Some(2), + file_sequence_number: Some(3), + ..entry(ManifestStatus::Added, added) + }, + entry(ManifestStatus::Existing, position_deletes), + ManifestEntry { + snapshot_id: Some(5), + ..entry(ManifestStatus::Deleted, equality_deletes) + }, + ]; + let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), entries)).unwrap(); + assert_eq!(batch.num_rows(), 3); + + let column = |name: &str| batch.column_by_name(name).unwrap().clone(); + assert_eq!( + column("status").as_ref(), + &Int32Array::from(vec![1, 0, 2]) as &dyn Array + ); + assert_eq!( + column("snapshot_id").as_ref(), + &Int64Array::from(vec![Some(1), None, Some(5)]) as &dyn Array + ); + assert_eq!( + column("sequence_number").as_ref(), + &Int64Array::from(vec![Some(2), None, None]) as &dyn Array + ); + assert_eq!( + column("file_sequence_number").as_ref(), + &Int64Array::from(vec![Some(3), None, None]) as &dyn Array + ); + + let column = |name: &str| data_file_column(&batch, name); + assert_eq!( + column("content").as_ref(), + &Int32Array::from(vec![0, 1, 2]) as &dyn Array + ); + assert_eq!( + column("file_path").as_ref(), + &StringArray::from(vec![ + "s3://b/data/0.parquet", + "s3://b/data/1.puffin", + "s3://b/data/2.parquet" + ]) as &dyn Array + ); + // Writers, including `ManifestWriter`, store the format name in upper case. + assert_eq!( + column("file_format").as_ref(), + &StringArray::from(vec!["PARQUET", "PUFFIN", "PARQUET"]) as &dyn Array + ); + assert_eq!( + column("record_count").as_ref(), + &Int64Array::from(vec![10, 1, 1]) as &dyn Array + ); + assert_eq!( + column("file_size_in_bytes").as_ref(), + &Int64Array::from(vec![100, 10, 10]) as &dyn Array + ); + for (name, expected) in [ + ("column_sizes", (1, 11)), + ("value_counts", (1, 10)), + ("null_value_counts", (1, 1)), + ("nan_value_counts", (4, 0)), + ] { + let map = column(name); + assert_eq!(map.null_count(), 0, "{name}"); + assert_eq!(long_map(&map, 0), vec![expected], "{name}"); + assert_eq!(long_map(&map, 1), vec![], "{name}"); + assert_eq!(long_map(&map, 2), vec![], "{name}"); + } + assert_eq!(binary_map(&column("lower_bounds"), 0), vec![( + 2, + (-1i32).to_le_bytes().to_vec() + )]); + assert_eq!(binary_map(&column("upper_bounds"), 0), vec![( + 2, + 7i32.to_le_bytes().to_vec() + )]); + assert_eq!(binary_map(&column("lower_bounds"), 1), vec![]); + assert_eq!( + column("key_metadata").as_ref(), + &LargeBinaryArray::from(vec![Some(&[1u8, 2][..]), None, None]) as &dyn Array + ); + let split_offsets = column("split_offsets"); + let split_offsets = split_offsets.as_list::(); + assert!(split_offsets.is_valid(0) && split_offsets.is_null(1) && split_offsets.is_valid(2)); + assert_eq!( + split_offsets.value(0).as_ref(), + &Int64Array::from(vec![4, 100]) as &dyn Array + ); + assert_eq!(split_offsets.value(2).len(), 0); + let equality_ids = column("equality_ids"); + let equality_ids = equality_ids.as_list::(); + assert!(equality_ids.is_null(0) && equality_ids.is_null(1) && equality_ids.is_valid(2)); + assert_eq!( + equality_ids.value(2).as_ref(), + &Int32Array::from(vec![1, 3]) as &dyn Array + ); + assert_eq!( + column("sort_order_id").as_ref(), + &Int32Array::from(vec![Some(0), None, None]) as &dyn Array + ); + assert_eq!( + column("first_row_id").as_ref(), + &Int64Array::from(vec![Some(1000), None, None]) as &dyn Array + ); + assert_eq!( + column("referenced_data_file").as_ref(), + &StringArray::from(vec![None, Some("s3://b/data/0.parquet"), None]) as &dyn Array + ); + assert_eq!( + column("content_offset").as_ref(), + &Int64Array::from(vec![None, Some(4), None]) as &dyn Array + ); + assert_eq!( + column("content_size_in_bytes").as_ref(), + &Int64Array::from(vec![None, Some(40), None]) as &dyn Array + ); + } + + #[test] + fn test_metric_maps_with_several_columns() { + let uuid = Uuid::parse_str(UUID).unwrap(); + let file = DataFile { + value_counts: HashMap::from([(1, 5), (3, 6), (13, 7)]), + lower_bounds: HashMap::from([ + (1, Datum::bool(false)), + (3, Datum::long(-2)), + (13, Datum::string("abc")), + (14, Datum::uuid(uuid)), + ]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }; + let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), vec![entry( + ManifestStatus::Added, + file, + )])) + .unwrap(); + + assert_eq!( + long_map(&data_file_column(&batch, "value_counts"), 0), + vec![(1, 5), (3, 6), (13, 7)] + ); + assert_eq!( + binary_map(&data_file_column(&batch, "lower_bounds"), 0), + vec![ + (1, vec![0]), + (3, (-2i64).to_le_bytes().to_vec()), + (13, b"abc".to_vec()), + (14, uuid.as_bytes().to_vec()), + ] + ); + } + + #[test] + fn test_partition_values() { + let identity_columns = [ + "b", "i", "l", "f", "d", "dec", "date", "time", "ts", "tstz", "ts_ns", "tstz_ns", "s", + "u", "fx", "bin", + ]; + let mut fields: Vec<_> = identity_columns + .iter() + .map(|column| (*column, Transform::Identity)) + .collect(); + fields.extend([ + ("i", Transform::Bucket(4)), + ("ts", Transform::Day), + ("s", Transform::Truncate(2)), + ("l", Transform::Void), + ]); + let spec = partition_spec(&fields); + + let uuid = Uuid::parse_str(UUID).unwrap(); + let values = Struct::from_iter([ + Some(Literal::bool(true)), + Some(Literal::int(-7)), + Some(Literal::long(1i64 << 40)), + Some(Literal::float(1.5f32)), + Some(Literal::double(-2.25)), + Some(Literal::decimal(-12_345)), + Some(Literal::date(19_000)), + Some(Literal::time(3_600_000_000)), + Some(Literal::timestamp(1_700_000_000_000_000)), + Some(Literal::timestamptz(1_700_000_000_000_001)), + Some(Literal::timestamp_nano(1_700_000_000_000_000_002)), + Some(Literal::timestamptz_nano(1_700_000_000_000_000_003)), + Some(Literal::string("abc")), + Some(Literal::uuid(uuid)), + Some(Literal::fixed([1u8, 2, 3])), + Some(Literal::binary([9u8])), + Some(Literal::int(3)), + Some(Literal::date(19_675)), + Some(Literal::string("ab")), + None, + ]); + let nulls = Struct::from_iter(fields.iter().map(|_| None)); + let entries = [values, nulls] + .into_iter() + .map(|partition| { + entry(ManifestStatus::Added, DataFile { + partition, + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }) + }) + .collect(); + let batch = manifest_to_record_batch(&manifest(spec, entries)).unwrap(); + + let partition = data_file_column(&batch, "partition"); + let partition = partition.as_struct(); + assert_eq!(partition.null_count(), 0); + let expected: Vec = vec![ + Arc::new(BooleanArray::from(vec![Some(true), None])), + Arc::new(Int32Array::from(vec![Some(-7), None])), + Arc::new(Int64Array::from(vec![Some(1i64 << 40), None])), + Arc::new(Float32Array::from(vec![Some(1.5), None])), + Arc::new(Float64Array::from(vec![Some(-2.25), None])), + Arc::new( + Decimal128Array::from(vec![Some(-12_345), None]) + .with_precision_and_scale(38, 10) + .unwrap(), + ), + Arc::new(Date32Array::from(vec![Some(19_000), None])), + Arc::new(Time64MicrosecondArray::from(vec![ + Some(3_600_000_000), + None, + ])), + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(1_700_000_000_000_000), + None, + ])), + Arc::new( + TimestampMicrosecondArray::from(vec![Some(1_700_000_000_000_001), None]) + .with_timezone("+00:00"), + ), + Arc::new(TimestampNanosecondArray::from(vec![ + Some(1_700_000_000_000_000_002), + None, + ])), + Arc::new( + TimestampNanosecondArray::from(vec![Some(1_700_000_000_000_000_003), None]) + .with_timezone("+00:00"), + ), + Arc::new(StringArray::from(vec![Some("abc"), None])), + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(uuid.as_bytes().to_vec()), None].into_iter(), + 16, + ) + .unwrap(), + ), + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some(vec![1u8, 2, 3]), None].into_iter(), + 3, + ) + .unwrap(), + ), + Arc::new(LargeBinaryArray::from(vec![Some(&[9u8][..]), None])), + Arc::new(Int32Array::from(vec![Some(3), None])), + Arc::new(Date32Array::from(vec![Some(19_675), None])), + Arc::new(StringArray::from(vec![Some("ab"), None])), + Arc::new(Int64Array::from(vec![None, None])), + ]; + assert_eq!(partition.num_columns(), expected.len()); + for (index, expected) in expected.iter().enumerate() { + assert_eq!( + partition.column(index).as_ref(), + expected.as_ref(), + "partition field {}", + partition.fields()[index].name() + ); + } + } + + #[test] + fn test_unpartitioned_partition_column() { + let entries = (0..3) + .map(|i| { + entry( + ManifestStatus::Added, + data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) + .build() + .unwrap(), + ) + }) + .collect(); + let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), entries)).unwrap(); + let partition = data_file_column(&batch, "partition"); + let partition: &StructArray = partition.as_struct(); + assert_eq!(partition.num_columns(), 0); + assert_eq!(partition.len(), 3); + assert_eq!(partition.null_count(), 0); + } + + #[test] + fn test_unknown_partition_field() { + let spec = || partition_spec(&[("unk", Transform::Void)]); + let entries = vec![entry(ManifestStatus::Added, DataFile { + partition: Struct::from_iter([None]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + })]; + let batch = manifest_to_record_batch(&manifest(spec(), entries)).unwrap(); + let partition = data_file_column(&batch, "partition"); + assert_eq!( + partition.as_struct().column(0).as_ref(), + &NullArray::new(1) as &dyn Array + ); + + let entries = vec![entry(ManifestStatus::Added, DataFile { + partition: Struct::from_iter([Some(Literal::int(1))]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + })]; + let err = manifest_to_record_batch(&manifest(spec(), entries)).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[test] + fn test_partition_value_of_wrong_type_is_an_error() { + let file = DataFile { + partition: Struct::from_iter([Some(Literal::string("not an int"))]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }; + let manifest = manifest(partition_spec(&[("i", Transform::Identity)]), vec![entry( + ManifestStatus::Added, + file, + )]); + let err = manifest_to_record_batch(&manifest).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[test] + fn test_fixed_partition_value_of_wrong_width_is_an_error() { + let file = DataFile { + partition: Struct::from_iter([Some(Literal::fixed([1u8, 2]))]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }; + let manifest = manifest(partition_spec(&[("fx", Transform::Identity)]), vec![entry( + ManifestStatus::Added, + file, + )]); + let err = manifest_to_record_batch(&manifest).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[test] + fn test_partition_value_with_wrong_arity_is_an_error() { + let file = DataFile { + partition: Struct::from_iter([Some(Literal::int(1)), Some(Literal::int(2))]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }; + let manifest = manifest(partition_spec(&[("i", Transform::Identity)]), vec![entry( + ManifestStatus::Added, + file, + )]); + let err = manifest_to_record_batch(&manifest).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } +} diff --git a/crates/iceberg/src/arrow/mod.rs b/crates/iceberg/src/arrow/mod.rs index 28cf6e94a4..daaa1516e6 100644 --- a/crates/iceberg/src/arrow/mod.rs +++ b/crates/iceberg/src/arrow/mod.rs @@ -28,6 +28,8 @@ pub mod delete_file_loader; pub(crate) mod delete_filter; mod int96; +mod manifest; +pub use manifest::*; mod reader; /// RecordBatch projection utilities pub mod record_batch_projector; diff --git a/crates/iceberg/src/spec/manifest/entry.rs b/crates/iceberg/src/spec/manifest/entry.rs index 555d86a25c..bbd2965d53 100644 --- a/crates/iceberg/src/spec/manifest/entry.rs +++ b/crates/iceberg/src/spec/manifest/entry.rs @@ -535,6 +535,24 @@ fn data_file_fields_v3(partition_type: &StructType) -> Vec { ] } +/// Fields of the `manifest_entry` struct in the latest format version. A +/// [`ManifestEntry`] of any version fits them: earlier versions lack only +/// optional fields, apart from the deprecated v1 `block_size_in_bytes`, which +/// [`DataFile`] doesn't keep. +pub(crate) fn manifest_entry_fields(partition_type: &StructType) -> Vec { + vec![ + STATUS.clone(), + SNAPSHOT_ID_V2.clone(), + SEQUENCE_NUMBER.clone(), + FILE_SEQUENCE_NUMBER.clone(), + Arc::new(NestedField::required( + 2, + "data_file", + Type::Struct(StructType::new(data_file_fields_v3(partition_type))), + )), + ] +} + pub(super) fn data_file_schema_v3(partition_type: &StructType) -> Result { let schema = Schema::builder() .with_fields(data_file_fields_v3(partition_type)) From 38e3a0145e073eebbc50cd5bedc543c29f7682ba Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 9 Oct 2026 17:41:20 -0700 Subject: [PATCH 2/7] test(arrow): cover manifest versions, count overflow, and a PyIceberg manifest --- crates/iceberg/src/arrow/manifest.rs | 125 +++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index e12d3a150b..e38f9de743 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -50,6 +50,11 @@ use crate::{Error, ErrorKind, Result}; /// entries in no particular order, and a metric map with no entries is empty /// rather than null. /// +/// Entries are converted as they are in `manifest`. Fields that readers +/// inherit from the manifest list, such as `snapshot_id` and +/// `sequence_number`, stay null in a manifest from [`Manifest::parse_avro`]. +/// [`ManifestReader::read`](crate::spec::ManifestReader::read) fills them in. +/// /// # Example /// /// ``` @@ -510,6 +515,8 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use apache_avro::Reader as AvroReader; + use apache_avro::types::Value as AvroValue; use arrow_array::cast::AsArray; use arrow_array::types::{Int32Type, Int64Type}; use arrow_array::{ @@ -1159,4 +1166,122 @@ mod tests { let err = manifest_to_record_batch(&manifest).unwrap_err(); assert_eq!(err.kind(), ErrorKind::DataInvalid); } + + #[test] + fn test_every_format_version_has_the_same_schema() { + let schemas: Vec<_> = [FormatVersion::V1, FormatVersion::V2, FormatVersion::V3] + .into_iter() + .map(|format_version| { + let manifest = Manifest::new( + ManifestMetadata { + format_version, + ..manifest(partition_spec(&[("i", Transform::Identity)]), vec![]) + .metadata() + .clone() + }, + vec![], + ); + manifest_to_record_batch(&manifest).unwrap().schema() + }) + .collect(); + assert_eq!(schemas[0], schemas[1]); + assert_eq!(schemas[1], schemas[2]); + } + + #[test] + fn test_count_beyond_long_range_is_an_error() { + let file = DataFile { + column_sizes: HashMap::from([(1, u64::MAX)]), + ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap() + }; + let manifest = manifest(partition_spec(&[]), vec![entry( + ManifestStatus::Added, + file, + )]); + let err = manifest_to_record_batch(&manifest).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + /// `(key, value)` pairs of the map field `name` of an Avro `data_file` + /// record, sorted by key. + fn avro_map( + data_file: &[(String, AvroValue)], + name: &str, + value: impl Fn(&AvroValue) -> T, + ) -> Vec<(i32, T)> { + let (_, AvroValue::Union(_, map)) = + data_file.iter().find(|(field, _)| field == name).unwrap() + else { + panic!("{name} is not a union"); + }; + let AvroValue::Array(pairs) = map.as_ref() else { + panic!("{name} is not an array"); + }; + let mut pairs: Vec<_> = pairs + .iter() + .map(|pair| { + let AvroValue::Record(pair) = pair else { + panic!("{name} entry is not a record"); + }; + let AvroValue::Int(key) = pair[0].1 else { + panic!("{name} key is not an int"); + }; + (key, value(&pair[1].1)) + }) + .collect(); + pairs.sort_unstable(); + pairs + } + + #[test] + fn test_manifest_written_by_pyiceberg() { + let bs = std::fs::read(format!( + "{}/testdata/manifests/pyiceberg-v2-data.avro", + env!("CARGO_MANIFEST_DIR") + )) + .unwrap(); + let batch = manifest_to_record_batch(&Manifest::parse_avro(&bs).unwrap()).unwrap(); + + let records: Vec<_> = AvroReader::new(bs.as_slice()) + .unwrap() + .map(|record| match record.unwrap() { + AvroValue::Record(fields) => fields, + _ => panic!("manifest entry is not a record"), + }) + .collect(); + assert_eq!(batch.num_rows(), records.len()); + for (row, record) in records.iter().enumerate() { + let Some((_, AvroValue::Record(data_file))) = + record.iter().find(|(name, _)| name == "data_file") + else { + panic!("data_file is not a record"); + }; + let long = |value: &AvroValue| match value { + AvroValue::Long(v) => *v, + _ => panic!("map value is not a long"), + }; + let bytes = |value: &AvroValue| match value { + AvroValue::Bytes(v) => v.clone(), + _ => panic!("map value is not bytes"), + }; + for name in ["column_sizes", "value_counts", "null_value_counts"] { + assert_eq!( + long_map(&data_file_column(&batch, name), row), + avro_map(data_file, name, long), + "{name}" + ); + } + for name in ["lower_bounds", "upper_bounds"] { + let expected = avro_map(data_file, name, bytes); + assert!(!expected.is_empty(), "{name}"); + assert_eq!( + binary_map(&data_file_column(&batch, name), row), + expected, + "{name}" + ); + } + } + } } From d3b06003b6c29283769e2641b257fbddf8b8e0f5 Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 9 Oct 2026 17:49:15 -0700 Subject: [PATCH 3/7] refactor(arrow): assert the shape of the generated manifest_entry schema --- crates/iceberg/src/arrow/manifest.rs | 81 +++++++++++----------------- 1 file changed, 30 insertions(+), 51 deletions(-) diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index e38f9de743..ba8a92e5ae 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -29,7 +29,7 @@ use arrow_array::{ Float64Array, Int32Array, Int64Array, LargeBinaryArray, NullArray, RecordBatch, StringArray, StructArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, }; -use arrow_schema::{DataType, Field, FieldRef, Fields}; +use arrow_schema::{DataType, Field, FieldRef}; use crate::arrow::{UTC_TIME_ZONE, schema_to_arrow_schema}; use crate::error::invalid_data; @@ -111,13 +111,11 @@ pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { .with_fields(manifest_entry_fields(&partition_type)) .build()?, )?); - let data_file_fields = struct_fields(schema.field_with_name("data_file")?)?; - let child = |name: &str| -> Result<&FieldRef> { - data_file_fields - .find(name) - .map(|(_, field)| field) - .ok_or_else(|| Error::new(ErrorKind::Unexpected, format!("no data_file field {name}"))) + let DataType::Struct(data_file_fields) = schema.field_with_name("data_file")?.data_type() + else { + unreachable!() }; + let child = |name: &str| data_file_fields.find(name).unwrap().1; let entries = manifest.entries(); let len = entries.len(); @@ -130,17 +128,17 @@ pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { let mut file_format = StringBuilder::new(); let mut record_count = Int64Builder::with_capacity(len); let mut file_size_in_bytes = Int64Builder::with_capacity(len); - let mut column_sizes = map_builder(child("column_sizes")?, Int64Builder::new())?; - let mut value_counts = map_builder(child("value_counts")?, Int64Builder::new())?; - let mut null_value_counts = map_builder(child("null_value_counts")?, Int64Builder::new())?; - let mut nan_value_counts = map_builder(child("nan_value_counts")?, Int64Builder::new())?; - let mut lower_bounds = map_builder(child("lower_bounds")?, LargeBinaryBuilder::new())?; - let mut upper_bounds = map_builder(child("upper_bounds")?, LargeBinaryBuilder::new())?; + let mut column_sizes = map_builder(child("column_sizes"), Int64Builder::new()); + let mut value_counts = map_builder(child("value_counts"), Int64Builder::new()); + let mut null_value_counts = map_builder(child("null_value_counts"), Int64Builder::new()); + let mut nan_value_counts = map_builder(child("nan_value_counts"), Int64Builder::new()); + let mut lower_bounds = map_builder(child("lower_bounds"), LargeBinaryBuilder::new()); + let mut upper_bounds = map_builder(child("upper_bounds"), LargeBinaryBuilder::new()); let mut key_metadata = LargeBinaryBuilder::new(); let mut split_offsets = - ListBuilder::new(Int64Builder::new()).with_field(list_element(child("split_offsets")?)?); + ListBuilder::new(Int64Builder::new()).with_field(list_element(child("split_offsets"))); let mut equality_ids = - ListBuilder::new(Int32Builder::new()).with_field(list_element(child("equality_ids")?)?); + ListBuilder::new(Int32Builder::new()).with_field(list_element(child("equality_ids"))); let mut sort_order_id = Int32Builder::with_capacity(len); let mut first_row_id = Int64Builder::with_capacity(len); let mut referenced_data_file = StringBuilder::new(); @@ -193,7 +191,7 @@ pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { content_size_in_bytes.append_option(data_file.content_size_in_bytes); } - let partition = partition_column(entries, &partition_type, child("partition")?)?; + let partition = partition_column(entries, &partition_type, child("partition"))?; let data_file = StructArray::try_new( data_file_fields.clone(), vec![ @@ -230,49 +228,31 @@ pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { ])?) } -fn struct_fields(field: &Field) -> Result<&Fields> { - match field.data_type() { - DataType::Struct(fields) => Ok(fields), - data_type => Err(Error::new( - ErrorKind::Unexpected, - format!("field {} is {data_type}, not a struct", field.name()), - )), - } -} - -fn list_element(field: &Field) -> Result { - match field.data_type() { - DataType::List(element) => Ok(element.clone()), - data_type => Err(Error::new( - ErrorKind::Unexpected, - format!("field {} is {data_type}, not a list", field.name()), - )), - } +fn list_element(field: &Field) -> FieldRef { + let DataType::List(element) = field.data_type() else { + unreachable!() + }; + element.clone() } /// Creates a builder for the `map` column `field` that produces the /// key and value fields of its schema, including their field IDs. -fn map_builder(field: &Field, values: V) -> Result> { +fn map_builder(field: &Field, values: V) -> MapBuilder { let DataType::Map(entries, _) = field.data_type() else { - return Err(Error::new( - ErrorKind::Unexpected, - format!("field {} is {}, not a map", field.name(), field.data_type()), - )); + unreachable!() }; - let [key, value] = &struct_fields(entries)?[..] else { - return Err(Error::new( - ErrorKind::Unexpected, - format!("map {} doesn't have two entry fields", field.name()), - )); + let DataType::Struct(key_value) = entries.data_type() else { + unreachable!() }; + let (key, value) = (&key_value[0], &key_value[1]); let names = MapFieldNames { entry: entries.name().clone(), key: key.name().clone(), value: value.name().clone(), }; - Ok(MapBuilder::new(Some(names), Int32Builder::new(), values) + MapBuilder::new(Some(names), Int32Builder::new(), values) .with_keys_field(key.clone()) - .with_values_field(value.clone())) + .with_values_field(value.clone()) } fn append_counts( @@ -315,7 +295,9 @@ fn partition_column( )); } - let fields = struct_fields(field)?; + let DataType::Struct(fields) = field.data_type() else { + unreachable!() + }; if fields.is_empty() { return Ok(Arc::new(StructArray::new_empty_fields(entries.len(), None))); } @@ -448,10 +430,7 @@ fn partition_array( )?), PrimitiveType::Decimal { .. } => { let DataType::Decimal128(precision, scale) = *data_type else { - return Err(Error::new( - ErrorKind::Unexpected, - format!("decimal partition field {} is {data_type}", field.name), - )); + unreachable!() }; Arc::new( collect::<_, Decimal128Array>(field, values, |literal| match literal { From ff0ee7a76eb6459715d3c1cf4e06f3a852df44cf Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Fri, 9 Oct 2026 18:45:08 -0700 Subject: [PATCH 4/7] feat(arrow): read manifest entries into Arrow with arrow-avro --- Cargo.lock | 27 +- Cargo.toml | 5 + crates/iceberg/Cargo.toml | 1 + crates/iceberg/public-api.txt | 2 +- crates/iceberg/src/arrow/manifest.rs | 1562 +++++++++++++------------- 5 files changed, 816 insertions(+), 781 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f3694ad2e2..bc09e7cadb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -278,6 +278,28 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-avro" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9145b685d586482102e4fe3d40349fe00086c0b219d51cae9daaead4698849ea" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-schema", + "bytes", + "crc", + "flate2", + "indexmap 2.14.0", + "rand 0.9.5", + "serde", + "serde_json", + "snap", + "strum_macros 0.28.0", + "uuid", + "zstd", +] + [[package]] name = "arrow-buffer" version = "59.3.0" @@ -1413,7 +1435,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -2716,6 +2738,7 @@ dependencies = [ "array-init", "arrow-arith", "arrow-array", + "arrow-avro", "arrow-buffer", "arrow-cast", "arrow-ord", @@ -7153,7 +7176,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index bb9770275a..13c44d06c5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,11 @@ array-init = "2" arrow = "59.2" arrow-arith = "59.2" arrow-array = "59.2" +arrow-avro = { version = "59.2", default-features = false, features = [ + "deflate", + "snappy", + "zstd", +] } arrow-buffer = "59.2" arrow-cast = "59.2" arrow-ord = "59.2" diff --git a/crates/iceberg/Cargo.toml b/crates/iceberg/Cargo.toml index f22ed59b9b..6dfc76ecb9 100644 --- a/crates/iceberg/Cargo.toml +++ b/crates/iceberg/Cargo.toml @@ -39,6 +39,7 @@ apache-avro = { workspace = true } array-init = { workspace = true } arrow-arith = { workspace = true } arrow-array = { workspace = true } +arrow-avro = { workspace = true } arrow-buffer = { workspace = true } arrow-cast = { workspace = true } arrow-ord = { workspace = true } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 90e899cb13..9d4667bbcc 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -130,7 +130,7 @@ pub fn iceberg::arrow::arrow_schema_to_schema_auto_assign_ids(schema: &arrow_sch pub fn iceberg::arrow::arrow_struct_to_literal(struct_array: &arrow_array::array::ArrayRef, ty: &iceberg::spec::StructType) -> iceberg::Result>> pub fn iceberg::arrow::arrow_type_to_type(ty: &arrow_schema::datatype::DataType) -> iceberg::Result pub fn iceberg::arrow::datum_to_arrow_type_with_ree(datum: &iceberg::spec::Datum) -> arrow_schema::datatype::DataType -pub fn iceberg::arrow::manifest_to_record_batch(manifest: &iceberg::spec::Manifest) -> iceberg::Result +pub fn iceberg::arrow::read_manifest_entries(bytes: &[u8], manifest_file: &iceberg::spec::ManifestFile) -> iceberg::Result pub fn iceberg::arrow::schema_to_arrow_schema(schema: &iceberg::spec::Schema) -> iceberg::Result pub fn iceberg::arrow::strip_metadata_from_schema(schema: &arrow_schema::schema::Schema) -> iceberg::Result pub fn iceberg::arrow::type_to_arrow_type(ty: &iceberg::spec::Type) -> iceberg::Result diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index ba8a92e5ae..15879870d4 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -15,94 +15,112 @@ // specific language governing permissions and limitations // under the License. -//! Conversion of manifest entries to Arrow. +//! Reading manifest entries into Arrow. use std::collections::HashMap; use std::sync::Arc; -use arrow_array::builder::{ - ArrayBuilder, Int32Builder, Int64Builder, LargeBinaryBuilder, ListBuilder, MapBuilder, - MapFieldNames, StringBuilder, -}; +use arrow_arith::boolean::{and, is_null}; +use arrow_array::cast::AsArray; +use arrow_array::types::{TimestampMicrosecondType, TimestampNanosecondType}; use arrow_array::{ - ArrayRef, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, Float32Array, - Float64Array, Int32Array, Int64Array, LargeBinaryArray, NullArray, RecordBatch, StringArray, - StructArray, Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, + Array, ArrayRef, BooleanArray, Int32Array, Int64Array, ListArray, MapArray, RecordBatch, + StructArray, new_null_array, }; -use arrow_schema::{DataType, Field, FieldRef}; +use arrow_avro::reader::ReaderBuilder; +use arrow_avro::schema::AvroSchema; +use arrow_cast::{CastOptions, cast_with_options}; +use arrow_ord::cmp::eq; +use arrow_schema::{ArrowError, DataType, Field, Fields, TimeUnit}; +use arrow_select::zip::zip; +use parquet::arrow::PARQUET_FIELD_ID_META_KEY; +use serde_json::Value as JsonValue; -use crate::arrow::{UTC_TIME_ZONE, schema_to_arrow_schema}; +use crate::arrow::schema_to_arrow_schema; use crate::error::invalid_data; use crate::spec::{ - DataFileFormat, Datum, Literal, Manifest, ManifestEntryRef, NestedField, PrimitiveLiteral, - PrimitiveType, Schema, StructType, manifest_entry_fields, + ManifestContentType, ManifestFile, ManifestMetadata, ManifestStatus, Schema, + manifest_entry_fields, }; -use crate::{Error, ErrorKind, Result}; +use crate::{Error, Result}; + +/// Field ID of `data_file`. +const DATA_FILE_FIELD_ID: i64 = 2; +/// Field ID of `data_file.partition`. +const PARTITION_FIELD_ID: i64 = 102; +/// Field ID of `data_file.content`. +const CONTENT_FIELD_ID: i32 = 134; +/// Field ID of `sequence_number`. +const SEQUENCE_NUMBER_FIELD_ID: i32 = 3; +/// Field ID of `file_sequence_number`. +const FILE_SEQUENCE_NUMBER_FIELD_ID: i32 = 4; -/// Converts the entries of `manifest` into a [`RecordBatch`] with one row per -/// entry. +/// Reads the entries of a manifest file into a [`RecordBatch`] with one row +/// per entry. /// -/// The schema is the spec's `manifest_entry` struct for the latest format -/// version, with field IDs in the Arrow field metadata, so manifests of every -/// format version convert to the same schema apart from `data_file.partition`, -/// which follows the manifest's partition spec. Lower and upper bounds use the -/// spec's binary single-value serialization. The metric maps list their -/// entries in no particular order, and a metric map with no entries is empty -/// rather than null. +/// `bytes` is the manifest file and `manifest_file` is its entry in the +/// manifest list. The schema is the spec's `manifest_entry` struct for the +/// latest format version, with field IDs in the Arrow field metadata, so +/// manifests of every format version read into the same schema apart from +/// `data_file.partition`, which follows the manifest's partition spec. /// -/// Entries are converted as they are in `manifest`. Fields that readers -/// inherit from the manifest list, such as `snapshot_id` and -/// `sequence_number`, stay null in a manifest from [`Manifest::parse_avro`]. -/// [`ManifestReader::read`](crate::spec::ManifestReader::read) fills them in. +/// Fields are matched to the manifest's Avro schema by field ID. A v1 +/// manifest reads with the spec's v2+ defaults, and null snapshot IDs, +/// sequence numbers, and first row IDs are inherited from `manifest_file`. +/// Lower and upper bounds keep the binary single-value serialization they were +/// written with. /// /// # Example /// /// ``` -/// use std::sync::Arc; +/// use iceberg::arrow::read_manifest_entries; +/// use iceberg::spec::{ManifestContentType, ManifestFile}; /// -/// use iceberg::arrow::manifest_to_record_batch; -/// use iceberg::spec::{ -/// DataContentType, DataFileBuilder, DataFileFormat, FormatVersion, Manifest, -/// ManifestContentType, ManifestEntry, ManifestMetadata, ManifestStatus, NestedField, -/// PartitionSpec, PrimitiveType, Schema, Type, +/// let bytes = std::fs::read(concat!( +/// env!("CARGO_MANIFEST_DIR"), +/// "/testdata/manifests/pyiceberg-v2-data.avro" +/// ))?; +/// let manifest_file = ManifestFile { +/// manifest_path: "s3://bucket/metadata/manifest.avro".to_string(), +/// manifest_length: bytes.len() as i64, +/// partition_spec_id: 0, +/// content: ManifestContentType::Data, +/// sequence_number: 1, +/// min_sequence_number: 1, +/// added_snapshot_id: 1, +/// added_files_count: None, +/// existing_files_count: None, +/// deleted_files_count: None, +/// added_rows_count: None, +/// existing_rows_count: None, +/// deleted_rows_count: None, +/// partitions: None, +/// key_metadata: None, +/// first_row_id: None, /// }; /// -/// let schema = Arc::new( -/// Schema::builder() -/// .with_fields(vec![ -/// NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(), -/// ]) -/// .build()?, -/// ); -/// let data_file = DataFileBuilder::default() -/// .content(DataContentType::Data) -/// .file_path("s3://bucket/data/0.parquet".to_string()) -/// .file_format(DataFileFormat::Parquet) -/// .record_count(100) -/// .file_size_in_bytes(4096) -/// .build()?; -/// let manifest = Manifest::new( -/// ManifestMetadata { -/// schema: schema.clone(), -/// schema_id: 0, -/// partition_spec: PartitionSpec::builder(schema).build()?, -/// format_version: FormatVersion::V2, -/// content: ManifestContentType::Data, -/// }, -/// vec![ -/// ManifestEntry::builder() -/// .status(ManifestStatus::Added) -/// .data_file(data_file) -/// .build(), -/// ], -/// ); -/// -/// let batch = manifest_to_record_batch(&manifest)?; -/// assert_eq!(batch.num_rows(), 1); +/// let batch = read_manifest_entries(&bytes, &manifest_file)?; +/// assert_eq!(batch.num_rows(), 2); /// # Ok::<(), Box>(()) /// ``` -pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { - let metadata = manifest.metadata(); +pub fn read_manifest_entries(bytes: &[u8], manifest_file: &ManifestFile) -> Result { + // usize::MAX decodes the whole manifest into one batch. First row ID + // inheritance is a running sum over every entry in the file, and an empty + // manifest still returns its schema. arrow-avro sizes its buffers + // independently of the batch size, so this doesn't preallocate. + let builder = || ReaderBuilder::new().with_batch_size(usize::MAX); + let mut reader = builder().build(bytes).map_err(invalid_manifest)?; + let header = reader.avro_header(); + let metadata: HashMap> = header + .metadata() + .map(|(key, value)| (String::from_utf8_lossy(key).into_owned(), value.to_vec())) + .collect(); + let metadata = ManifestMetadata::parse(&metadata)?; + let mut avro_schema: JsonValue = serde_json::from_slice( + header + .get("avro.schema") + .ok_or_else(|| invalid_data!("manifest has no avro.schema"))?, + )?; let partition_type = metadata .partition_spec() .partition_type(metadata.schema())?; @@ -111,382 +129,299 @@ pub fn manifest_to_record_batch(manifest: &Manifest) -> Result { .with_fields(manifest_entry_fields(&partition_type)) .build()?, )?); - let DataType::Struct(data_file_fields) = schema.field_with_name("data_file")?.data_type() - else { - unreachable!() - }; - let child = |name: &str| data_file_fields.find(name).unwrap().1; - - let entries = manifest.entries(); - let len = entries.len(); - let mut status = Int32Builder::with_capacity(len); - let mut snapshot_id = Int64Builder::with_capacity(len); - let mut sequence_number = Int64Builder::with_capacity(len); - let mut file_sequence_number = Int64Builder::with_capacity(len); - let mut content = Int32Builder::with_capacity(len); - let mut file_path = StringBuilder::new(); - let mut file_format = StringBuilder::new(); - let mut record_count = Int64Builder::with_capacity(len); - let mut file_size_in_bytes = Int64Builder::with_capacity(len); - let mut column_sizes = map_builder(child("column_sizes"), Int64Builder::new()); - let mut value_counts = map_builder(child("value_counts"), Int64Builder::new()); - let mut null_value_counts = map_builder(child("null_value_counts"), Int64Builder::new()); - let mut nan_value_counts = map_builder(child("nan_value_counts"), Int64Builder::new()); - let mut lower_bounds = map_builder(child("lower_bounds"), LargeBinaryBuilder::new()); - let mut upper_bounds = map_builder(child("upper_bounds"), LargeBinaryBuilder::new()); - let mut key_metadata = LargeBinaryBuilder::new(); - let mut split_offsets = - ListBuilder::new(Int64Builder::new()).with_field(list_element(child("split_offsets"))); - let mut equality_ids = - ListBuilder::new(Int32Builder::new()).with_field(list_element(child("equality_ids"))); - let mut sort_order_id = Int32Builder::with_capacity(len); - let mut first_row_id = Int64Builder::with_capacity(len); - let mut referenced_data_file = StringBuilder::new(); - let mut content_offset = Int64Builder::with_capacity(len); - let mut content_size_in_bytes = Int64Builder::with_capacity(len); - for entry in entries { - status.append_value(entry.status as i32); - snapshot_id.append_option(entry.snapshot_id); - sequence_number.append_option(entry.sequence_number); - file_sequence_number.append_option(entry.file_sequence_number); - - let data_file = &entry.data_file; - content.append_value(data_file.content as i32); - file_path.append_value(&data_file.file_path); - // Writers store the format name in upper case. - file_format.append_value(match data_file.file_format { - DataFileFormat::Avro => "AVRO", - DataFileFormat::Orc => "ORC", - DataFileFormat::Parquet => "PARQUET", - DataFileFormat::Puffin => "PUFFIN", - }); - record_count.append_value(data_file.record_count.try_into()?); - file_size_in_bytes.append_value(data_file.file_size_in_bytes.try_into()?); - append_counts(&mut column_sizes, &data_file.column_sizes)?; - append_counts(&mut value_counts, &data_file.value_counts)?; - append_counts(&mut null_value_counts, &data_file.null_value_counts)?; - append_counts(&mut nan_value_counts, &data_file.nan_value_counts)?; - append_bounds(&mut lower_bounds, &data_file.lower_bounds)?; - append_bounds(&mut upper_bounds, &data_file.upper_bounds)?; - key_metadata.append_option(data_file.key_metadata.as_deref()); - match &data_file.split_offsets { - Some(offsets) => { - split_offsets.values().append_slice(offsets); - split_offsets.append(true); - } - None => split_offsets.append(false), - } - match &data_file.equality_ids { - Some(ids) => { - equality_ids.values().append_slice(ids); - equality_ids.append(true); - } - None => equality_ids.append(false), - } - sort_order_id.append_option(data_file.sort_order_id); - first_row_id.append_option(data_file.first_row_id); - referenced_data_file.append_option(data_file.referenced_data_file.as_deref()); - content_offset.append_option(data_file.content_offset); - content_size_in_bytes.append_option(data_file.content_size_in_bytes); + // TODO(#3257): arrow-avro 59 can't decode a record with no fields, which + // is the partition of every unpartitioned manifest, so project it out and + // build it in `missing_field`. arrow-avro 60 fixes this + // (apache/arrow-rs#10771). + if remove_empty_partition(&mut avro_schema) { + reader = builder() + .with_reader_schema(AvroSchema::new(avro_schema.to_string())) + .build(bytes) + .map_err(invalid_manifest)?; } - let partition = partition_column(entries, &partition_type, child("partition"))?; - let data_file = StructArray::try_new( - data_file_fields.clone(), - vec![ - Arc::new(content.finish()), - Arc::new(file_path.finish()), - Arc::new(file_format.finish()), - partition, - Arc::new(record_count.finish()), - Arc::new(file_size_in_bytes.finish()), - Arc::new(column_sizes.finish()), - Arc::new(value_counts.finish()), - Arc::new(null_value_counts.finish()), - Arc::new(nan_value_counts.finish()), - Arc::new(lower_bounds.finish()), - Arc::new(upper_bounds.finish()), - Arc::new(key_metadata.finish()), - Arc::new(split_offsets.finish()), - Arc::new(equality_ids.finish()), - Arc::new(sort_order_id.finish()), - Arc::new(first_row_id.finish()), - Arc::new(referenced_data_file.finish()), - Arc::new(content_offset.finish()), - Arc::new(content_size_in_bytes.finish()), - ], - None, + let Some(batch) = reader.next() else { + return Ok(RecordBatch::new_empty(schema)); + }; + let entries = project_struct( + &StructArray::from(batch.map_err(invalid_manifest)?), + &avro_schema, + schema.fields(), )?; - - Ok(RecordBatch::try_new(schema, vec![ - Arc::new(status.finish()), - Arc::new(snapshot_id.finish()), - Arc::new(sequence_number.finish()), - Arc::new(file_sequence_number.finish()), - Arc::new(data_file), - ])?) + Ok(RecordBatch::from(inherit(entries, manifest_file)?)) } -fn list_element(field: &Field) -> FieldRef { - let DataType::List(element) = field.data_type() else { - unreachable!() +/// Removes `data_file.partition` from a manifest's Avro schema if it has no +/// fields, and returns whether it did. +fn remove_empty_partition(avro_schema: &mut JsonValue) -> bool { + let has_field_id = |field: &JsonValue, id: i64| field["field-id"].as_i64() == Some(id); + let Some(data_file_fields) = avro_schema["fields"] + .as_array_mut() + .and_then(|fields| { + fields + .iter_mut() + .find(|field| has_field_id(field, DATA_FILE_FIELD_ID)) + }) + .and_then(|data_file| data_file["type"]["fields"].as_array_mut()) + else { + return false; }; - element.clone() + let len = data_file_fields.len(); + data_file_fields.retain(|field| { + !(has_field_id(field, PARTITION_FIELD_ID) + && field["type"]["fields"] + .as_array() + .is_some_and(Vec::is_empty)) + }); + data_file_fields.len() != len } -/// Creates a builder for the `map` column `field` that produces the -/// key and value fields of its schema, including their field IDs. -fn map_builder(field: &Field, values: V) -> MapBuilder { - let DataType::Map(entries, _) = field.data_type() else { - unreachable!() - }; - let DataType::Struct(key_value) = entries.data_type() else { - unreachable!() - }; - let (key, value) = (&key_value[0], &key_value[1]); - let names = MapFieldNames { - entry: entries.name().clone(), - key: key.name().clone(), - value: value.name().clone(), - }; - MapBuilder::new(Some(names), Int32Builder::new(), values) - .with_keys_field(key.clone()) - .with_values_field(value.clone()) +fn invalid_manifest(err: ArrowError) -> Error { + invalid_data!("failed to read manifest").with_source(err) } -fn append_counts( - builder: &mut MapBuilder, - counts: &HashMap, -) -> Result<()> { - for (&field_id, &count) in counts { - builder.keys().append_value(field_id); - builder.values().append_value(count.try_into()?); - } - Ok(builder.append(true)?) +/// Returns the field ID of a field in the schema built from +/// `manifest_entry_fields`, which gives every field one. +fn field_id(field: &Field) -> i32 { + field.metadata()[PARQUET_FIELD_ID_META_KEY].parse().unwrap() } -fn append_bounds( - builder: &mut MapBuilder, - bounds: &HashMap, -) -> Result<()> { - for (&field_id, bound) in bounds { - builder.keys().append_value(field_id); - builder.values().append_value(bound.to_bytes()?); +/// Returns the type of a nullable Avro field, `["null", T]`, as `T`. +fn non_null_type(avro_type: &JsonValue) -> &JsonValue { + match avro_type.as_array().map(Vec::as_slice) { + Some([JsonValue::String(null), branch] | [branch, JsonValue::String(null)]) + if null == "null" => + { + branch + } + _ => avro_type, } - Ok(builder.append(true)?) } -/// Builds the `data_file.partition` column, one child array per partition field. -fn partition_column( - entries: &[ManifestEntryRef], - partition_type: &StructType, - field: &Field, -) -> Result { - if let Some(entry) = entries - .iter() - .find(|entry| entry.data_file.partition.fields().len() != partition_type.fields().len()) - { - return Err(invalid_data!( - "partition of {} has {} values, but the partition spec has {} fields", - entry.data_file.file_path, - entry.data_file.partition.fields().len(), - partition_type.fields().len() - )); - } - - let DataType::Struct(fields) = field.data_type() else { - unreachable!() - }; - if fields.is_empty() { - return Ok(Arc::new(StructArray::new_empty_fields(entries.len(), None))); - } - let columns = partition_type - .fields() +/// Builds each of `fields` from the column of `array` that the Avro record +/// `record` gives the same field ID. +fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> Result { + let avro_fields = record["fields"] + .as_array() + .ok_or_else(|| invalid_data!("manifest Avro type {record} is not a record"))?; + let columns = fields .iter() - .zip(fields) - .enumerate() - .map(|(index, (partition_field, arrow_field))| { - let values: Vec<_> = entries + .map(|field| { + let id = field_id(field); + match avro_fields .iter() - .map(|entry| match &entry.data_file.partition[index] { - None => Ok(None), - Some(Literal::Primitive(value)) => Ok(Some(value)), - Some(value) => Err(invalid_data!( - "partition field {} has non-primitive value {value:?}", - partition_field.name - )), - }) - .collect::>()?; - partition_array(partition_field, arrow_field.data_type(), &values) + .position(|avro_field| avro_field["field-id"].as_i64() == Some(id.into())) + { + Some(index) => project( + array.column(index), + non_null_type(&avro_fields[index]["type"]), + field, + ), + None => missing_field(field, array.len()), + } }) .collect::>()?; - Ok(Arc::new(StructArray::try_new( - fields.clone(), - columns, - None, - )?)) + StructArray::try_new_with_length(fields.clone(), columns, array.nulls().cloned(), array.len()) + .map_err(invalid_manifest) } -/// Builds the array of one partition field from its values in each entry. -fn partition_array( - field: &NestedField, - data_type: &DataType, - values: &[Option<&PrimitiveLiteral>], -) -> Result { - let primitive_type = field.field_type.as_primitive_type().ok_or_else(|| { +/// Builds the column for a field that the manifest doesn't have. +fn missing_field(field: &Field, len: usize) -> Result { + // v1 manifests have no content or sequence number columns, which v2+ + // readers read as 0. + match field_id(field) { + CONTENT_FIELD_ID => Ok(Arc::new(Int32Array::from_value(0, len))), + SEQUENCE_NUMBER_FIELD_ID | FILE_SEQUENCE_NUMBER_FIELD_ID => { + Ok(Arc::new(Int64Array::from_value(0, len))) + } + _ if field.is_nullable() => Ok(new_null_array(field.data_type(), len)), + _ if matches!(field.data_type(), DataType::Struct(fields) if fields.is_empty()) => { + Ok(Arc::new(StructArray::new_empty_fields(len, None))) + } + id => Err(invalid_data!( + "manifest has no required field {} with field ID {id}", + field.name() + )), + } +} + +/// Converts the decoded `array` of Avro type `avro_type` to `field`, the +/// spec's Arrow field for it. +fn project(array: &ArrayRef, avro_type: &JsonValue, field: &Field) -> Result { + let type_mismatch = || { invalid_data!( - "partition field {} has non-primitive type {}", - field.name, - field.field_type + "manifest field {} has Avro type {avro_type}, which doesn't match {}", + field.name(), + field.data_type() ) - })?; - // Collects the values, returning an error for a value of another type. - fn collect<'a, T, A: FromIterator>>( - field: &NestedField, - values: &[Option<&'a PrimitiveLiteral>], - value: impl Fn(&'a PrimitiveLiteral) -> Option, - ) -> Result { - values - .iter() - .map(|literal| { - literal - .map(|literal| { - value(literal).ok_or_else(|| { - invalid_data!( - "partition field {} of type {} can't hold {literal:?}", - field.name, - field.field_type - ) - }) - }) - .transpose() - }) - .collect() - } - fn int(literal: &PrimitiveLiteral) -> Option { - match literal { - PrimitiveLiteral::Int(v) => Some(*v), - _ => None, + }; + Ok(match field.data_type() { + DataType::Struct(fields) => Arc::new(project_struct( + array.as_struct_opt().ok_or_else(type_mismatch)?, + avro_type, + fields, + )?), + DataType::List(element) => { + let list = array.as_list_opt::().ok_or_else(type_mismatch)?; + let values = project(list.values(), non_null_type(&avro_type["items"]), element)?; + Arc::new( + ListArray::try_new( + element.clone(), + list.offsets().clone(), + values, + list.nulls().cloned(), + ) + .map_err(invalid_manifest)?, + ) } - } - fn long(literal: &PrimitiveLiteral) -> Option { - match literal { - PrimitiveLiteral::Long(v) => Some(*v), - _ => None, + // Avro stores maps with int keys as arrays of key-value records. + DataType::Map(entries, _) => { + let list = array.as_list_opt::().ok_or_else(type_mismatch)?; + let DataType::Struct(key_value) = entries.data_type() else { + unreachable!() + }; + let entries_array = project_struct( + list.values().as_struct_opt().ok_or_else(type_mismatch)?, + &avro_type["items"], + key_value, + )?; + Arc::new( + MapArray::try_new( + entries.clone(), + list.offsets().clone(), + entries_array, + list.nulls().cloned(), + false, + ) + .map_err(invalid_manifest)?, + ) } - } - fn binary(literal: &PrimitiveLiteral) -> Option<&[u8]> { - match literal { - PrimitiveLiteral::Binary(v) => Some(v), - _ => None, + // arrow-avro gives every timestamp a UTC time zone, so only the type + // changes. + DataType::Timestamp(TimeUnit::Microsecond, time_zone) => Arc::new( + array + .as_primitive_opt::() + .ok_or_else(type_mismatch)? + .clone() + .with_timezone_opt(time_zone.clone()), + ), + DataType::Timestamp(TimeUnit::Nanosecond, time_zone) => Arc::new( + array + .as_primitive_opt::() + .ok_or_else(type_mismatch)? + .clone() + .with_timezone_opt(time_zone.clone()), + ), + data_type if array.data_type() == data_type => array.clone(), + // Without `safe: false`, a value that doesn't fit the spec's type would + // become null instead of an error. + data_type => cast_with_options(array, data_type, &CastOptions { + safe: false, + ..Default::default() + }) + .map_err(invalid_manifest)?, + }) +} + +/// Fills in the values that the spec says entries inherit from the manifest +/// list. +fn inherit(entries: StructArray, manifest_file: &ManifestFile) -> Result { + let (fields, mut columns, nulls) = entries.into_parts(); + let index = |name: &str| fields.find(name).unwrap().0; + let (status, snapshot_id, sequence_number, file_sequence_number, data_file) = ( + index("status"), + index("snapshot_id"), + index("sequence_number"), + index("file_sequence_number"), + index("data_file"), + ); + + columns[snapshot_id] = + fill_nulls(&columns[snapshot_id], None, manifest_file.added_snapshot_id)?; + if columns[sequence_number].null_count() > 0 || columns[file_sequence_number].null_count() > 0 { + let added = eq( + &columns[status], + &Int32Array::new_scalar(ManifestStatus::Added as i32), + )?; + for column in [sequence_number, file_sequence_number] { + columns[column] = fill_nulls( + &columns[column], + Some(&added), + manifest_file.sequence_number, + )?; } } - let array: ArrayRef = - match primitive_type { - PrimitiveType::Boolean => Arc::new(collect::<_, BooleanArray>( - field, - values, - |literal| match literal { - PrimitiveLiteral::Boolean(v) => Some(*v), - _ => None, - }, - )?), - PrimitiveType::Int => Arc::new(collect::<_, Int32Array>(field, values, int)?), - PrimitiveType::Date => Arc::new(collect::<_, Date32Array>(field, values, int)?), - PrimitiveType::Long => Arc::new(collect::<_, Int64Array>(field, values, long)?), - PrimitiveType::Time => { - Arc::new(collect::<_, Time64MicrosecondArray>(field, values, long)?) - } - PrimitiveType::Timestamp => Arc::new(collect::<_, TimestampMicrosecondArray>( - field, values, long, - )?), - PrimitiveType::Timestamptz => Arc::new( - collect::<_, TimestampMicrosecondArray>(field, values, long)? - .with_timezone(UTC_TIME_ZONE), - ), - PrimitiveType::TimestampNs => { - Arc::new(collect::<_, TimestampNanosecondArray>(field, values, long)?) - } - PrimitiveType::TimestamptzNs => Arc::new( - collect::<_, TimestampNanosecondArray>(field, values, long)? - .with_timezone(UTC_TIME_ZONE), - ), - PrimitiveType::Float => Arc::new(collect::<_, Float32Array>( - field, - values, - |literal| match literal { - PrimitiveLiteral::Float(v) => Some(v.0), - _ => None, - }, - )?), - PrimitiveType::Double => Arc::new(collect::<_, Float64Array>( - field, - values, - |literal| match literal { - PrimitiveLiteral::Double(v) => Some(v.0), - _ => None, - }, - )?), - PrimitiveType::Decimal { .. } => { - let DataType::Decimal128(precision, scale) = *data_type else { - unreachable!() - }; - Arc::new( - collect::<_, Decimal128Array>(field, values, |literal| match literal { - PrimitiveLiteral::Int128(v) => Some(*v), - _ => None, - })? - .with_precision_and_scale(precision, scale)?, - ) - } - PrimitiveType::String => Arc::new(collect::<_, StringArray>( - field, - values, - |literal| match literal { - PrimitiveLiteral::String(v) => Some(v.as_str()), - _ => None, - }, - )?), - PrimitiveType::Uuid => { - let uuids: Vec<_> = collect(field, values, |literal| match literal { - PrimitiveLiteral::UInt128(v) => Some(v.to_be_bytes()), - _ => None, - })?; - Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size( - uuids.into_iter(), - 16, - )?) - } - PrimitiveType::Fixed(_) => { - let DataType::FixedSizeBinary(width) = *data_type else { - return Err(Error::new( - ErrorKind::FeatureUnsupported, - format!("fixed partition field {} is {data_type}", field.name), - )); - }; - let bytes: Vec<_> = collect(field, values, binary)?; - Arc::new( - FixedSizeBinaryArray::try_from_sparse_iter_with_size(bytes.into_iter(), width) - .map_err(|err| { - invalid_data!( - "partition field {} has a value of the wrong width", - field.name - ) - .with_source(err) - })?, - ) - } - PrimitiveType::Binary => { - Arc::new(collect::<_, LargeBinaryArray>(field, values, binary)?) - } - PrimitiveType::Unknown => { - // Only nulls fit, so this returns an error for any value. - collect::<(), Vec<_>>(field, values, |_| None)?; - Arc::new(NullArray::new(values.len())) - } - }; - Ok(array) + if manifest_file.content == ManifestContentType::Data { + let (data_file_fields, mut data_file_columns, data_file_nulls) = + columns[data_file].as_struct().clone().into_parts(); + let index = |name: &str| data_file_fields.find(name).unwrap().0; + let first_row_id = index("first_row_id"); + data_file_columns[first_row_id] = inherit_first_row_ids( + data_file_columns[first_row_id].as_primitive(), + columns[status].as_primitive(), + data_file_columns[index("record_count")].as_primitive(), + manifest_file, + )?; + columns[data_file] = Arc::new(StructArray::new( + data_file_fields, + data_file_columns, + data_file_nulls, + )); + } + Ok(StructArray::new(fields, columns, nulls)) +} + +/// Replaces the nulls of `array` with `value`, only in the rows that `rows` +/// selects when given. +fn fill_nulls(array: &ArrayRef, rows: Option<&BooleanArray>, value: i64) -> Result { + if array.null_count() == 0 { + return Ok(array.clone()); + } + let mask = match rows { + Some(rows) => and(&is_null(array)?, rows)?, + None => is_null(array)?, + }; + Ok(zip(&mask, &Int64Array::new_scalar(value), array)?) +} + +/// Assigns first row IDs to the data files that lack one, following first row +/// ID inheritance. As in Java's `ManifestReader`, deleted entries take no row +/// IDs, and a manifest without a first row ID clears them all. +fn inherit_first_row_ids( + first_row_ids: &Int64Array, + status: &Int32Array, + record_counts: &Int64Array, + manifest_file: &ManifestFile, +) -> Result { + let Some(manifest_first_row_id) = manifest_file.first_row_id else { + return Ok(new_null_array(&DataType::Int64, first_row_ids.len())); + }; + if first_row_ids.null_count() == 0 { + return Ok(Arc::new(first_row_ids.clone())); + } + let mut next_row_id = i64::try_from(manifest_first_row_id)?; + let first_row_ids: Int64Array = first_row_ids + .iter() + .zip(status.values()) + .zip(record_counts.values()) + .map( + |((first_row_id, &status), &record_count)| match first_row_id { + None if status != ManifestStatus::Deleted as i32 => { + let assigned = next_row_id; + next_row_id = next_row_id.checked_add(record_count).ok_or_else(|| { + invalid_data!( + "row ID overflow assigning first row IDs in {}", + manifest_file.manifest_path + ) + })?; + Ok(Some(assigned)) + } + first_row_id => Ok(first_row_id), + }, + ) + .collect::>()?; + Ok(Arc::new(first_row_ids)) } #[cfg(test)] @@ -494,28 +429,26 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; - use apache_avro::Reader as AvroReader; use apache_avro::types::Value as AvroValue; use arrow_array::cast::AsArray; use arrow_array::types::{Int32Type, Int64Type}; use arrow_array::{ Array, ArrayRef, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray, - Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray, NullArray, - StringArray, StructArray, Time64MicrosecondArray, TimestampMicrosecondArray, - TimestampNanosecondArray, + Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray, StringArray, + Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, }; - use uuid::Uuid; use super::*; + use crate::ErrorKind; use crate::arrow::arrow_schema_to_schema; + use crate::io::FileIO; use crate::spec::{ - DataContentType, DataFile, DataFileBuilder, DataFileFormat, Datum, FormatVersion, ListType, - Literal, ManifestContentType, ManifestEntry, ManifestMetadata, ManifestStatus, MapType, - NestedField, PartitionSpec, PrimitiveType, Schema, Struct, StructType, Transform, Type, + DataContentType, DataFile, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, + Manifest, ManifestEntry, ManifestStatus, ManifestWriter, ManifestWriterBuilder, MapType, + NestedField, PartitionSpec, PrimitiveLiteral, PrimitiveType, Schema, Struct, StructType, + Transform, Type, }; - const UUID: &str = "f79c3e09-677c-4bbd-a479-3f349cb785e7"; - fn table_schema() -> Schema { let columns = [ (1, "b", PrimitiveType::Boolean), @@ -534,10 +467,8 @@ mod tests { (11, "ts_ns", PrimitiveType::TimestampNs), (12, "tstz_ns", PrimitiveType::TimestamptzNs), (13, "s", PrimitiveType::String), - (14, "u", PrimitiveType::Uuid), - (15, "fx", PrimitiveType::Fixed(3)), - (16, "bin", PrimitiveType::Binary), - (17, "unk", PrimitiveType::Unknown), + (14, "fx", PrimitiveType::Fixed(3)), + (15, "bin", PrimitiveType::Binary), ]; Schema::builder() .with_fields( @@ -564,19 +495,6 @@ mod tests { .unwrap() } - fn manifest(partition_spec: PartitionSpec, entries: Vec) -> Manifest { - Manifest::new( - ManifestMetadata { - schema: Arc::new(table_schema()), - schema_id: 0, - partition_spec, - format_version: FormatVersion::V2, - content: ManifestContentType::Data, - }, - entries, - ) - } - fn data_file(content: DataContentType, path: &str) -> DataFileBuilder { let mut builder = DataFileBuilder::default(); builder @@ -588,20 +506,40 @@ mod tests { builder } - fn entry(status: ManifestStatus, data_file: DataFile) -> ManifestEntry { - ManifestEntry { - status, - snapshot_id: None, - sequence_number: None, - file_sequence_number: None, - data_file, - } + /// Writes a manifest with `ManifestWriter` and returns its bytes and its + /// manifest list entry. + async fn write_manifest( + spec: PartitionSpec, + snapshot_id: Option, + build: impl FnOnce(ManifestWriterBuilder) -> ManifestWriter, + add: impl FnOnce(&mut ManifestWriter), + ) -> (Vec, ManifestFile) { + let io = FileIO::new_with_memory(); + let output = io.new_output("memory:///manifest.avro").unwrap(); + let mut writer = build(ManifestWriterBuilder::new( + output, + snapshot_id, + Arc::new(table_schema()), + spec, + )); + add(&mut writer); + let manifest_file = writer.write_manifest_file().await.unwrap(); + let bytes = io + .new_input("memory:///manifest.avro") + .unwrap() + .read() + .await + .unwrap() + .to_vec(); + (bytes, manifest_file) + } + + fn column(batch: &RecordBatch, name: &str) -> ArrayRef { + batch.column_by_name(name).unwrap().clone() } fn data_file_column(batch: &RecordBatch, name: &str) -> ArrayRef { - batch - .column_by_name("data_file") - .unwrap() + column(batch, "data_file") .as_struct() .column_by_name(name) .unwrap() @@ -651,10 +589,80 @@ mod tests { )) } - #[test] - fn test_schema_is_spec_manifest_entry() { + /// Rewrites a manifest with `edit_schema` applied to its Avro schema and + /// `edit_record` applied to each record. + fn rewrite_manifest( + bytes: &[u8], + edit_schema: impl Fn(&mut serde_json::Value), + edit_record: impl Fn(&mut AvroValue), + ) -> Vec { + let reader = apache_avro::Reader::new(bytes).unwrap(); + let mut schema = serde_json::to_value(reader.writer_schema()).unwrap(); + edit_schema(&mut schema); + let schema = apache_avro::Schema::parse(&schema).unwrap(); + let metadata = reader.user_metadata().clone(); + let mut writer = apache_avro::Writer::new(&schema, Vec::new()).unwrap(); + for (key, value) in metadata { + writer.add_user_metadata(key, value).unwrap(); + } + for record in reader { + let mut record = record.unwrap(); + edit_record(&mut record); + writer.append_value(record).unwrap(); + } + writer.into_inner().unwrap() + } + + /// Calls `edit` on every record field object in an Avro schema. + fn for_each_avro_field(schema: &mut serde_json::Value, edit: &impl Fn(&mut serde_json::Value)) { + match schema { + serde_json::Value::Object(object) => { + if let Some(serde_json::Value::Array(fields)) = object.get_mut("fields") { + for field in fields { + edit(field); + for_each_avro_field(field.get_mut("type").unwrap(), edit); + } + } + if let Some(items) = object.get_mut("items") { + for_each_avro_field(items, edit); + } + } + serde_json::Value::Array(branches) => { + for branch in branches { + for_each_avro_field(branch, edit); + } + } + _ => {} + } + } + + /// Renames the record field `from` to `to` in every Avro record value. + fn rename_avro_field(value: &mut AvroValue, from: &str, to: &str) { + match value { + AvroValue::Record(fields) => { + for (name, field) in fields { + if name == from { + *name = to.to_string(); + } + rename_avro_field(field, from, to); + } + } + AvroValue::Union(_, value) => rename_avro_field(value, from, to), + AvroValue::Array(items) => { + for item in items { + rename_avro_field(item, from, to); + } + } + _ => {} + } + } + + #[tokio::test] + async fn test_schema_is_spec_manifest_entry() { let spec = partition_spec(&[("i", Transform::Identity), ("s", Transform::Truncate(2))]); - let batch = manifest_to_record_batch(&manifest(spec, vec![])).unwrap(); + let (bytes, manifest_file) = + write_manifest(spec, Some(1), |b| b.build_v2_data(), |_| {}).await; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); assert_eq!(batch.num_rows(), 0); let long = || Type::Primitive(PrimitiveType::Long); @@ -737,204 +745,167 @@ mod tests { ); } - #[test] - fn test_entry_and_data_file_columns() { + #[tokio::test] + async fn test_data_file_columns() { let added = DataFile { column_sizes: HashMap::from([(1, 11)]), - value_counts: HashMap::from([(1, 10)]), + value_counts: HashMap::from([(1, 10), (3, 9)]), null_value_counts: HashMap::from([(1, 1)]), nan_value_counts: HashMap::from([(4, 0)]), - lower_bounds: HashMap::from([(2, Datum::int(-1))]), + lower_bounds: HashMap::from([(2, Datum::int(-1)), (13, Datum::string("abc"))]), upper_bounds: HashMap::from([(2, Datum::int(7))]), key_metadata: Some(vec![1, 2]), split_offsets: Some(vec![4, 100]), sort_order_id: Some(0), - first_row_id: Some(1000), ..data_file(DataContentType::Data, "s3://b/data/0.parquet") .record_count(10) .file_size_in_bytes(100) .build() .unwrap() }; - let position_deletes = DataFile { - file_format: DataFileFormat::Puffin, - referenced_data_file: Some("s3://b/data/0.parquet".to_string()), - content_offset: Some(4), - content_size_in_bytes: Some(40), - ..data_file(DataContentType::PositionDeletes, "s3://b/data/1.puffin") - .build() - .unwrap() - }; - let equality_deletes = DataFile { - equality_ids: Some(vec![1, 3]), - split_offsets: Some(vec![]), - ..data_file(DataContentType::EqualityDeletes, "s3://b/data/2.parquet") - .build() - .unwrap() - }; - let entries = vec![ - ManifestEntry { - snapshot_id: Some(1), - sequence_number: Some(2), - file_sequence_number: Some(3), - ..entry(ManifestStatus::Added, added) - }, - entry(ManifestStatus::Existing, position_deletes), - ManifestEntry { - snapshot_id: Some(5), - ..entry(ManifestStatus::Deleted, equality_deletes) + let existing = data_file(DataContentType::Data, "s3://b/data/1.parquet") + .build() + .unwrap(); + let (bytes, manifest_file) = write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v2_data(), + |writer| { + writer.add_file(added, 2).unwrap(); + writer.add_existing_file(existing, 5, 3, Some(3)).unwrap(); }, - ]; - let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), entries)).unwrap(); - assert_eq!(batch.num_rows(), 3); + ) + .await; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + assert_eq!(batch.num_rows(), 2); - let column = |name: &str| batch.column_by_name(name).unwrap().clone(); assert_eq!( - column("status").as_ref(), - &Int32Array::from(vec![1, 0, 2]) as &dyn Array + column(&batch, "status").as_ref(), + &Int32Array::from(vec![1, 0]) as &dyn Array ); assert_eq!( - column("snapshot_id").as_ref(), - &Int64Array::from(vec![Some(1), None, Some(5)]) as &dyn Array + column(&batch, "snapshot_id").as_ref(), + &Int64Array::from(vec![1, 5]) as &dyn Array ); assert_eq!( - column("sequence_number").as_ref(), - &Int64Array::from(vec![Some(2), None, None]) as &dyn Array - ); - assert_eq!( - column("file_sequence_number").as_ref(), - &Int64Array::from(vec![Some(3), None, None]) as &dyn Array + column(&batch, "sequence_number").as_ref(), + &Int64Array::from(vec![2, 3]) as &dyn Array ); let column = |name: &str| data_file_column(&batch, name); assert_eq!( column("content").as_ref(), - &Int32Array::from(vec![0, 1, 2]) as &dyn Array + &Int32Array::from(vec![0, 0]) as &dyn Array ); assert_eq!( column("file_path").as_ref(), - &StringArray::from(vec![ - "s3://b/data/0.parquet", - "s3://b/data/1.puffin", - "s3://b/data/2.parquet" - ]) as &dyn Array + &StringArray::from(vec!["s3://b/data/0.parquet", "s3://b/data/1.parquet"]) + as &dyn Array ); - // Writers, including `ManifestWriter`, store the format name in upper case. assert_eq!( column("file_format").as_ref(), - &StringArray::from(vec!["PARQUET", "PUFFIN", "PARQUET"]) as &dyn Array + &StringArray::from(vec!["PARQUET", "PARQUET"]) as &dyn Array ); assert_eq!( column("record_count").as_ref(), - &Int64Array::from(vec![10, 1, 1]) as &dyn Array + &Int64Array::from(vec![10, 1]) as &dyn Array ); assert_eq!( column("file_size_in_bytes").as_ref(), - &Int64Array::from(vec![100, 10, 10]) as &dyn Array + &Int64Array::from(vec![100, 10]) as &dyn Array ); - for (name, expected) in [ - ("column_sizes", (1, 11)), - ("value_counts", (1, 10)), - ("null_value_counts", (1, 1)), - ("nan_value_counts", (4, 0)), - ] { - let map = column(name); - assert_eq!(map.null_count(), 0, "{name}"); - assert_eq!(long_map(&map, 0), vec![expected], "{name}"); - assert_eq!(long_map(&map, 1), vec![], "{name}"); - assert_eq!(long_map(&map, 2), vec![], "{name}"); - } - assert_eq!(binary_map(&column("lower_bounds"), 0), vec![( - 2, - (-1i32).to_le_bytes().to_vec() - )]); + assert_eq!(long_map(&column("column_sizes"), 0), vec![(1, 11)]); + assert_eq!(long_map(&column("value_counts"), 0), vec![(1, 10), (3, 9)]); + assert_eq!(long_map(&column("null_value_counts"), 0), vec![(1, 1)]); + assert_eq!(long_map(&column("nan_value_counts"), 0), vec![(4, 0)]); + assert_eq!(long_map(&column("value_counts"), 1), vec![]); + assert_eq!(binary_map(&column("lower_bounds"), 0), vec![ + (2, (-1i32).to_le_bytes().to_vec()), + (13, b"abc".to_vec()), + ]); assert_eq!(binary_map(&column("upper_bounds"), 0), vec![( 2, 7i32.to_le_bytes().to_vec() )]); - assert_eq!(binary_map(&column("lower_bounds"), 1), vec![]); assert_eq!( column("key_metadata").as_ref(), - &LargeBinaryArray::from(vec![Some(&[1u8, 2][..]), None, None]) as &dyn Array + &LargeBinaryArray::from(vec![Some(&[1u8, 2][..]), None]) as &dyn Array ); let split_offsets = column("split_offsets"); - let split_offsets = split_offsets.as_list::(); - assert!(split_offsets.is_valid(0) && split_offsets.is_null(1) && split_offsets.is_valid(2)); assert_eq!( - split_offsets.value(0).as_ref(), + split_offsets.as_list::().value(0).as_ref(), &Int64Array::from(vec![4, 100]) as &dyn Array ); - assert_eq!(split_offsets.value(2).len(), 0); - let equality_ids = column("equality_ids"); - let equality_ids = equality_ids.as_list::(); - assert!(equality_ids.is_null(0) && equality_ids.is_null(1) && equality_ids.is_valid(2)); assert_eq!( - equality_ids.value(2).as_ref(), - &Int32Array::from(vec![1, 3]) as &dyn Array + column("sort_order_id").as_ref(), + &Int32Array::from(vec![Some(0), None]) as &dyn Array ); + } + + #[tokio::test] + async fn test_delete_file_columns() { + let position_deletes = DataFile { + file_format: DataFileFormat::Puffin, + referenced_data_file: Some("s3://b/data/0.parquet".to_string()), + content_offset: Some(4), + content_size_in_bytes: Some(40), + ..data_file(DataContentType::PositionDeletes, "s3://b/data/1.puffin") + .build() + .unwrap() + }; + let equality_deletes = DataFile { + equality_ids: Some(vec![1, 3]), + ..data_file(DataContentType::EqualityDeletes, "s3://b/data/2.parquet") + .build() + .unwrap() + }; + let (bytes, manifest_file) = write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v3_deletes(), + |writer| { + writer.add_file(position_deletes, 2).unwrap(); + writer.add_file(equality_deletes, 2).unwrap(); + }, + ) + .await; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + + let column = |name: &str| data_file_column(&batch, name); assert_eq!( - column("sort_order_id").as_ref(), - &Int32Array::from(vec![Some(0), None, None]) as &dyn Array + column("content").as_ref(), + &Int32Array::from(vec![1, 2]) as &dyn Array ); assert_eq!( - column("first_row_id").as_ref(), - &Int64Array::from(vec![Some(1000), None, None]) as &dyn Array + column("file_format").as_ref(), + &StringArray::from(vec!["PUFFIN", "PARQUET"]) as &dyn Array ); assert_eq!( column("referenced_data_file").as_ref(), - &StringArray::from(vec![None, Some("s3://b/data/0.parquet"), None]) as &dyn Array + &StringArray::from(vec![Some("s3://b/data/0.parquet"), None]) as &dyn Array ); assert_eq!( column("content_offset").as_ref(), - &Int64Array::from(vec![None, Some(4), None]) as &dyn Array + &Int64Array::from(vec![Some(4), None]) as &dyn Array ); assert_eq!( column("content_size_in_bytes").as_ref(), - &Int64Array::from(vec![None, Some(40), None]) as &dyn Array - ); - } - - #[test] - fn test_metric_maps_with_several_columns() { - let uuid = Uuid::parse_str(UUID).unwrap(); - let file = DataFile { - value_counts: HashMap::from([(1, 5), (3, 6), (13, 7)]), - lower_bounds: HashMap::from([ - (1, Datum::bool(false)), - (3, Datum::long(-2)), - (13, Datum::string("abc")), - (14, Datum::uuid(uuid)), - ]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") - .build() - .unwrap() - }; - let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), vec![entry( - ManifestStatus::Added, - file, - )])) - .unwrap(); - - assert_eq!( - long_map(&data_file_column(&batch, "value_counts"), 0), - vec![(1, 5), (3, 6), (13, 7)] + &Int64Array::from(vec![Some(40), None]) as &dyn Array ); + let equality_ids = column("equality_ids"); + let equality_ids = equality_ids.as_list::(); + assert!(equality_ids.is_null(0)); assert_eq!( - binary_map(&data_file_column(&batch, "lower_bounds"), 0), - vec![ - (1, vec![0]), - (3, (-2i64).to_le_bytes().to_vec()), - (13, b"abc".to_vec()), - (14, uuid.as_bytes().to_vec()), - ] + equality_ids.value(1).as_ref(), + &Int32Array::from(vec![1, 3]) as &dyn Array ); } - #[test] - fn test_partition_values() { + #[tokio::test] + async fn test_partition_values() { let identity_columns = [ "b", "i", "l", "f", "d", "dec", "date", "time", "ts", "tstz", "ts_ns", "tstz_ns", "s", - "u", "fx", "bin", + "fx", "bin", ]; let mut fields: Vec<_> = identity_columns .iter() @@ -948,7 +919,6 @@ mod tests { ]); let spec = partition_spec(&fields); - let uuid = Uuid::parse_str(UUID).unwrap(); let values = Struct::from_iter([ Some(Literal::bool(true)), Some(Literal::int(-7)), @@ -963,7 +933,6 @@ mod tests { Some(Literal::timestamp_nano(1_700_000_000_000_000_002)), Some(Literal::timestamptz_nano(1_700_000_000_000_000_003)), Some(Literal::string("abc")), - Some(Literal::uuid(uuid)), Some(Literal::fixed([1u8, 2, 3])), Some(Literal::binary([9u8])), Some(Literal::int(3)), @@ -972,22 +941,27 @@ mod tests { None, ]); let nulls = Struct::from_iter(fields.iter().map(|_| None)); - let entries = [values, nulls] - .into_iter() - .map(|partition| { - entry(ManifestStatus::Added, DataFile { - partition, - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") - .build() - .unwrap() - }) - }) - .collect(); - let batch = manifest_to_record_batch(&manifest(spec, entries)).unwrap(); + let (bytes, manifest_file) = write_manifest( + spec, + Some(1), + |b| b.build_v2_data(), + |writer| { + for (i, partition) in [values, nulls].into_iter().enumerate() { + let file = DataFile { + partition, + ..data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) + .build() + .unwrap() + }; + writer.add_file(file, 1).unwrap(); + } + }, + ) + .await; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); let partition = data_file_column(&batch, "partition"); let partition = partition.as_struct(); - assert_eq!(partition.null_count(), 0); let expected: Vec = vec![ Arc::new(BooleanArray::from(vec![Some(true), None])), Arc::new(Int32Array::from(vec![Some(-7), None])), @@ -1021,13 +995,6 @@ mod tests { .with_timezone("+00:00"), ), Arc::new(StringArray::from(vec![Some("abc"), None])), - Arc::new( - FixedSizeBinaryArray::try_from_sparse_iter_with_size( - [Some(uuid.as_bytes().to_vec()), None].into_iter(), - 16, - ) - .unwrap(), - ), Arc::new( FixedSizeBinaryArray::try_from_sparse_iter_with_size( [Some(vec![1u8, 2, 3]), None].into_iter(), @@ -1052,215 +1019,254 @@ mod tests { } } - #[test] - fn test_unpartitioned_partition_column() { - let entries = (0..3) - .map(|i| { - entry( - ManifestStatus::Added, - data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) - .build() - .unwrap(), - ) - }) - .collect(); - let batch = manifest_to_record_batch(&manifest(partition_spec(&[]), entries)).unwrap(); - let partition = data_file_column(&batch, "partition"); - let partition: &StructArray = partition.as_struct(); - assert_eq!(partition.num_columns(), 0); - assert_eq!(partition.len(), 3); - assert_eq!(partition.null_count(), 0); - } + #[tokio::test] + async fn test_v1_manifest_reads_with_v2_defaults() { + let file = data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap(); + let (bytes, mut manifest_file) = write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v1(), + |writer| writer.add_file(file, -1).unwrap(), + ) + .await; + // The v1 default must win over inheritance, so give the manifest list + // entry a sequence number that inheritance would use. + manifest_file.sequence_number = 9; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); - #[test] - fn test_unknown_partition_field() { - let spec = || partition_spec(&[("unk", Transform::Void)]); - let entries = vec![entry(ManifestStatus::Added, DataFile { - partition: Struct::from_iter([None]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") - .build() - .unwrap() - })]; - let batch = manifest_to_record_batch(&manifest(spec(), entries)).unwrap(); - let partition = data_file_column(&batch, "partition"); assert_eq!( - partition.as_struct().column(0).as_ref(), - &NullArray::new(1) as &dyn Array + data_file_column(&batch, "content").as_ref(), + &Int32Array::from(vec![0]) as &dyn Array + ); + for name in ["sequence_number", "file_sequence_number"] { + assert_eq!( + column(&batch, name).as_ref(), + &Int64Array::from(vec![0]) as &dyn Array, + "{name}" + ); + } + assert_eq!( + column(&batch, "snapshot_id").as_ref(), + &Int64Array::from(vec![1]) as &dyn Array ); - - let entries = vec![entry(ManifestStatus::Added, DataFile { - partition: Struct::from_iter([Some(Literal::int(1))]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") - .build() - .unwrap() - })]; - let err = manifest_to_record_batch(&manifest(spec(), entries)).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::DataInvalid); } - #[test] - fn test_partition_value_of_wrong_type_is_an_error() { - let file = DataFile { - partition: Struct::from_iter([Some(Literal::string("not an int"))]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + #[tokio::test] + async fn test_snapshot_id_and_sequence_number_inheritance() { + let file = |i: usize| { + data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) .build() .unwrap() }; - let manifest = manifest(partition_spec(&[("i", Transform::Identity)]), vec![entry( - ManifestStatus::Added, - file, - )]); - let err = manifest_to_record_batch(&manifest).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::DataInvalid); + let (bytes, mut manifest_file) = write_manifest( + partition_spec(&[]), + None, + |b| b.build_v2_data(), + |writer| { + writer.add_file(file(0), -1).unwrap(); + writer.add_existing_file(file(1), 3, 4, Some(4)).unwrap(); + }, + ) + .await; + manifest_file.added_snapshot_id = 42; + manifest_file.sequence_number = 7; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + + assert_eq!( + column(&batch, "snapshot_id").as_ref(), + &Int64Array::from(vec![42, 3]) as &dyn Array + ); + for name in ["sequence_number", "file_sequence_number"] { + assert_eq!( + column(&batch, name).as_ref(), + &Int64Array::from(vec![7, 4]) as &dyn Array, + "{name}" + ); + } } - #[test] - fn test_fixed_partition_value_of_wrong_width_is_an_error() { - let file = DataFile { - partition: Struct::from_iter([Some(Literal::fixed([1u8, 2]))]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") + async fn first_row_id_manifest() -> (Vec, ManifestFile) { + let file = |i: usize, record_count: u64, first_row_id: Option| DataFile { + first_row_id, + ..data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) + .record_count(record_count) .build() .unwrap() }; - let manifest = manifest(partition_spec(&[("fx", Transform::Identity)]), vec![entry( - ManifestStatus::Added, - file, - )]); - let err = manifest_to_record_batch(&manifest).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::DataInvalid); + let entry = |status, data_file| ManifestEntry { + status, + snapshot_id: Some(1), + sequence_number: Some(1), + file_sequence_number: Some(1), + data_file, + }; + write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v3_data(), + |writer| { + writer.add_file(file(0, 10, None), 1).unwrap(); + writer + .add_existing_entry(entry(ManifestStatus::Existing, file(1, 3, Some(5)))) + .unwrap(); + writer + .add_delete_entry(entry(ManifestStatus::Deleted, file(2, 4, None))) + .unwrap(); + writer + .add_existing_entry(entry(ManifestStatus::Existing, file(3, 6, None))) + .unwrap(); + writer.add_file(file(4, 2, None), 1).unwrap(); + }, + ) + .await } - #[test] - fn test_partition_value_with_wrong_arity_is_an_error() { - let file = DataFile { - partition: Struct::from_iter([Some(Literal::int(1)), Some(Literal::int(2))]), - ..data_file(DataContentType::Data, "s3://b/data/0.parquet") - .build() - .unwrap() - }; - let manifest = manifest(partition_spec(&[("i", Transform::Identity)]), vec![entry( - ManifestStatus::Added, - file, - )]); - let err = manifest_to_record_batch(&manifest).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::DataInvalid); + #[tokio::test] + async fn test_first_row_id_inheritance() { + let (bytes, mut manifest_file) = first_row_id_manifest().await; + manifest_file.first_row_id = Some(100); + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + // Deleted entries don't take row IDs, and entries with a first_row_id keep it. + assert_eq!( + data_file_column(&batch, "first_row_id").as_ref(), + &Int64Array::from(vec![Some(100), Some(5), None, Some(110), Some(116)]) as &dyn Array + ); + + manifest_file.first_row_id = None; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + assert_eq!(data_file_column(&batch, "first_row_id").null_count(), 5); } - #[test] - fn test_every_format_version_has_the_same_schema() { - let schemas: Vec<_> = [FormatVersion::V1, FormatVersion::V2, FormatVersion::V3] - .into_iter() - .map(|format_version| { - let manifest = Manifest::new( - ManifestMetadata { - format_version, - ..manifest(partition_spec(&[("i", Transform::Identity)]), vec![]) - .metadata() - .clone() - }, - vec![], - ); - manifest_to_record_batch(&manifest).unwrap().schema() - }) - .collect(); - assert_eq!(schemas[0], schemas[1]); - assert_eq!(schemas[1], schemas[2]); + #[tokio::test] + async fn test_first_row_id_overflow_is_an_error() { + let (bytes, mut manifest_file) = first_row_id_manifest().await; + manifest_file.first_row_id = Some(i64::MAX as u64 - 5); + let err = read_manifest_entries(&bytes, &manifest_file).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); } - #[test] - fn test_count_beyond_long_range_is_an_error() { + #[tokio::test] + async fn test_fields_are_matched_by_field_id() { let file = DataFile { - column_sizes: HashMap::from([(1, u64::MAX)]), + partition: Struct::from_iter([Some(Literal::int(7))]), ..data_file(DataContentType::Data, "s3://b/data/0.parquet") .build() .unwrap() }; - let manifest = manifest(partition_spec(&[]), vec![entry( - ManifestStatus::Added, - file, - )]); - let err = manifest_to_record_batch(&manifest).unwrap_err(); - assert_eq!(err.kind(), ErrorKind::DataInvalid); + let (bytes, manifest_file) = write_manifest( + partition_spec(&[("i", Transform::Identity)]), + Some(1), + |b| b.build_v2_data(), + |writer| writer.add_file(file, 1).unwrap(), + ) + .await; + let renames = [("file_path", "path"), ("i_identity", "renamed_partition")]; + let bytes = rewrite_manifest( + &bytes, + |schema| { + for_each_avro_field(schema, &|field| { + for (from, to) in renames { + if field["name"] == from { + field["name"] = to.into(); + } + } + }) + }, + |record| { + for (from, to) in renames { + rename_avro_field(record, from, to); + } + }, + ); + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + + assert_eq!( + data_file_column(&batch, "file_path").as_ref(), + &StringArray::from(vec!["s3://b/data/0.parquet"]) as &dyn Array + ); + let partition = data_file_column(&batch, "partition"); + assert_eq!( + partition.as_struct().column(0).as_ref(), + &Int32Array::from(vec![7]) as &dyn Array + ); } - /// `(key, value)` pairs of the map field `name` of an Avro `data_file` - /// record, sorted by key. - fn avro_map( - data_file: &[(String, AvroValue)], - name: &str, - value: impl Fn(&AvroValue) -> T, - ) -> Vec<(i32, T)> { - let (_, AvroValue::Union(_, map)) = - data_file.iter().find(|(field, _)| field == name).unwrap() - else { - panic!("{name} is not a union"); - }; - let AvroValue::Array(pairs) = map.as_ref() else { - panic!("{name} is not an array"); - }; - let mut pairs: Vec<_> = pairs - .iter() - .map(|pair| { - let AvroValue::Record(pair) = pair else { - panic!("{name} entry is not a record"); - }; - let AvroValue::Int(key) = pair[0].1 else { - panic!("{name} key is not an int"); - }; - (key, value(&pair[1].1)) - }) - .collect(); - pairs.sort_unstable(); - pairs + #[tokio::test] + async fn test_missing_required_field_is_an_error() { + let file = data_file(DataContentType::Data, "s3://b/data/0.parquet") + .build() + .unwrap(); + let (bytes, manifest_file) = write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v2_data(), + |writer| writer.add_file(file, 1).unwrap(), + ) + .await; + // A file_path under another field ID leaves the required file_path missing. + let bytes = rewrite_manifest( + &bytes, + |schema| { + for_each_avro_field(schema, &|field| { + if field["name"] == "file_path" { + field["field-id"] = 999.into(); + } + }) + }, + |_| {}, + ); + let err = read_manifest_entries(&bytes, &manifest_file).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); } - #[test] - fn test_manifest_written_by_pyiceberg() { - let bs = std::fs::read(format!( + #[tokio::test] + async fn test_manifest_written_by_pyiceberg() { + let bytes = std::fs::read(format!( "{}/testdata/manifests/pyiceberg-v2-data.avro", env!("CARGO_MANIFEST_DIR") )) .unwrap(); - let batch = manifest_to_record_batch(&Manifest::parse_avro(&bs).unwrap()).unwrap(); + let manifest = Manifest::parse_avro(&bytes).unwrap(); + let (_, manifest_file) = + write_manifest(partition_spec(&[]), Some(1), |b| b.build_v2_data(), |_| {}).await; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); - let records: Vec<_> = AvroReader::new(bs.as_slice()) - .unwrap() - .map(|record| match record.unwrap() { - AvroValue::Record(fields) => fields, - _ => panic!("manifest entry is not a record"), - }) - .collect(); - assert_eq!(batch.num_rows(), records.len()); - for (row, record) in records.iter().enumerate() { - let Some((_, AvroValue::Record(data_file))) = - record.iter().find(|(name, _)| name == "data_file") + assert_eq!(batch.num_rows(), manifest.entries().len()); + let paths = data_file_column(&batch, "file_path"); + let partition = data_file_column(&batch, "partition"); + let category = partition.as_struct().column(0).as_string::().clone(); + for (row, entry) in manifest.entries().iter().enumerate() { + let data_file = entry.data_file(); + assert_eq!(paths.as_string::().value(row), data_file.file_path()); + let Some(Literal::Primitive(PrimitiveLiteral::String(expected))) = + &data_file.partition()[0] else { - panic!("data_file is not a record"); - }; - let long = |value: &AvroValue| match value { - AvroValue::Long(v) => *v, - _ => panic!("map value is not a long"), - }; - let bytes = |value: &AvroValue| match value { - AvroValue::Bytes(v) => v.clone(), - _ => panic!("map value is not bytes"), + panic!("partition value is not a string"); }; - for name in ["column_sizes", "value_counts", "null_value_counts"] { - assert_eq!( - long_map(&data_file_column(&batch, name), row), - avro_map(data_file, name, long), - "{name}" - ); - } - for name in ["lower_bounds", "upper_bounds"] { - let expected = avro_map(data_file, name, bytes); - assert!(!expected.is_empty(), "{name}"); - assert_eq!( - binary_map(&data_file_column(&batch, name), row), - expected, - "{name}" - ); - } + assert_eq!(category.value(row), expected); + let mut lower_bounds: Vec<_> = data_file + .lower_bounds() + .iter() + .map(|(id, bound)| (*id, bound.to_bytes().unwrap().to_vec())) + .collect(); + lower_bounds.sort_unstable(); + assert!(!lower_bounds.is_empty()); + assert_eq!( + binary_map(&data_file_column(&batch, "lower_bounds"), row), + lower_bounds + ); + let mut value_counts: Vec<_> = data_file + .value_counts() + .iter() + .map(|(id, count)| (*id, *count as i64)) + .collect(); + value_counts.sort_unstable(); + assert_eq!( + long_map(&data_file_column(&batch, "value_counts"), row), + value_counts + ); } } } From 9c6325ff83a428b2f4d42adc85e12595346d074d Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 12 Oct 2026 11:13:22 -0400 Subject: [PATCH 5/7] address PR feedback --- crates/iceberg/src/arrow/manifest.rs | 145 +++++++++++++++++---------- 1 file changed, 94 insertions(+), 51 deletions(-) diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index 15879870d4..21c5658d87 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -21,6 +21,7 @@ use std::collections::HashMap; use std::sync::Arc; use arrow_arith::boolean::{and, is_null}; +use arrow_array::builder::Int64Builder; use arrow_array::cast::AsArray; use arrow_array::types::{TimestampMicrosecondType, TimestampNanosecondType}; use arrow_array::{ @@ -205,10 +206,10 @@ fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> R let avro_fields = record["fields"] .as_array() .ok_or_else(|| invalid_data!("manifest Avro type {record} is not a record"))?; - let columns = fields - .iter() - .map(|field| { - let id = field_id(field); + let mut columns = Vec::with_capacity(fields.len()); + for field in fields { + let id = field_id(field); + columns.push( match avro_fields .iter() .position(|avro_field| avro_field["field-id"].as_i64() == Some(id.into())) @@ -217,11 +218,11 @@ fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> R array.column(index), non_null_type(&avro_fields[index]["type"]), field, - ), - None => missing_field(field, array.len()), - } - }) - .collect::>()?; + )?, + None => missing_field(field, array.len())?, + }, + ); + } StructArray::try_new_with_length(fields.clone(), columns, array.nulls().cloned(), array.len()) .map_err(invalid_manifest) } @@ -401,27 +402,26 @@ fn inherit_first_row_ids( return Ok(Arc::new(first_row_ids.clone())); } let mut next_row_id = i64::try_from(manifest_first_row_id)?; - let first_row_ids: Int64Array = first_row_ids + let mut builder = Int64Builder::with_capacity(first_row_ids.len()); + for ((first_row_id, &status), &record_count) in first_row_ids .iter() .zip(status.values()) .zip(record_counts.values()) - .map( - |((first_row_id, &status), &record_count)| match first_row_id { - None if status != ManifestStatus::Deleted as i32 => { - let assigned = next_row_id; - next_row_id = next_row_id.checked_add(record_count).ok_or_else(|| { - invalid_data!( - "row ID overflow assigning first row IDs in {}", - manifest_file.manifest_path - ) - })?; - Ok(Some(assigned)) - } - first_row_id => Ok(first_row_id), - }, - ) - .collect::>()?; - Ok(Arc::new(first_row_ids)) + { + match first_row_id { + None if status != ManifestStatus::Deleted as i32 => { + builder.append_value(next_row_id); + next_row_id = next_row_id.checked_add(record_count).ok_or_else(|| { + invalid_data!( + "row ID overflow assigning first row IDs in {}", + manifest_file.manifest_path + ) + })?; + } + first_row_id => builder.append_option(first_row_id), + } + } + Ok(Arc::new(builder.finish())) } #[cfg(test)] @@ -430,6 +430,7 @@ mod tests { use std::sync::Arc; use apache_avro::types::Value as AvroValue; + use apache_avro::{Codec, DeflateSettings, ZstandardSettings}; use arrow_array::cast::AsArray; use arrow_array::types::{Int32Type, Int64Type}; use arrow_array::{ @@ -589,10 +590,11 @@ mod tests { )) } - /// Rewrites a manifest with `edit_schema` applied to its Avro schema and - /// `edit_record` applied to each record. + /// Rewrites a manifest with `codec`, `edit_schema` applied to its Avro + /// schema, and `edit_record` applied to each record. fn rewrite_manifest( bytes: &[u8], + codec: Codec, edit_schema: impl Fn(&mut serde_json::Value), edit_record: impl Fn(&mut AvroValue), ) -> Vec { @@ -601,7 +603,7 @@ mod tests { edit_schema(&mut schema); let schema = apache_avro::Schema::parse(&schema).unwrap(); let metadata = reader.user_metadata().clone(); - let mut writer = apache_avro::Writer::new(&schema, Vec::new()).unwrap(); + let mut writer = apache_avro::Writer::with_codec(&schema, Vec::new(), codec).unwrap(); for (key, value) in metadata { writer.add_user_metadata(key, value).unwrap(); } @@ -1066,22 +1068,41 @@ mod tests { |b| b.build_v2_data(), |writer| { writer.add_file(file(0), -1).unwrap(); - writer.add_existing_file(file(1), 3, 4, Some(4)).unwrap(); + writer.add_delete_file(file(1), 4, Some(4)).unwrap(); + writer.add_existing_file(file(2), 3, 5, Some(5)).unwrap(); }, ) .await; + // ManifestWriter rejects an existing entry without sequence numbers. + let bytes = rewrite_manifest( + &bytes, + Codec::Null, + |_| {}, + |record| { + let AvroValue::Record(fields) = record else { + panic!("manifest entry is not a record"); + }; + if fields.contains(&("status".to_string(), AvroValue::Int(0))) { + for (name, value) in fields { + if name.ends_with("sequence_number") { + *value = AvroValue::Union(0, Box::new(AvroValue::Null)); + } + } + } + }, + ); manifest_file.added_snapshot_id = 42; manifest_file.sequence_number = 7; let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); assert_eq!( column(&batch, "snapshot_id").as_ref(), - &Int64Array::from(vec![42, 3]) as &dyn Array + &Int64Array::from(vec![42, 42, 3]) as &dyn Array ); for name in ["sequence_number", "file_sequence_number"] { assert_eq!( column(&batch, name).as_ref(), - &Int64Array::from(vec![7, 4]) as &dyn Array, + &Int64Array::from(vec![Some(7), Some(4), None]) as &dyn Array, "{name}" ); } @@ -1137,16 +1158,30 @@ mod tests { manifest_file.first_row_id = None; let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); assert_eq!(data_file_column(&batch, "first_row_id").null_count(), 5); - } - #[tokio::test] - async fn test_first_row_id_overflow_is_an_error() { - let (bytes, mut manifest_file) = first_row_id_manifest().await; manifest_file.first_row_id = Some(i64::MAX as u64 - 5); let err = read_manifest_entries(&bytes, &manifest_file).unwrap_err(); assert_eq!(err.kind(), ErrorKind::DataInvalid); } + #[tokio::test] + async fn test_compressed_manifests() { + let (bytes, manifest_file) = first_row_id_manifest().await; + let expected = read_manifest_entries(&bytes, &manifest_file).unwrap(); + for codec in [ + Codec::Deflate(DeflateSettings::default()), + Codec::Snappy, + Codec::Zstandard(ZstandardSettings::default()), + ] { + let bytes = rewrite_manifest(&bytes, codec, |_| {}, |_| {}); + assert_eq!( + read_manifest_entries(&bytes, &manifest_file).unwrap(), + expected, + "{codec:?}" + ); + } + } + #[tokio::test] async fn test_fields_are_matched_by_field_id() { let file = DataFile { @@ -1165,6 +1200,7 @@ mod tests { let renames = [("file_path", "path"), ("i_identity", "renamed_partition")]; let bytes = rewrite_manifest( &bytes, + Codec::Null, |schema| { for_each_avro_field(schema, &|field| { for (from, to) in renames { @@ -1208,6 +1244,7 @@ mod tests { // A file_path under another field ID leaves the required file_path missing. let bytes = rewrite_manifest( &bytes, + Codec::Null, |schema| { for_each_avro_field(schema, &|field| { if field["name"] == "file_path" { @@ -1221,16 +1258,32 @@ mod tests { assert_eq!(err.kind(), ErrorKind::DataInvalid); } - #[tokio::test] - async fn test_manifest_written_by_pyiceberg() { + #[test] + fn test_manifest_written_by_pyiceberg() { let bytes = std::fs::read(format!( "{}/testdata/manifests/pyiceberg-v2-data.avro", env!("CARGO_MANIFEST_DIR") )) .unwrap(); let manifest = Manifest::parse_avro(&bytes).unwrap(); - let (_, manifest_file) = - write_manifest(partition_spec(&[]), Some(1), |b| b.build_v2_data(), |_| {}).await; + let manifest_file = ManifestFile { + manifest_path: "s3://bucket/metadata/manifest.avro".to_string(), + manifest_length: bytes.len() as i64, + partition_spec_id: 0, + content: ManifestContentType::Data, + sequence_number: 1, + min_sequence_number: 1, + added_snapshot_id: 1, + added_files_count: None, + existing_files_count: None, + deleted_files_count: None, + added_rows_count: None, + existing_rows_count: None, + deleted_rows_count: None, + partitions: None, + key_metadata: None, + first_row_id: None, + }; let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); assert_eq!(batch.num_rows(), manifest.entries().len()); @@ -1257,16 +1310,6 @@ mod tests { binary_map(&data_file_column(&batch, "lower_bounds"), row), lower_bounds ); - let mut value_counts: Vec<_> = data_file - .value_counts() - .iter() - .map(|(id, count)| (*id, *count as i64)) - .collect(); - value_counts.sort_unstable(); - assert_eq!( - long_map(&data_file_column(&batch, "value_counts"), row), - value_counts - ); } } } From 3311b892417540e9d6ad910c840460c385e679ee Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Mon, 12 Oct 2026 17:16:31 -0400 Subject: [PATCH 6/7] Return an error if arrow-avro splits a manifest into more than one batch --- crates/iceberg/src/arrow/manifest.rs | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index 21c5658d87..17ad866f44 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -145,11 +145,15 @@ pub fn read_manifest_entries(bytes: &[u8], manifest_file: &ManifestFile) -> Resu let Some(batch) = reader.next() else { return Ok(RecordBatch::new_empty(schema)); }; - let entries = project_struct( - &StructArray::from(batch.map_err(invalid_manifest)?), - &avro_schema, - schema.fields(), - )?; + let batch = batch.map_err(invalid_manifest)?; + // Taking only the first batch relies on arrow-avro filling it to the batch + // size, so fail instead of dropping entries if that changes. + if reader.next().is_some() { + return Err(invalid_data!( + "arrow-avro split the manifest into more than one batch" + )); + } + let entries = project_struct(&StructArray::from(batch), &avro_schema, schema.fields())?; Ok(RecordBatch::from(inherit(entries, manifest_file)?)) } From 7c962377c46293dd87b409068fe4869a353f10bc Mon Sep 17 00:00:00 2001 From: Matt Butrovich Date: Tue, 20 Oct 2026 13:29:20 -0400 Subject: [PATCH 7/7] fix v2 discrepancy --- crates/iceberg/src/arrow/manifest.rs | 125 ++++++++++++++++++++++++--- 1 file changed, 114 insertions(+), 11 deletions(-) diff --git a/crates/iceberg/src/arrow/manifest.rs b/crates/iceberg/src/arrow/manifest.rs index 17ad866f44..8e06d9ea56 100644 --- a/crates/iceberg/src/arrow/manifest.rs +++ b/crates/iceberg/src/arrow/manifest.rs @@ -40,7 +40,7 @@ use serde_json::Value as JsonValue; use crate::arrow::schema_to_arrow_schema; use crate::error::invalid_data; use crate::spec::{ - ManifestContentType, ManifestFile, ManifestMetadata, ManifestStatus, Schema, + FormatVersion, ManifestContentType, ManifestFile, ManifestMetadata, ManifestStatus, Schema, manifest_entry_fields, }; use crate::{Error, Result}; @@ -153,7 +153,12 @@ pub fn read_manifest_entries(bytes: &[u8], manifest_file: &ManifestFile) -> Resu "arrow-avro split the manifest into more than one batch" )); } - let entries = project_struct(&StructArray::from(batch), &avro_schema, schema.fields())?; + let entries = project_struct( + &StructArray::from(batch), + &avro_schema, + schema.fields(), + *metadata.format_version(), + )?; Ok(RecordBatch::from(inherit(entries, manifest_file)?)) } @@ -206,7 +211,12 @@ fn non_null_type(avro_type: &JsonValue) -> &JsonValue { /// Builds each of `fields` from the column of `array` that the Avro record /// `record` gives the same field ID. -fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> Result { +fn project_struct( + array: &StructArray, + record: &JsonValue, + fields: &Fields, + format_version: FormatVersion, +) -> Result { let avro_fields = record["fields"] .as_array() .ok_or_else(|| invalid_data!("manifest Avro type {record} is not a record"))?; @@ -222,8 +232,9 @@ fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> R array.column(index), non_null_type(&avro_fields[index]["type"]), field, + format_version, )?, - None => missing_field(field, array.len())?, + None => missing_field(field, array.len(), format_version)?, }, ); } @@ -232,12 +243,15 @@ fn project_struct(array: &StructArray, record: &JsonValue, fields: &Fields) -> R } /// Builds the column for a field that the manifest doesn't have. -fn missing_field(field: &Field, len: usize) -> Result { +fn missing_field(field: &Field, len: usize, format_version: FormatVersion) -> Result { // v1 manifests have no content or sequence number columns, which v2+ - // readers read as 0. + // readers read as 0. A v2+ manifest may lack file_sequence_number (Java + // wrote none before 1.1.0), which reads as null and is inherited. match field_id(field) { CONTENT_FIELD_ID => Ok(Arc::new(Int32Array::from_value(0, len))), - SEQUENCE_NUMBER_FIELD_ID | FILE_SEQUENCE_NUMBER_FIELD_ID => { + SEQUENCE_NUMBER_FIELD_ID | FILE_SEQUENCE_NUMBER_FIELD_ID + if format_version == FormatVersion::V1 => + { Ok(Arc::new(Int64Array::from_value(0, len))) } _ if field.is_nullable() => Ok(new_null_array(field.data_type(), len)), @@ -253,7 +267,12 @@ fn missing_field(field: &Field, len: usize) -> Result { /// Converts the decoded `array` of Avro type `avro_type` to `field`, the /// spec's Arrow field for it. -fn project(array: &ArrayRef, avro_type: &JsonValue, field: &Field) -> Result { +fn project( + array: &ArrayRef, + avro_type: &JsonValue, + field: &Field, + format_version: FormatVersion, +) -> Result { let type_mismatch = || { invalid_data!( "manifest field {} has Avro type {avro_type}, which doesn't match {}", @@ -266,10 +285,16 @@ fn project(array: &ArrayRef, avro_type: &JsonValue, field: &Field) -> Result { let list = array.as_list_opt::().ok_or_else(type_mismatch)?; - let values = project(list.values(), non_null_type(&avro_type["items"]), element)?; + let values = project( + list.values(), + non_null_type(&avro_type["items"]), + element, + format_version, + )?; Arc::new( ListArray::try_new( element.clone(), @@ -290,6 +315,7 @@ fn project(array: &ArrayRef, avro_type: &JsonValue, field: &Field) -> Result, value: i64) -> Resu } /// Assigns first row IDs to the data files that lack one, following first row -/// ID inheritance. As in Java's `ManifestReader`, deleted entries take no row -/// IDs, and a manifest without a first row ID clears them all. +/// ID inheritance. Deleted entries take no row IDs, because the spec's +/// suggested manifest `first_row_id` increment, `added_rows_count` plus +/// `existing_rows_count`, leaves no room for them. A manifest without a first +/// row ID belongs to a snapshot without row lineage, so all its row IDs are +/// null. fn inherit_first_row_ids( first_row_ids: &Int64Array, status: &Int32Array, @@ -414,6 +443,12 @@ fn inherit_first_row_ids( { match first_row_id { None if status != ManifestStatus::Deleted as i32 => { + if record_count < 0 { + return Err(invalid_data!( + "negative record_count {record_count} assigning first row IDs in {}", + manifest_file.manifest_path + )); + } builder.append_value(next_row_id); next_row_id = next_row_id.checked_add(record_count).ok_or_else(|| { invalid_data!( @@ -1112,6 +1147,49 @@ mod tests { } } + #[tokio::test] + async fn test_v2_manifest_without_file_sequence_number_inherits() { + let file = |i: usize| { + data_file(DataContentType::Data, &format!("s3://b/data/{i}.parquet")) + .build() + .unwrap() + }; + let (bytes, mut manifest_file) = write_manifest( + partition_spec(&[]), + Some(1), + |b| b.build_v2_data(), + |writer| { + writer.add_file(file(0), -1).unwrap(); + writer.add_existing_file(file(1), 3, 5, Some(5)).unwrap(); + }, + ) + .await; + // Java wrote v2 manifests without file_sequence_number before 1.1.0. + let bytes = rewrite_manifest( + &bytes, + Codec::Null, + |schema| { + schema["fields"] + .as_array_mut() + .unwrap() + .retain(|field| field["name"] != "file_sequence_number") + }, + |record| { + let AvroValue::Record(fields) = record else { + panic!("manifest entry is not a record"); + }; + fields.retain(|(name, _)| name != "file_sequence_number"); + }, + ); + manifest_file.sequence_number = 7; + let batch = read_manifest_entries(&bytes, &manifest_file).unwrap(); + + assert_eq!( + column(&batch, "file_sequence_number").as_ref(), + &Int64Array::from(vec![Some(7), None]) as &dyn Array + ); + } + async fn first_row_id_manifest() -> (Vec, ManifestFile) { let file = |i: usize, record_count: u64, first_row_id: Option| DataFile { first_row_id, @@ -1166,6 +1244,31 @@ mod tests { manifest_file.first_row_id = Some(i64::MAX as u64 - 5); let err = read_manifest_entries(&bytes, &manifest_file).unwrap_err(); assert_eq!(err.kind(), ErrorKind::DataInvalid); + + // A negative record_count would assign overlapping row IDs. + let bytes = rewrite_manifest( + &bytes, + Codec::Null, + |_| {}, + |record| { + let AvroValue::Record(fields) = record else { + panic!("manifest entry is not a record"); + }; + let Some((_, AvroValue::Record(data_file))) = + fields.iter_mut().find(|(name, _)| name == "data_file") + else { + panic!("manifest entry has no data_file record"); + }; + for (name, value) in data_file { + if name == "record_count" && *value == AvroValue::Long(10) { + *value = AvroValue::Long(-10); + } + } + }, + ); + manifest_file.first_row_id = Some(100); + let err = read_manifest_entries(&bytes, &manifest_file).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); } #[tokio::test]