mbutrovich commented on code in PR #3388: URL: https://github.com/apache/iceberg-rust/pull/3388#discussion_r4249587989
########## 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 { + 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, + )?; + } + } + 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<ArrayRef> { + 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<ArrayRef> { + 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::<Result<_>>()?; Review Comment: https://github.com/apache/iceberg-rust/pull/3388#discussion_r4238992793 > Nit: `collect::<Result<Int64Array>>()` doesn't pre-size (unverified, based on the size hint through the `Result` adapter), and the length is known. A plain loop over an `Int64Builder::with_capacity(first_row_ids.len())` with `append_option` avoids both the regrowth and the `Result` collect. The `collect::<Result<_>>()` of columns in `project_struct` can likewise push into `Vec::with_capacity(fields.len())`. You're right. The `Result` adapter reports a lower size hint of 0, and [`PrimitiveArray::from_iter`](https://github.com/apache/arrow-rs/blob/f90e061326bd821a7af09281d9e92de6f3b603d9/arrow-array/src/array/primitive_array.rs#L1482-L1489) sizes its buffers from that lower bound. First row IDs now go into an `Int64Builder::with_capacity`, and `project_struct` pushes into a `Vec::with_capacity(fields.len())`. -- 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]
