vendor/iceberg-arrow59/src/writer/base_writer/equality_delete_writer.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.
//! This module provide `EqualityDeleteWriter`.
use std::collections::HashSet;
use std::sync::Arc;
use arrow_array::RecordBatch;
use arrow_schema::{DataType, Field, SchemaRef as ArrowSchemaRef};
use itertools::Itertools;
use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
use crate::arrow::record_batch_projector::RecordBatchProjector;
use crate::arrow::schema_to_arrow_schema;
use crate::error::invalid_data;
use crate::spec::{DataFile, PartitionKey, Schema, SchemaRef};
use crate::writer::file_writer::FileWriterBuilder;
use crate::writer::file_writer::location_generator::{FileNameGenerator, LocationGenerator};
use crate::writer::file_writer::rolling_writer::{RollingFileWriter, RollingFileWriterBuilder};
use crate::writer::{IcebergWriter, IcebergWriterBuilder};
use crate::{Error, ErrorKind, Result};
/// Builder for `EqualityDeleteWriter`.
#[derive(Debug)]
pub struct EqualityDeleteFileWriterBuilder<
B: FileWriterBuilder,
L: LocationGenerator,
F: FileNameGenerator,
> {
inner: RollingFileWriterBuilder<B, L, F>,
config: EqualityDeleteWriterConfig,
}
impl<B, L, F> EqualityDeleteFileWriterBuilder<B, L, F>
where
B: FileWriterBuilder,
L: LocationGenerator,
F: FileNameGenerator,
{
/// Create a new `EqualityDeleteFileWriterBuilder` using a `RollingFileWriterBuilder`.
pub fn new(
inner: RollingFileWriterBuilder<B, L, F>,
config: EqualityDeleteWriterConfig,
) -> Self {
Self { inner, config }
}
}
/// Config for `EqualityDeleteWriter`.
#[derive(Debug)]
pub struct EqualityDeleteWriterConfig {
// Field ids used to determine row equality in equality delete files.
equality_ids: Vec<i32>,
// Projector used to project the data chunk into specific fields.
projector: RecordBatchProjector,
}
/// Validates equality delete field ids before constructing the writer config.
///
/// In addition to rejecting an empty id list, this check also rejects duplicate
/// ids and ids that do not exist in the table schema. These two checks are
/// intentionally stricter than the current Java reference implementation, which
/// only guards against null/empty lists; catching duplicates and missing ids
/// here produces clearer errors and avoids generating meaningless delete files.
fn validate_equality_ids(equality_ids: &[i32], original_schema: &Schema) -> Result<()> {
if equality_ids.is_empty() {
return Err(invalid_data!(
"Equality delete field ids must not be empty."
));
}
let mut seen = HashSet::with_capacity(equality_ids.len());
for id in equality_ids {
if !seen.insert(*id) {
return Err(invalid_data!("Duplicate equality delete field id: {id}")
.with_context("field_id", id.to_string()));
}
if original_schema.field_by_id(*id).is_none() {
return Err(invalid_data!("Invalid equality delete field id: {id}")
.with_context("field_id", id.to_string()));
}
}
Ok(())
}
impl EqualityDeleteWriterConfig {
/// Create a new `EqualityDeleteWriterConfig` with equality ids.
pub fn new(equality_ids: Vec<i32>, original_schema: SchemaRef) -> Result<Self> {
validate_equality_ids(&equality_ids, &original_schema)?;
let original_arrow_schema = Arc::new(schema_to_arrow_schema(&original_schema)?);
let projector = RecordBatchProjector::new(
original_arrow_schema,
&equality_ids,
// Equality delete fields follow identifier-field type restrictions,
// except optional columns and columns nested under optional structs are allowed.
// Project only primitive, non-floating fields; RecordBatchProjector's traversal
// keeps fields under maps and lists unreachable.
|field| {
// Only primitive type is allowed to be used for identifier field ids
if field.data_type().is_nested()
|| matches!(
field.data_type(),
DataType::Float16 | DataType::Float32 | DataType::Float64
)
{
return Ok(None);
}
Ok(Some(
field
.metadata()
.get(PARQUET_FIELD_ID_META_KEY)
.ok_or_else(|| {
Error::new(ErrorKind::Unexpected, "Field metadata is missing.")
})?
.parse::<i64>()
.map_err(|e| Error::new(ErrorKind::Unexpected, e.to_string()))?,
))
},
|_field: &Field| true,
)?;
Ok(Self {
equality_ids,
projector,
})
}
/// Return projected Schema
pub fn projected_arrow_schema_ref(&self) -> &ArrowSchemaRef {
self.projector.projected_schema_ref()
}
}
#[async_trait::async_trait]
impl<B, L, F> IcebergWriterBuilder for EqualityDeleteFileWriterBuilder<B, L, F>
where
B: FileWriterBuilder,
L: LocationGenerator,
F: FileNameGenerator,
{
type R = EqualityDeleteFileWriter<B, L, F>;
async fn build(&self, partition_key: Option<PartitionKey>) -> Result<Self::R> {
Ok(EqualityDeleteFileWriter {
inner: Some(self.inner.build()),
projector: self.config.projector.clone(),
equality_ids: self.config.equality_ids.clone(),
partition_key,
})
}
}
/// Writer used to write equality delete files.
#[derive(Debug)]
pub struct EqualityDeleteFileWriter<
B: FileWriterBuilder,
L: LocationGenerator,
F: FileNameGenerator,
> {
inner: Option<RollingFileWriter<B, L, F>>,
projector: RecordBatchProjector,
equality_ids: Vec<i32>,
partition_key: Option<PartitionKey>,
}
#[async_trait::async_trait]
impl<B, L, F> IcebergWriter for EqualityDeleteFileWriter<B, L, F>
where
B: FileWriterBuilder,
L: LocationGenerator,
F: FileNameGenerator,
{
async fn write(&mut self, batch: RecordBatch) -> Result<()> {
let batch = self.projector.project_batch(batch)?;
if let Some(writer) = self.inner.as_mut() {
writer.write(&self.partition_key, &batch).await
} else {
Err(Error::new(
ErrorKind::Unexpected,
"Equality delete inner writer has been closed.",
))
}
}
async fn close(&mut self) -> Result<Vec<DataFile>> {
if let Some(writer) = self.inner.take() {
writer
.close()
.await?
.into_iter()
.map(|mut res| {
res.content(crate::spec::DataContentType::EqualityDeletes);
res.equality_ids(Some(self.equality_ids.iter().copied().collect_vec()));
if let Some(pk) = self.partition_key.as_ref() {
res.partition(pk.data().clone());
res.partition_spec_id(pk.spec().spec_id());
}
res.build()
.map_err(|e| invalid_data!("Failed to build data file: {e}"))
})
.collect()
} else {
Err(Error::new(
ErrorKind::Unexpected,
"Equality delete inner writer has been closed.",
))
}
}
}
#[cfg(test)]
mod test {
use std::collections::HashMap;
use std::sync::Arc;
use arrow_array::types::Int32Type;
use arrow_array::{ArrayRef, BooleanArray, Int32Array, Int64Array, RecordBatch, StructArray};
use arrow_buffer::NullBuffer;
use arrow_schema::{DataType, Field, Fields};
use arrow_select::concat::concat_batches;
use itertools::Itertools;
use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::file::properties::WriterProperties;
use tempfile::TempDir;
use uuid::Uuid;
use crate::ErrorKind;
use crate::arrow::{arrow_schema_to_schema, schema_to_arrow_schema};
use crate::io::FileIO;
use crate::spec::{
DataFile, DataFileFormat, ListType, MapType, NestedField, PrimitiveType, Schema,
StructType, Type,
};
use crate::writer::base_writer::equality_delete_writer::{
EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig,
};
use crate::writer::file_writer::ParquetWriterBuilder;
use crate::writer::file_writer::location_generator::{
DefaultFileNameGenerator, DefaultLocationGenerator,
};
use crate::writer::file_writer::rolling_writer::RollingFileWriterBuilder;
use crate::writer::{IcebergWriter, IcebergWriterBuilder};
async fn check_parquet_data_file_with_equality_delete_write(
file_io: &FileIO,
data_file: &DataFile,
batch: &RecordBatch,
) {
assert_eq!(data_file.file_format, DataFileFormat::Parquet);
// read the written file
let input_file = file_io.new_input(data_file.file_path.clone()).unwrap();
// read the written file
let input_content = input_file.read().await.unwrap();
let reader_builder =
ParquetRecordBatchReaderBuilder::try_new(input_content.clone()).unwrap();
let metadata = reader_builder.metadata().clone();
// check data
let reader = reader_builder.build().unwrap();
let batches = reader.map(|batch| batch.unwrap()).collect::<Vec<_>>();
let res = concat_batches(&batch.schema(), &batches).unwrap();
assert_eq!(*batch, res);
// check metadata
let expect_column_num = batch.num_columns();
assert_eq!(
data_file.record_count,
metadata
.row_groups()
.iter()
.map(|group| group.num_rows())
.sum::<i64>() as u64
);
assert_eq!(data_file.file_size_in_bytes, input_content.len() as u64);
assert_eq!(data_file.column_sizes.len(), expect_column_num);
for (index, id) in data_file.column_sizes().keys().sorted().enumerate() {
metadata
.row_groups()
.iter()
.map(|group| group.columns())
.for_each(|column| {
assert_eq!(
*data_file.column_sizes.get(id).unwrap() as i64,
column.get(index).unwrap().compressed_size()
);
});
}
assert_eq!(data_file.value_counts.len(), expect_column_num);
data_file.value_counts.iter().for_each(|(_, &v)| {
let expect = metadata
.row_groups()
.iter()
.map(|group| group.num_rows())
.sum::<i64>() as u64;
assert_eq!(v, expect);
});
for (index, id) in data_file.null_value_counts().keys().enumerate() {
let expect = batch.column(index).null_count() as u64;
assert_eq!(*data_file.null_value_counts.get(id).unwrap(), expect);
}
let split_offsets = data_file
.split_offsets
.as_ref()
.expect("split_offsets should be set");
assert_eq!(split_offsets.len(), metadata.num_row_groups());
split_offsets.iter().enumerate().for_each(|(i, &v)| {
let expect = metadata.row_groups()[i].file_offset().unwrap();
assert_eq!(v, expect);
});
}
#[tokio::test]
async fn test_equality_delete_writer() -> Result<(), anyhow::Error> {
let temp_dir = TempDir::new().unwrap();
let file_io = FileIO::new_with_fs();
let location_gen = DefaultLocationGenerator::with_data_location(
temp_dir.path().to_str().unwrap().to_string(),
);
let file_name_gen =
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
// prepare data
// Int, Struct(Int), String, List(Int), Struct(Struct(Int))
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(0, "col0", Type::Primitive(PrimitiveType::Int)).into(),
NestedField::required(
1,
"col1",
Type::Struct(StructType::new(vec![
NestedField::required(5, "sub_col", Type::Primitive(PrimitiveType::Int))
.into(),
])),
)
.into(),
NestedField::required(2, "col2", Type::Primitive(PrimitiveType::String)).into(),
NestedField::required(
3,
"col3",
Type::List(ListType::new(
NestedField::required(6, "element", Type::Primitive(PrimitiveType::Int))
.into(),
)),
)
.into(),
NestedField::required(
4,
"col4",
Type::Struct(StructType::new(vec![
NestedField::required(
7,
"sub_col",
Type::Struct(StructType::new(vec![
NestedField::required(
8,
"sub_sub_col",
Type::Primitive(PrimitiveType::Int),
)
.into(),
])),
)
.into(),
])),
)
.into(),
])
.build()
.unwrap();
let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap());
let col0 = Arc::new(Int32Array::from_iter_values(vec![1; 1024])) as ArrayRef;
let col1 = Arc::new(StructArray::new(
if let DataType::Struct(fields) = arrow_schema.fields.get(1).unwrap().data_type() {
fields.clone()
} else {
unreachable!()
},
vec![Arc::new(Int32Array::from_iter_values(vec![1; 1024]))],
None,
));
let col2 = Arc::new(arrow_array::StringArray::from_iter_values(vec![
"test";
1024
])) as ArrayRef;
let col3 = Arc::new({
let list_parts = arrow_array::ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
Some(
vec![Some(1),]
);
1024
])
.into_parts();
arrow_array::ListArray::new(
if let DataType::List(field) = arrow_schema.fields.get(3).unwrap().data_type() {
field.clone()
} else {
unreachable!()
},
list_parts.1,
list_parts.2,
list_parts.3,
)
}) as ArrayRef;
let col4 = Arc::new(StructArray::new(
if let DataType::Struct(fields) = arrow_schema.fields.get(4).unwrap().data_type() {
fields.clone()
} else {
unreachable!()
},
vec![Arc::new(StructArray::new(
if let DataType::Struct(fields) = arrow_schema.fields.get(4).unwrap().data_type() {
if let DataType::Struct(fields) = fields.first().unwrap().data_type() {
fields.clone()
} else {
unreachable!()
}
} else {
unreachable!()
},
vec![Arc::new(Int32Array::from_iter_values(vec![1; 1024]))],
None,
))],
None,
));
let columns = vec![col0, col1, col2, col3, col4];
let to_write = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap();
let equality_ids = vec![0_i32, 8];
let equality_config =
EqualityDeleteWriterConfig::new(equality_ids, Arc::new(schema)).unwrap();
let delete_schema =
arrow_schema_to_schema(equality_config.projected_arrow_schema_ref()).unwrap();
let projector = equality_config.projector.clone();
// prepare writer
let pb =
ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(delete_schema));
let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
pb,
file_io.clone(),
location_gen,
file_name_gen,
);
let mut equality_delete_writer =
EqualityDeleteFileWriterBuilder::new(rolling_writer_builder, equality_config)
.build(None)
.await?;
// write
equality_delete_writer.write(to_write.clone()).await?;
let res = equality_delete_writer.close().await?;
assert_eq!(res.len(), 1);
let data_file = res.into_iter().next().unwrap();
// check
let to_write_projected = projector.project_batch(to_write)?;
check_parquet_data_file_with_equality_delete_write(
&file_io,
&data_file,
&to_write_projected,
)
.await;
Ok(())
}
fn equality_id_validation_schema() -> Arc<Schema> {
Arc::new(
Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
NestedField::optional(2, "name", Type::Primitive(PrimitiveType::String)).into(),
])
.build()
.unwrap(),
)
}
#[test]
fn test_equality_delete_rejects_empty_equality_ids() {
let err =
EqualityDeleteWriterConfig::new(vec![], equality_id_validation_schema()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::DataInvalid);
assert!(
err.to_string()
.contains("Equality delete field ids must not be empty."),
"{err}"
);
}
#[test]
fn test_equality_delete_rejects_duplicate_equality_ids() {
let err = EqualityDeleteWriterConfig::new(vec![1, 1], equality_id_validation_schema())
.unwrap_err();
assert_eq!(err.kind(), ErrorKind::DataInvalid);
assert!(
err.to_string()
.contains("Duplicate equality delete field id: 1"),
"{err}"
);
assert!(err.to_string().contains("field_id: 1"), "{err}");
}
#[test]
fn test_equality_delete_rejects_missing_equality_id() {
let err =
EqualityDeleteWriterConfig::new(vec![99], equality_id_validation_schema()).unwrap_err();
assert_eq!(err.kind(), ErrorKind::DataInvalid);
assert!(
err.to_string()
.contains("Invalid equality delete field id: 99"),
"{err}"
);
assert!(err.to_string().contains("field_id: 99"), "{err}");
}
#[tokio::test]
async fn test_equality_delete_unreachable_column() -> Result<(), anyhow::Error> {
let schema = Arc::new(
Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(0, "col0", Type::Primitive(PrimitiveType::Float)).into(),
NestedField::required(1, "col1", Type::Primitive(PrimitiveType::Double)).into(),
NestedField::optional(2, "col2", Type::Primitive(PrimitiveType::Int)).into(),
NestedField::required(
3,
"col3",
Type::Struct(StructType::new(vec![
NestedField::required(
4,
"sub_col",
Type::Primitive(PrimitiveType::Int),
)
.into(),
])),
)
.into(),
NestedField::optional(
5,
"col4",
Type::Struct(StructType::new(vec![
NestedField::required(
6,
"sub_col2",
Type::Primitive(PrimitiveType::Int),
)
.into(),
])),
)
.into(),
NestedField::required(
7,
"col5",
Type::Map(MapType::new(
Arc::new(NestedField::required(
8,
"key",
Type::Primitive(PrimitiveType::String),
)),
Arc::new(NestedField::required(
9,
"value",
Type::Primitive(PrimitiveType::Int),
)),
)),
)
.into(),
NestedField::required(
10,
"col6",
Type::List(ListType::new(Arc::new(NestedField::required(
11,
"element",
Type::Primitive(PrimitiveType::Int),
)))),
)
.into(),
])
.build()
.unwrap(),
);
// Float and Double are not allowed to be used for equality delete
assert!(EqualityDeleteWriterConfig::new(vec![0], schema.clone()).is_err());
assert!(EqualityDeleteWriterConfig::new(vec![1], schema.clone()).is_err());
// Struct is not allowed to be used for equality delete
assert!(EqualityDeleteWriterConfig::new(vec![3], schema.clone()).is_err());
// Nested field of struct is allowed to be used for equality delete
assert!(EqualityDeleteWriterConfig::new(vec![4], schema.clone()).is_ok());
// Nested field of map is not allowed to be used for equality delete
assert!(EqualityDeleteWriterConfig::new(vec![7], schema.clone()).is_err());
assert!(EqualityDeleteWriterConfig::new(vec![8], schema.clone()).is_err());
assert!(EqualityDeleteWriterConfig::new(vec![9], schema.clone()).is_err());
// Nested field of list is not allowed to be used for equality delete
assert!(EqualityDeleteWriterConfig::new(vec![10], schema.clone()).is_err());
assert!(EqualityDeleteWriterConfig::new(vec![11], schema.clone()).is_err());
Ok(())
}
#[tokio::test]
async fn test_equality_delete_with_primitive_type() -> Result<(), anyhow::Error> {
let temp_dir = TempDir::new().unwrap();
let file_io = FileIO::new_with_fs();
let location_gen = DefaultLocationGenerator::with_data_location(
temp_dir.path().to_str().unwrap().to_string(),
);
let file_name_gen =
DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
let schema = Arc::new(
Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::required(0, "col0", Type::Primitive(PrimitiveType::Boolean))
.into(),
NestedField::required(1, "col1", Type::Primitive(PrimitiveType::Int)).into(),
NestedField::required(2, "col2", Type::Primitive(PrimitiveType::Long)).into(),
NestedField::required(
3,
"col3",
Type::Primitive(PrimitiveType::Decimal {
precision: 38,
scale: 5,
}),
)
.into(),
NestedField::required(4, "col4", Type::Primitive(PrimitiveType::Date)).into(),
NestedField::required(5, "col5", Type::Primitive(PrimitiveType::Time)).into(),
NestedField::required(6, "col6", Type::Primitive(PrimitiveType::Timestamp))
.into(),
NestedField::required(7, "col7", Type::Primitive(PrimitiveType::Timestamptz))
.into(),
NestedField::required(8, "col8", Type::Primitive(PrimitiveType::TimestampNs))
.into(),
NestedField::required(9, "col9", Type::Primitive(PrimitiveType::TimestamptzNs))
.into(),
NestedField::required(10, "col10", Type::Primitive(PrimitiveType::String))
.into(),
NestedField::required(11, "col11", Type::Primitive(PrimitiveType::Uuid)).into(),
NestedField::required(12, "col12", Type::Primitive(PrimitiveType::Fixed(10)))
.into(),
NestedField::required(13, "col13", Type::Primitive(PrimitiveType::Binary))
.into(),
])
.build()
.unwrap(),
);
let equality_ids = vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
let config = EqualityDeleteWriterConfig::new(equality_ids, schema.clone()).unwrap();
let delete_arrow_schema = config.projected_arrow_schema_ref().clone();
let delete_schema = arrow_schema_to_schema(&delete_arrow_schema).unwrap();
let pb =
ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(delete_schema));
let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
pb,
file_io.clone(),
location_gen,
file_name_gen,
);
let mut equality_delete_writer =
EqualityDeleteFileWriterBuilder::new(rolling_writer_builder, config)
.build(None)
.await?;
// prepare data
let col0 = Arc::new(BooleanArray::from(vec![
Some(true),
Some(false),
Some(true),
])) as ArrayRef;
let col1 = Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(4)])) as ArrayRef;
let col2 = Arc::new(Int64Array::from(vec![Some(1), Some(2), Some(4)])) as ArrayRef;
let col3 = Arc::new(
arrow_array::Decimal128Array::from(vec![Some(1), Some(2), Some(4)])
.with_precision_and_scale(38, 5)
.unwrap(),
) as ArrayRef;
let col4 = Arc::new(arrow_array::Date32Array::from(vec![
Some(0),
Some(1),
Some(3),
])) as ArrayRef;
let col5 = Arc::new(arrow_array::Time64MicrosecondArray::from(vec![
Some(0),
Some(1),
Some(3),
])) as ArrayRef;
let col6 = Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![
Some(0),
Some(1),
Some(3),
])) as ArrayRef;
let col7 = Arc::new(
arrow_array::TimestampMicrosecondArray::from(vec![Some(0), Some(1), Some(3)])
.with_timezone_utc(),
) as ArrayRef;
let col8 = Arc::new(arrow_array::TimestampNanosecondArray::from(vec![
Some(0),
Some(1),
Some(3),
])) as ArrayRef;
let col9 = Arc::new(
arrow_array::TimestampNanosecondArray::from(vec![Some(0), Some(1), Some(3)])
.with_timezone_utc(),
) as ArrayRef;
let col10 = Arc::new(arrow_array::StringArray::from(vec![
Some("a"),
Some("b"),
Some("d"),
])) as ArrayRef;
let col11 = Arc::new(
arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
vec![
Some(Uuid::from_u128(0).as_bytes().to_vec()),
Some(Uuid::from_u128(1).as_bytes().to_vec()),
Some(Uuid::from_u128(3).as_bytes().to_vec()),
]
.into_iter(),
16,
)
.unwrap(),
) as ArrayRef;
let col12 = Arc::new(
arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
vec![
Some(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
Some(vec![11, 12, 13, 14, 15, 16, 17, 18, 19, 20]),
Some(vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30]),
]
.into_iter(),
10,
)
.unwrap(),
) as ArrayRef;
let col13 = Arc::new(arrow_array::LargeBinaryArray::from_opt_vec(vec![
Some(b"one"),
Some(b""),
Some(b"zzzz"),
])) as ArrayRef;
let to_write = RecordBatch::try_new(delete_arrow_schema.clone(), vec![
col0, col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12, col13,
])
.unwrap();
equality_delete_writer.write(to_write.clone()).await?;
let res = equality_delete_writer.close().await?;
assert_eq!(res.len(), 1);
check_parquet_data_file_with_equality_delete_write(
&file_io,
&res.into_iter().next().unwrap(),
&to_write,
)
.await;
Ok(())
}
#[tokio::test]
async fn test_equality_delete_with_nullable_field() -> Result<(), anyhow::Error> {
// prepare data
// Int, Struct(Int), Struct(Struct(Int))
let schema = Schema::builder()
.with_schema_id(1)
.with_fields(vec![
NestedField::optional(0, "col0", Type::Primitive(PrimitiveType::Int)).into(),
NestedField::optional(
1,
"col1",
Type::Struct(StructType::new(vec![
NestedField::optional(2, "sub_col", Type::Primitive(PrimitiveType::Int))
.into(),
])),
)
.into(),
NestedField::optional(
3,
"col2",
Type::Struct(StructType::new(vec![
NestedField::optional(
4,
"sub_struct_col",
Type::Struct(StructType::new(vec![
NestedField::optional(
5,
"sub_sub_col",
Type::Primitive(PrimitiveType::Int),
)
.into(),
])),
)
.into(),
])),
)
.into(),
])
.build()
.unwrap();
let arrow_schema = Arc::new(schema_to_arrow_schema(&schema).unwrap());
// null 1 null(struct)
// 2 null(struct) null(sub_struct_col)
// 3 null(field) null(sub_sub_col)
let col0 = Arc::new(Int32Array::from(vec![None, Some(2), Some(3)])) as ArrayRef;
let col1 = {
let nulls = NullBuffer::from(vec![true, false, true]);
Arc::new(StructArray::new(
if let DataType::Struct(fields) = arrow_schema.fields.get(1).unwrap().data_type() {
fields.clone()
} else {
unreachable!()
},
vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), None]))],
Some(nulls),
))
};
let col2 = {
let inner_col = {
let nulls = NullBuffer::from(vec![true, false, true]);
Arc::new(StructArray::new(
Fields::from(vec![
Field::new("sub_sub_col", DataType::Int32, true).with_metadata(
HashMap::from([(
PARQUET_FIELD_ID_META_KEY.to_string(),
"5".to_string(),
)]),
),
]),
vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), None]))],
Some(nulls),
))
};
let nulls = NullBuffer::from(vec![false, true, true]);
Arc::new(StructArray::new(
if let DataType::Struct(fields) = arrow_schema.fields.get(2).unwrap().data_type() {
fields.clone()
} else {
unreachable!()
},
vec![inner_col],
Some(nulls),
))
};
let columns = vec![col0, col1, col2];
let to_write = RecordBatch::try_new(arrow_schema.clone(), columns).unwrap();
let equality_ids = vec![0_i32, 2, 5];
let equality_config =
EqualityDeleteWriterConfig::new(equality_ids, Arc::new(schema)).unwrap();
let projector = equality_config.projector.clone();
// check
let to_write_projected = projector.project_batch(to_write)?;
let expect_batch =
RecordBatch::try_new(equality_config.projected_arrow_schema_ref().clone(), vec![
Arc::new(Int32Array::from(vec![None, Some(2), Some(3)])) as ArrayRef,
Arc::new(Int32Array::from(vec![Some(1), None, None])) as ArrayRef,
Arc::new(Int32Array::from(vec![None, None, None])) as ArrayRef,
])
.unwrap();
assert_eq!(to_write_projected, expect_batch);
Ok(())
}
}