mbutrovich commented on code in PR #3388: URL: https://github.com/apache/iceberg-rust/pull/3388#discussion_r4249585655
########## crates/iceberg/src/arrow/manifest.rs: ########## @@ -0,0 +1,1272 @@ +// 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. + +//! Reading manifest entries into Arrow. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow_arith::boolean::{and, is_null}; +use arrow_array::cast::AsArray; +use arrow_array::types::{TimestampMicrosecondType, TimestampNanosecondType}; +use arrow_array::{ + Array, ArrayRef, BooleanArray, Int32Array, Int64Array, ListArray, MapArray, RecordBatch, + StructArray, new_null_array, +}; +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::schema_to_arrow_schema; +use crate::error::invalid_data; +use crate::spec::{ + ManifestContentType, ManifestFile, ManifestMetadata, ManifestStatus, Schema, + manifest_entry_fields, +}; +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; + +/// Reads the entries of a manifest file into a [`RecordBatch`] with one row +/// per entry. +/// +/// `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. +/// +/// 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 iceberg::arrow::read_manifest_entries; +/// use iceberg::spec::{ManifestContentType, ManifestFile}; +/// +/// 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 batch = read_manifest_entries(&bytes, &manifest_file)?; +/// assert_eq!(batch.num_rows(), 2); +/// # Ok::<(), Box<dyn std::error::Error>>(()) +/// ``` +pub fn read_manifest_entries(bytes: &[u8], manifest_file: &ManifestFile) -> Result<RecordBatch> { + // 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<String, Vec<u8>> = 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())?; + let schema = Arc::new(schema_to_arrow_schema( + &Schema::builder() + .with_fields(manifest_entry_fields(&partition_type)) + .build()?, + )?); + + // 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 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::from(inherit(entries, manifest_file)?)) +} + +/// 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; + }; + 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 +} + +fn invalid_manifest(err: ArrowError) -> Error { + invalid_data!("failed to read manifest").with_source(err) +} + +/// 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() +} + +/// 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, + } +} + +/// 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<StructArray> { + 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); + match avro_fields + .iter() + .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::<Result<_>>()?; + StructArray::try_new_with_length(fields.clone(), columns, array.nulls().cloned(), array.len()) + .map_err(invalid_manifest) +} + +/// Builds the column for a field that the manifest doesn't have. +fn missing_field(field: &Field, len: usize) -> Result<ArrayRef> { + // 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<ArrayRef> { + let type_mismatch = || { + invalid_data!( + "manifest field {} has Avro type {avro_type}, which doesn't match {}", + field.name(), + field.data_type() + ) + }; + 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::<i32>().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)?, + ) + } + // Avro stores maps with int keys as arrays of key-value records. + DataType::Map(entries, _) => { + let list = array.as_list_opt::<i32>().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)?, + ) + } + // 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::<TimestampMicrosecondType>() + .ok_or_else(type_mismatch)? + .clone() + .with_timezone_opt(time_zone.clone()), + ), + DataType::Timestamp(TimeUnit::Nanosecond, time_zone) => Arc::new( + array + .as_primitive_opt::<TimestampNanosecondType>() + .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<StructArray> { + 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 { Review Comment: > Sequence number inheritance only fills `ADDED` rows. Java's `InheritableMetadataFactory` and `ManifestEntry::inherit_data` also fill any null when the manifest list entry's sequence number is 0 (`sequenceNumber == 0 || status == ADDED`). Here, a v2-layout manifest listed with sequence number 0 returns nulls for `EXISTING` and `DELETED` entries, where the `Manifest` path returns 0. Adding `|| manifest_file.sequence_number == INITIAL_SEQUENCE_NUMBER` to the mask would match. I'd like to keep the `ADDED`-only mask, because it follows the spec's [sequence number inheritance](https://github.com/apache/iceberg/blob/5e7169168db3d34e29354c6f59ec4d6e420b8d2d/format/spec.md?plain=1#L927-L938) rules. The [`manifest_entry` fields](https://github.com/apache/iceberg/blob/5e7169168db3d34e29354c6f59ec4d6e420b8d2d/format/spec.md?plain=1#L692-L693) say both sequence numbers are "Inherited when null and status is 1 (added)", and the [entry rules](https://github.com/apache/iceberg/blob/5e7169168db3d34e29354c6f59ec4d6e420b8d2d/format/spec.md?plain=1#L705) say "If the entry status is 0 (existing) or 2 (deleted), the entry must include both sequence numbers explicitly." The spec asks for 0 only when [reading v1 manifests](https://github.com/apache/iceberg/blob/5e7169168db3d34e29354c6f59ec4d6e420b8d2d/format/spec.md?plain=1#L938): "When reading v1 manifests with no sequence number column, sequence numbers for all files must default to 0." This reader fills the missing colu mn with 0 in that case, and `test_v1_manifest_reads_with_v2_defaults` covers it. A v2-layout manifest with a null sequence number on an existing or deleted entry doesn't follow the spec, so I'd rather not add Java's fallback for it here. Java's [`V2Metadata`](https://github.com/apache/iceberg/blob/5e7169168db3d34e29354c6f59ec4d6e420b8d2d/core/src/main/java/org/apache/iceberg/V2Metadata.java#L330-L333) and PyIceberg's [`ManifestWriterV2`](https://github.com/apache/iceberg-python/blob/0d584073f67948c46f34d4c2a1672455bc10b5c9/pyiceberg/manifest.py#L1324-L1329) both reject a null sequence number on a non-added entry, so I don't expect these manifests from either writer. Do you know of another writer that produces them? > The `ADDED`-only restriction is also untested. Replacing `Some(&added)` with `None` here leaves all 11 tests green (I ran it), because `test_snapshot_id_and_sequence_number_inheritance` only uses an `EXISTING` entry with explicit values. Could you add an `EXISTING` entry with null sequence numbers that must stay null (`add_existing_entry`, as in `first_row_id_manifest`), plus the sequence number 0 case? Thanks, I added it. `test_snapshot_id_and_sequence_number_inheritance` now has an added entry that inherits, a deleted entry with explicit values, and an existing entry whose null sequence numbers stay null. `ManifestWriter` rejects an existing entry without sequence numbers, including through `add_existing_entry`, so the test nulls them with `rewrite_manifest`. I left out the sequence number 0 case for the reason above. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
