vendor/iceberg-arrow59/src/inspect/manifests.rs
// 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.
use std::collections::HashMap;
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_array::builder::{
BooleanBuilder, GenericListBuilder, ListBuilder, PrimitiveBuilder, StringBuilder, StructBuilder,
};
use arrow_array::types::{Int32Type, Int64Type};
use arrow_schema::{DataType, Field, Fields};
use futures::{StreamExt, stream};
use crate::Result;
use crate::arrow::schema_to_arrow_schema;
use crate::error::invalid_data;
use crate::scan::ArrowRecordBatchStream;
use crate::spec::{Datum, FieldSummary, ListType, NestedField, PrimitiveType, StructType, Type};
use crate::table::Table;
/// Manifests table.
pub struct ManifestsTable<'a> {
table: &'a Table,
}
impl<'a> ManifestsTable<'a> {
/// Create a new Manifests table instance.
pub fn new(table: &'a Table) -> Self {
Self { table }
}
/// Returns the iceberg schema of the manifests table.
pub fn schema(&self) -> crate::spec::Schema {
let fields = vec![
NestedField::new(14, "content", Type::Primitive(PrimitiveType::Int), true),
NestedField::new(1, "path", Type::Primitive(PrimitiveType::String), true),
NestedField::new(2, "length", Type::Primitive(PrimitiveType::Long), true),
NestedField::new(
3,
"partition_spec_id",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
4,
"added_snapshot_id",
Type::Primitive(PrimitiveType::Long),
true,
),
NestedField::new(
5,
"added_data_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
6,
"existing_data_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
7,
"deleted_data_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
15,
"added_delete_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
16,
"existing_delete_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
17,
"deleted_delete_files_count",
Type::Primitive(PrimitiveType::Int),
true,
),
NestedField::new(
8,
"partition_summaries",
Type::List(ListType {
element_field: Arc::new(NestedField::new(
9,
"item",
Type::Struct(StructType::new(vec![
Arc::new(NestedField::new(
10,
"contains_null",
Type::Primitive(PrimitiveType::Boolean),
true,
)),
Arc::new(NestedField::new(
11,
"contains_nan",
Type::Primitive(PrimitiveType::Boolean),
false,
)),
Arc::new(NestedField::new(
12,
"lower_bound",
Type::Primitive(PrimitiveType::String),
false,
)),
Arc::new(NestedField::new(
13,
"upper_bound",
Type::Primitive(PrimitiveType::String),
false,
)),
])),
true,
)),
}),
true,
),
];
crate::spec::Schema::builder()
.with_fields(fields.into_iter().map(|f| f.into()))
.build()
.unwrap()
}
/// Scans the manifests table.
pub async fn scan(&self) -> Result<ArrowRecordBatchStream> {
let schema = schema_to_arrow_schema(&self.schema())?;
let mut content = PrimitiveBuilder::<Int32Type>::new();
let mut path = StringBuilder::new();
let mut length = PrimitiveBuilder::<Int64Type>::new();
let mut partition_spec_id = PrimitiveBuilder::<Int32Type>::new();
let mut added_snapshot_id = PrimitiveBuilder::<Int64Type>::new();
let mut added_data_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut existing_data_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut deleted_data_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut added_delete_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut existing_delete_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut deleted_delete_files_count = PrimitiveBuilder::<Int32Type>::new();
let mut partition_summaries = self.partition_summary_builder()?;
if let Some(snapshot) = self.table.metadata().current_snapshot() {
let manifest_list = self.table.manifest_list_reader(snapshot).load().await?;
for manifest in manifest_list.entries() {
content.append_value(manifest.content as i32);
path.append_value(manifest.manifest_path.clone());
length.append_value(manifest.manifest_length);
partition_spec_id.append_value(manifest.partition_spec_id);
added_snapshot_id.append_value(manifest.added_snapshot_id);
added_data_files_count.append_value(manifest.added_files_count.unwrap_or(0) as i32);
existing_data_files_count
.append_value(manifest.existing_files_count.unwrap_or(0) as i32);
deleted_data_files_count
.append_value(manifest.deleted_files_count.unwrap_or(0) as i32);
added_delete_files_count
.append_value(manifest.added_files_count.unwrap_or(0) as i32);
existing_delete_files_count
.append_value(manifest.existing_files_count.unwrap_or(0) as i32);
deleted_delete_files_count
.append_value(manifest.deleted_files_count.unwrap_or(0) as i32);
let spec = self
.table
.metadata()
.partition_spec_by_id(manifest.partition_spec_id)
.ok_or_else(|| {
invalid_data!(
"Partition spec {} for manifest {} is not in table metadata",
manifest.partition_spec_id,
manifest.manifest_path
)
})?;
let spec_struct = spec.partition_type(self.table.metadata().current_schema())?;
self.append_partition_summaries(
&mut partition_summaries,
manifest.partitions.as_deref().unwrap_or(&[]),
spec_struct,
);
}
}
let batch = RecordBatch::try_new(Arc::new(schema), vec![
Arc::new(content.finish()),
Arc::new(path.finish()),
Arc::new(length.finish()),
Arc::new(partition_spec_id.finish()),
Arc::new(added_snapshot_id.finish()),
Arc::new(added_data_files_count.finish()),
Arc::new(existing_data_files_count.finish()),
Arc::new(deleted_data_files_count.finish()),
Arc::new(added_delete_files_count.finish()),
Arc::new(existing_delete_files_count.finish()),
Arc::new(deleted_delete_files_count.finish()),
Arc::new(partition_summaries.finish()),
])?;
Ok(stream::iter(vec![Ok(batch)]).boxed())
}
fn partition_summary_builder(&self) -> Result<GenericListBuilder<i32, StructBuilder>> {
let schema = schema_to_arrow_schema(&self.schema())?;
let partition_summary_fields =
match schema.field_with_name("partition_summaries")?.data_type() {
DataType::List(list_type) => match list_type.data_type() {
DataType::Struct(fields) => fields.to_vec(),
_ => unreachable!(),
},
_ => unreachable!(),
};
let partition_summaries = ListBuilder::new(StructBuilder::from_fields(
Fields::from(partition_summary_fields.clone()),
0,
))
.with_field(Arc::new(
Field::new_struct("item", partition_summary_fields, false).with_metadata(
HashMap::from([("PARQUET:field_id".to_string(), "9".to_string())]),
),
));
Ok(partition_summaries)
}
fn append_partition_summaries(
&self,
builder: &mut GenericListBuilder<i32, StructBuilder>,
partitions: &[FieldSummary],
partition_struct: StructType,
) {
let partition_summaries_builder = builder.values();
for (summary, field) in partitions.iter().zip(partition_struct.fields()) {
partition_summaries_builder
.field_builder::<BooleanBuilder>(0)
.unwrap()
.append_value(summary.contains_null);
partition_summaries_builder
.field_builder::<BooleanBuilder>(1)
.unwrap()
.append_option(summary.contains_nan);
let field_type = field.field_type.as_primitive_type().unwrap();
for (index, bound) in [(2, &summary.lower_bound), (3, &summary.upper_bound)] {
// Bounds cannot be decoded when the source column's type is unknown.
let bound = bound
.as_ref()
.filter(|_| *field_type != PrimitiveType::Unknown)
.map(|v| {
Datum::try_from_bytes(v, field_type.clone())
.unwrap()
.to_string()
});
partition_summaries_builder
.field_builder::<StringBuilder>(index)
.unwrap()
.append_option(bound);
}
partition_summaries_builder.append(true);
}
builder.append(true);
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow_array::{Array, ListArray, StructArray};
use expect_test::expect;
use futures::TryStreamExt;
use crate::spec::TableMetadata;
use crate::test_utils::check_record_batches;
use crate::test_utils::scan::TableTestFixture;
#[tokio::test]
async fn test_manifests_table() {
let mut fixture = TableTestFixture::new();
fixture.setup_manifest_files().await;
let record_batch = fixture.table.inspect().manifests().scan().await.unwrap();
check_record_batches(
record_batch.try_collect::<Vec<_>>().await.unwrap(),
expect![[r#"
Field { "content": Int32, metadata: {"PARQUET:field_id": "14"} },
Field { "path": Utf8, metadata: {"PARQUET:field_id": "1"} },
Field { "length": Int64, metadata: {"PARQUET:field_id": "2"} },
Field { "partition_spec_id": Int32, metadata: {"PARQUET:field_id": "3"} },
Field { "added_snapshot_id": Int64, metadata: {"PARQUET:field_id": "4"} },
Field { "added_data_files_count": Int32, metadata: {"PARQUET:field_id": "5"} },
Field { "existing_data_files_count": Int32, metadata: {"PARQUET:field_id": "6"} },
Field { "deleted_data_files_count": Int32, metadata: {"PARQUET:field_id": "7"} },
Field { "added_delete_files_count": Int32, metadata: {"PARQUET:field_id": "15"} },
Field { "existing_delete_files_count": Int32, metadata: {"PARQUET:field_id": "16"} },
Field { "deleted_delete_files_count": Int32, metadata: {"PARQUET:field_id": "17"} },
Field { "partition_summaries": List(non-null Struct("contains_null": non-null Boolean, metadata: {"PARQUET:field_id": "10"}, "contains_nan": Boolean, metadata: {"PARQUET:field_id": "11"}, "lower_bound": Utf8, metadata: {"PARQUET:field_id": "12"}, "upper_bound": Utf8, metadata: {"PARQUET:field_id": "13"}), metadata: {"PARQUET:field_id": "9"}), metadata: {"PARQUET:field_id": "8"} }"#]],
expect![[r#"
content: PrimitiveArray<Int32>
[
0,
],
path: (skipped),
length: (skipped),
partition_spec_id: PrimitiveArray<Int32>
[
0,
],
added_snapshot_id: PrimitiveArray<Int64>
[
3055729675574597004,
],
added_data_files_count: PrimitiveArray<Int32>
[
1,
],
existing_data_files_count: PrimitiveArray<Int32>
[
1,
],
deleted_data_files_count: PrimitiveArray<Int32>
[
1,
],
added_delete_files_count: PrimitiveArray<Int32>
[
1,
],
existing_delete_files_count: PrimitiveArray<Int32>
[
1,
],
deleted_delete_files_count: PrimitiveArray<Int32>
[
1,
],
partition_summaries: ListArray
[
StructArray
-- validity:
[
valid,
]
[
-- child 0: "contains_null" (Boolean)
BooleanArray
[
false,
]
-- child 1: "contains_nan" (Boolean)
BooleanArray
[
false,
]
-- child 2: "lower_bound" (Utf8)
StringArray
[
"100",
]
-- child 3: "upper_bound" (Utf8)
StringArray
[
"300",
]
],
]"#]],
&["path", "length"],
Some("path"),
);
}
#[tokio::test]
async fn test_manifests_table_with_dropped_partition_source_column() {
let mut fixture = TableTestFixture::new();
fixture.setup_manifest_files().await;
// Evolve the table so that the manifests reference a historical spec whose source
// column is no longer in the current schema: add an unpartitioned default spec, then
// drop the source column of the original spec.
let mut metadata = serde_json::to_value(fixture.table.metadata()).unwrap();
let current_schema_id = metadata["current-schema-id"].clone();
let schemas = metadata["schemas"].as_array_mut().unwrap();
for schema in schemas {
if schema["schema-id"] == current_schema_id {
let fields = schema["fields"].as_array_mut().unwrap();
fields.retain(|field| field["id"] != 1);
let identifier_ids = schema["identifier-field-ids"].as_array_mut().unwrap();
identifier_ids.retain(|id| *id != 1);
}
}
metadata["partition-specs"]
.as_array_mut()
.unwrap()
.push(serde_json::json!({"spec-id": 1, "fields": []}));
metadata["default-spec-id"] = serde_json::json!(1);
let metadata: TableMetadata = serde_json::from_value(metadata).unwrap();
let table = fixture.table.clone().with_metadata(Arc::new(metadata));
let batches: Vec<_> = table
.inspect()
.manifests()
.scan()
.await
.unwrap()
.try_collect()
.await
.unwrap();
assert_eq!(batches.len(), 1);
assert_eq!(batches[0].num_rows(), 1);
let summaries = batches[0]
.column_by_name("partition_summaries")
.unwrap()
.as_any()
.downcast_ref::<ListArray>()
.unwrap();
let summary = summaries.value(0);
let summary = summary.as_any().downcast_ref::<StructArray>().unwrap();
assert_eq!(summary.len(), 1);
assert!(summary.column_by_name("contains_null").unwrap().is_valid(0));
assert!(summary.column_by_name("contains_nan").unwrap().is_valid(0));
assert!(summary.column_by_name("lower_bound").unwrap().is_null(0));
assert!(summary.column_by_name("upper_bound").unwrap().is_null(0));
}
}