This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 2c6eca0 test: cover Mosaic read path combinations (#401)
2c6eca0 is described below
commit 2c6eca0c955906d97068465ff24208360b888935
Author: QuakeWang <[email protected]>
AuthorDate: Mon Jun 22 17:18:57 2026 +0800
test: cover Mosaic read path combinations (#401)
---
crates/paimon/src/arrow/format/mosaic.rs | 461 +++++++++++++++++++++++++++-
crates/paimon/src/table/data_file_reader.rs | 301 +++++++++++++++++-
2 files changed, 757 insertions(+), 5 deletions(-)
diff --git a/crates/paimon/src/arrow/format/mosaic.rs
b/crates/paimon/src/arrow/format/mosaic.rs
index f631fad..963226b 100644
--- a/crates/paimon/src/arrow/format/mosaic.rs
+++ b/crates/paimon/src/arrow/format/mosaic.rs
@@ -532,10 +532,15 @@ mod tests {
use super::*;
use crate::arrow::format::{FilePredicates, FormatFileReader};
use crate::spec::{
- ArrayType, DataType, Datum, IntType, Predicate, PredicateBuilder,
RowType, TimestampType,
- VarCharType,
+ ArrayType, BigIntType, BooleanType, DataType, DateType, Datum,
DecimalType, DoubleType,
+ FloatType, IntType, LocalZonedTimestampType, Predicate,
PredicateBuilder, RowType,
+ SmallIntType, TimeType, TimestampType, TinyIntType, VarBinaryType,
VarCharType,
+ };
+ use arrow_array::{
+ Array, BinaryArray, BooleanArray, Date32Array, Decimal128Array,
Float32Array, Float64Array,
+ Int16Array, Int32Array, Int64Array, Int8Array, StringArray,
Time32MillisecondArray,
+ TimestampMicrosecondArray, TimestampMillisecondArray,
TimestampNanosecondArray,
};
- use arrow_array::{Array, Int32Array, StringArray,
TimestampMicrosecondArray};
use arrow_schema::{DataType as ArrowDataType, Field, Schema};
use bytes::Bytes;
use futures::TryStreamExt;
@@ -1154,4 +1159,454 @@ mod tests {
matches!(err, Error::Unsupported { message } if
message.contains("Mosaic format does not support column 'nested'"))
);
}
+
+ fn full_type_fields() -> Vec<DataField> {
+ vec![
+ field(
+ 0,
+ "f_bool",
+ DataType::Boolean(BooleanType::with_nullable(true)),
+ ),
+ field(
+ 1,
+ "f_tinyint",
+ DataType::TinyInt(TinyIntType::with_nullable(true)),
+ ),
+ field(
+ 2,
+ "f_smallint",
+ DataType::SmallInt(SmallIntType::with_nullable(true)),
+ ),
+ field(3, "f_int", DataType::Int(IntType::with_nullable(false))),
+ field(
+ 4,
+ "f_bigint",
+ DataType::BigInt(BigIntType::with_nullable(true)),
+ ),
+ field(
+ 5,
+ "f_float",
+ DataType::Float(FloatType::with_nullable(true)),
+ ),
+ field(
+ 6,
+ "f_double",
+ DataType::Double(DoubleType::with_nullable(true)),
+ ),
+ field(7, "f_date", DataType::Date(DateType::with_nullable(true))),
+ field(
+ 8,
+ "f_time",
+ DataType::Time(TimeType::with_nullable(true, 3).unwrap()),
+ ),
+ field(
+ 9,
+ "f_string",
+ DataType::VarChar(VarCharType::with_nullable(true,
20).unwrap()),
+ ),
+ field(
+ 10,
+ "f_binary",
+ DataType::VarBinary(VarBinaryType::try_new(true, 20).unwrap()),
+ ),
+ field(
+ 11,
+ "f_decimal_compact",
+ DataType::Decimal(DecimalType::with_nullable(true, 5,
2).unwrap()),
+ ),
+ field(
+ 12,
+ "f_decimal_large",
+ DataType::Decimal(DecimalType::with_nullable(true, 20,
0).unwrap()),
+ ),
+ field(
+ 13,
+ "f_ts3",
+ DataType::Timestamp(TimestampType::with_nullable(true,
3).unwrap()),
+ ),
+ field(
+ 14,
+ "f_ts6",
+ DataType::Timestamp(TimestampType::with_nullable(true,
6).unwrap()),
+ ),
+ field(
+ 15,
+ "f_ts9",
+ DataType::Timestamp(TimestampType::with_nullable(true,
9).unwrap()),
+ ),
+ field(
+ 16,
+ "f_ltz",
+ DataType::LocalZonedTimestamp(
+ LocalZonedTimestampType::with_nullable(true, 6).unwrap(),
+ ),
+ ),
+ ]
+ }
+
+ /// Round-trips every scalar/temporal type Mosaic supports through write +
read,
+ /// asserting values survive the format. ARRAY/MAP are intentionally
excluded:
+ /// `paimon-mosaic-core` 0.1.0 does not support them and the reader
rejects them.
+ #[tokio::test]
+ async fn test_read_full_types() {
+ let fields = full_type_fields();
+ let schema = build_target_arrow_schema(&fields).unwrap();
+ let batch = RecordBatch::try_new(
+ schema,
+ vec![
+ Arc::new(BooleanArray::from(vec![Some(true), Some(false)])),
+ Arc::new(Int8Array::from(vec![Some(1i8), Some(-2)])),
+ Arc::new(Int16Array::from(vec![Some(100i16), Some(-200)])),
+ Arc::new(Int32Array::from(vec![10, 20])),
+ Arc::new(Int64Array::from(vec![Some(1_000i64), Some(-2_000)])),
+ Arc::new(Float32Array::from(vec![Some(1.5f32), Some(-2.5)])),
+ Arc::new(Float64Array::from(vec![Some(3.25f64), Some(-4.75)])),
+ Arc::new(Date32Array::from(vec![Some(18_000), Some(19_000)])),
+ Arc::new(Time32MillisecondArray::from(vec![
+ Some(3_600_000),
+ Some(7_200_000),
+ ])),
+ Arc::new(StringArray::from(vec![Some("hello"),
Some("mosaic")])),
+ Arc::new(BinaryArray::from_opt_vec(vec![
+ Some(b"ab".as_ref()),
+ Some(b"cd".as_ref()),
+ ])),
+ Arc::new(
+ Decimal128Array::from(vec![Some(12_345i128), Some(-678)])
+ .with_precision_and_scale(5, 2)
+ .unwrap(),
+ ),
+ Arc::new(
+
Decimal128Array::from(vec![Some(12_345_678_901_234_567_890i128), Some(-1)])
+ .with_precision_and_scale(20, 0)
+ .unwrap(),
+ ),
+ Arc::new(TimestampMillisecondArray::from(vec![
+ Some(1_700_000_000_000i64),
+ Some(-1),
+ ])),
+ Arc::new(TimestampMicrosecondArray::from(vec![
+ Some(1_700_000_000_000_000i64),
+ Some(-1),
+ ])),
+ Arc::new(TimestampNanosecondArray::from(vec![
+ Some(1_700_000_000_123_456_789i64),
+ Some(-1),
+ ])),
+ Arc::new(
+
TimestampMicrosecondArray::from(vec![Some(1_700_000_000_000_000i64), Some(-1)])
+ .with_timezone("UTC"),
+ ),
+ ],
+ )
+ .unwrap();
+
+ let data = write_mosaic(&batch);
+ let batches = read_batches(data, &fields, None).await.unwrap();
+
+ assert_eq!(batches.len(), 1);
+ let result = &batches[0];
+ assert_eq!(result.num_rows(), 2);
+ assert_eq!(result.num_columns(), fields.len());
+
+ let col = |i: usize| result.column(i);
+ let downcast = |i: usize| col(i).as_any();
+
+ assert!(downcast(0).downcast_ref::<BooleanArray>().unwrap().value(0));
+ assert!(!downcast(0).downcast_ref::<BooleanArray>().unwrap().value(1));
+ assert_eq!(
+ downcast(1).downcast_ref::<Int8Array>().unwrap().value(1),
+ -2
+ );
+ assert_eq!(
+ downcast(2).downcast_ref::<Int16Array>().unwrap().value(0),
+ 100
+ );
+ assert_eq!(
+ downcast(3).downcast_ref::<Int32Array>().unwrap().values(),
+ &[10, 20]
+ );
+ assert_eq!(
+ downcast(4).downcast_ref::<Int64Array>().unwrap().value(1),
+ -2_000
+ );
+ assert_eq!(
+ downcast(5).downcast_ref::<Float32Array>().unwrap().value(0),
+ 1.5
+ );
+ assert_eq!(
+ downcast(6).downcast_ref::<Float64Array>().unwrap().value(1),
+ -4.75
+ );
+ assert_eq!(
+ downcast(7).downcast_ref::<Date32Array>().unwrap().value(0),
+ 18_000
+ );
+ assert_eq!(
+ downcast(8)
+ .downcast_ref::<Time32MillisecondArray>()
+ .unwrap()
+ .value(0),
+ 3_600_000
+ );
+ assert_eq!(
+ downcast(9).downcast_ref::<StringArray>().unwrap().value(1),
+ "mosaic"
+ );
+ assert_eq!(
+ downcast(10).downcast_ref::<BinaryArray>().unwrap().value(0),
+ b"ab"
+ );
+ assert_eq!(
+ downcast(11)
+ .downcast_ref::<Decimal128Array>()
+ .unwrap()
+ .value(0),
+ 12_345
+ );
+ assert_eq!(
+ downcast(12)
+ .downcast_ref::<Decimal128Array>()
+ .unwrap()
+ .value(0),
+ 12_345_678_901_234_567_890
+ );
+ assert_eq!(
+ downcast(13)
+ .downcast_ref::<TimestampMillisecondArray>()
+ .unwrap()
+ .value(0),
+ 1_700_000_000_000
+ );
+ assert_eq!(
+ downcast(14)
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap()
+ .value(0),
+ 1_700_000_000_000_000
+ );
+ assert_eq!(
+ downcast(15)
+ .downcast_ref::<TimestampNanosecondArray>()
+ .unwrap()
+ .value(0),
+ 1_700_000_000_123_456_789
+ );
+ assert_eq!(
+ downcast(16)
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .unwrap()
+ .value(0),
+ 1_700_000_000_000_000
+ );
+ assert_eq!(
+ result.schema().field(16).data_type(),
+ &ArrowDataType::Timestamp(TimeUnit::Microsecond,
Some("UTC".into()))
+ );
+ }
+
+ /// Null values in nullable columns must round-trip as nulls.
+ #[tokio::test]
+ async fn test_read_null_values() {
+ let fields = data_fields();
+ let batch = RecordBatch::try_new(
+ arrow_schema(),
+ vec![
+ Arc::new(Int32Array::from(vec![1, 2, 3])),
+ Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
+ Arc::new(Int32Array::from(vec![Some(10), None, None])),
+ ],
+ )
+ .unwrap();
+ let data = write_mosaic(&batch);
+ let batches = read_batches(data, &fields, None).await.unwrap();
+
+ assert_eq!(batches.len(), 1);
+ let result = &batches[0];
+ assert_eq!(result.num_rows(), 3);
+
+ let names = result
+ .column(1)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .unwrap();
+ assert_eq!(names.value(0), "a");
+ assert!(names.is_null(1));
+ assert_eq!(names.value(2), "c");
+
+ let scores = result
+ .column(2)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap();
+ assert_eq!(scores.value(0), 10);
+ assert_eq!(scores.null_count(), 2);
+ }
+
+ #[test]
+ fn test_mosaic_value_to_datum_conversions() {
+ let ts = |p| DataType::Timestamp(TimestampType::new(p).unwrap());
+ let ltz = |p|
DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(p).unwrap());
+
+ // Each supported variant maps to the matching Datum.
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::Boolean(true),
+ &DataType::Boolean(BooleanType::new())
+ ),
+ Some(Datum::Bool(true))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::TinyInt(7),
+ &DataType::TinyInt(TinyIntType::new())
+ ),
+ Some(Datum::TinyInt(7))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::SmallInt(-9),
+ &DataType::SmallInt(SmallIntType::new())
+ ),
+ Some(Datum::SmallInt(-9))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::Integer(42),
&DataType::Int(IntType::new())),
+ Some(Datum::Int(42))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::BigInt(1_000),
+ &DataType::BigInt(BigIntType::new())
+ ),
+ Some(Datum::Long(1_000))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::Float(1.5),
&DataType::Float(FloatType::new())),
+ Some(Datum::Float(1.5))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::Double(2.5),
+ &DataType::Double(DoubleType::new())
+ ),
+ Some(Datum::Double(2.5))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::Date(100),
&DataType::Date(DateType::new())),
+ Some(Datum::Date(100))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::Time(200),
+ &DataType::Time(TimeType::new(3).unwrap())
+ ),
+ Some(Datum::Time(200))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::String(b"hi".to_vec()),
+ &DataType::VarChar(VarCharType::new(20).unwrap())
+ ),
+ Some(Datum::String("hi".to_string()))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::Bytes(vec![1, 2]),
+ &DataType::VarBinary(VarBinaryType::try_new(true, 20).unwrap())
+ ),
+ Some(Datum::Bytes(vec![1, 2]))
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::DecimalCompact(1_000),
+ &DataType::Decimal(DecimalType::new(5, 2).unwrap())
+ ),
+ Some(Datum::Decimal {
+ unscaled: 1_000,
+ precision: 5,
+ scale: 2,
+ })
+ );
+
+ // Timestamp precision boundaries select the matching Mosaic encoding.
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::TimestampMillis(5), &ts(3)),
+ Some(Datum::Timestamp {
+ millis: 5,
+ nanos: 0
+ })
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::TimestampMillis(5), <z(3)),
+ Some(Datum::LocalZonedTimestamp {
+ millis: 5,
+ nanos: 0
+ })
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::TimestampMicros(1_500),
&ts(6)),
+ Some(Datum::Timestamp {
+ millis: 1,
+ nanos: 500_000,
+ })
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::TimestampNanos {
+ millis: 1,
+ nanos_of_milli: 2,
+ },
+ &ts(9)
+ ),
+ Some(Datum::Timestamp {
+ millis: 1,
+ nanos: 2
+ })
+ );
+
+ // Ambiguous or unsupported inputs must fail open (None).
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::Null,
&DataType::Int(IntType::new())),
+ None
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::Integer(1),
+ &DataType::BigInt(BigIntType::new())
+ ),
+ None,
+ "type mismatch must not convert"
+ );
+ assert_eq!(
+ mosaic_value_to_datum(
+ &MosaicValue::DecimalLarge(vec![0, 0]),
+ &DataType::Decimal(DecimalType::new(20, 0).unwrap())
+ ),
+ None,
+ "large decimal stats are not converted"
+ );
+ assert_eq!(
+ mosaic_value_to_datum(&MosaicValue::TimestampMillis(5), &ts(6)),
+ None,
+ "millis encoding must not satisfy a micros-precision type"
+ );
+ }
+
+ #[test]
+ fn test_split_batch() {
+ let chunks = split_batch(sample_batch(), 2);
+ assert_eq!(chunks.len(), 3);
+ assert_eq!(
+ chunks.iter().map(RecordBatch::num_rows).collect::<Vec<_>>(),
+ vec![2, 2, 1]
+ );
+
+ let unsplit = split_batch(sample_batch(), 0);
+ assert_eq!(unsplit.len(), 1);
+ assert_eq!(unsplit[0].num_rows(), 5);
+
+ let whole = split_batch(sample_batch(), 10);
+ assert_eq!(whole.len(), 1);
+ }
}
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index f08ad91..d5d9f3e 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -470,13 +470,16 @@ mod tests {
use crate::arrow::build_target_arrow_schema;
use crate::io::FileIOBuilder;
use crate::spec::stats::BinaryTableStats;
- use crate::spec::{ArrayType, DataFileMeta, DataType, IntType, VarCharType};
- use crate::table::source::DataSplitBuilder;
+ use crate::spec::{
+ ArrayType, DataFileMeta, DataType, Datum, IntType, Predicate,
PredicateBuilder, VarCharType,
+ };
+ use crate::table::source::{DataSplitBuilder, DeletionFile};
use arrow_array::{Int32Array, StringArray};
use bytes::Bytes;
use futures::TryStreamExt;
use paimon_mosaic_core::spec::COMPRESSION_NONE;
use paimon_mosaic_core::writer::{MosaicWriter, OutputFile, WriterOptions};
+ use roaring::RoaringBitmap;
use std::io;
struct MemOutputFile {
@@ -644,4 +647,298 @@ mod tests {
assert_eq!(names.value(0), "a");
assert_eq!(names.value(2), "c");
}
+
+ fn pk_fields() -> Vec<DataField> {
+ vec![
+ data_field(0, "id", DataType::Int(IntType::with_nullable(false))),
+ data_field(
+ 1,
+ "name",
+ DataType::VarChar(VarCharType::with_nullable(true,
20).unwrap()),
+ ),
+ ]
+ }
+
+ fn pk_batch(ids: Vec<i32>, names: Vec<&str>) -> RecordBatch {
+ RecordBatch::try_new(
+ build_target_arrow_schema(&pk_fields()).unwrap(),
+ vec![
+ Arc::new(Int32Array::from(ids)),
+ Arc::new(StringArray::from(names)),
+ ],
+ )
+ .unwrap()
+ }
+
+ fn write_multi_row_group_mosaic(batches: &[RecordBatch], stats_columns:
Vec<String>) -> Bytes {
+ let out = MemOutputFile::new();
+ let mut writer = MosaicWriter::new(
+ out,
+ batches[0].schema().as_ref(),
+ WriterOptions {
+ compression: COMPRESSION_NONE,
+ num_buckets: 2,
+ // One row group per written batch, so each batch carries its
own stats.
+ row_group_max_size: 1,
+ stats_columns,
+ ..Default::default()
+ },
+ )
+ .unwrap();
+ for batch in batches {
+ writer.write_batch(batch).unwrap();
+ }
+ writer.close().unwrap();
+ Bytes::from(writer.output().data.to_vec())
+ }
+
+ fn write_parquet(batch: &RecordBatch) -> Bytes {
+ let mut buf = Vec::new();
+ let mut writer =
+ parquet::arrow::ArrowWriter::try_new(&mut buf, batch.schema(),
None).unwrap();
+ writer.write(batch).unwrap();
+ writer.close().unwrap();
+ Bytes::from(buf)
+ }
+
+ /// Writes a Paimon deletion-vector blob and returns the `DeletionFile`
pointing at it.
+ /// Layout matches [`DeletionVector::read_from_bytes`]:
+ /// `i32 bitmapLength (magic + bitmap) | i32 magic | bitmap bytes | i32
crc`.
+ async fn write_deletion_file(
+ file_io: &crate::io::FileIO,
+ path: &str,
+ deleted_rows: &[u32],
+ ) -> DeletionFile {
+ // BitmapDeletionVector.MAGIC_NUMBER, see crate::deletion_vector.
+ const MAGIC_NUMBER: i32 = 1581511376;
+ let mut bitmap = RoaringBitmap::new();
+ for row in deleted_rows {
+ bitmap.insert(*row);
+ }
+ let mut bitmap_bytes = Vec::new();
+ bitmap.serialize_into(&mut bitmap_bytes).unwrap();
+
+ let bitmap_length = 4 + bitmap_bytes.len() as i32;
+ let mut blob = Vec::new();
+ blob.extend_from_slice(&bitmap_length.to_be_bytes());
+ blob.extend_from_slice(&MAGIC_NUMBER.to_be_bytes());
+ blob.extend_from_slice(&bitmap_bytes);
+ blob.extend_from_slice(&0i32.to_be_bytes()); // crc, skipped on read
+ file_io
+ .new_output(path)
+ .unwrap()
+ .write(Bytes::from(blob))
+ .await
+ .unwrap();
+
+ DeletionFile::new(
+ path.to_string(),
+ 0,
+ bitmap_length as i64,
+ Some(deleted_rows.len() as i64),
+ )
+ }
+
+ fn collect_ids(batches: &[RecordBatch]) -> Vec<i32> {
+ batches
+ .iter()
+ .flat_map(|b| {
+ b.column(0)
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect()
+ }
+
+ /// Deletion vectors are applied format-agnostically by `DataFileReader`;
verify a
+ /// Mosaic file honors deleted rows end to end.
+ #[tokio::test]
+ async fn test_mosaic_with_deletion_vector() {
+ let fields = pk_fields();
+ let data = write_mosaic(&pk_batch(vec![1, 2, 3], vec!["a", "b", "c"]));
+
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let table_path = "memory:/mosaic_dv";
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_name = "part-0.mosaic";
+ file_io
+ .new_output(&format!("{bucket_path}/{file_name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+ // Delete row index 1 (id = 2).
+ let dv = write_deletion_file(&file_io,
&format!("{table_path}/index/dv-0"), &[1]).await;
+
+ let table_schema_id = 1;
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file(
+ file_name,
+ data.len() as i64,
+ 3,
+ table_schema_id,
+ )])
+ .with_data_deletion_files(vec![Some(dv)])
+ .build()
+ .unwrap();
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ table_schema_id,
+ fields.clone(),
+ fields.clone(),
+ Vec::new(),
+ );
+ let batches = reader
+ .read(&[split])
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(collect_ids(&batches), vec![1, 3]);
+ }
+
+ /// A Mosaic file and a Parquet file in the same split must both be read
and concatenated.
+ #[tokio::test]
+ async fn test_mosaic_mixed_format_read() {
+ let fields = pk_fields();
+ let mosaic_data = write_mosaic(&pk_batch(vec![1, 2], vec!["a", "b"]));
+ let parquet_data = write_parquet(&pk_batch(vec![3, 4], vec!["c",
"d"]));
+
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let table_path = "memory:/mosaic_mixed";
+ let bucket_path = format!("{table_path}/bucket-0");
+ for (name, data) in [
+ ("part-0.mosaic", &mosaic_data),
+ ("part-1.parquet", &parquet_data),
+ ] {
+ file_io
+ .new_output(&format!("{bucket_path}/{name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+ }
+
+ let table_schema_id = 1;
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![
+ data_file(
+ "part-0.mosaic",
+ mosaic_data.len() as i64,
+ 2,
+ table_schema_id,
+ ),
+ data_file(
+ "part-1.parquet",
+ parquet_data.len() as i64,
+ 2,
+ table_schema_id,
+ ),
+ ])
+ .build()
+ .unwrap();
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ table_schema_id,
+ fields.clone(),
+ fields.clone(),
+ Vec::new(),
+ );
+ let batches = reader
+ .read(&[split])
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ let mut ids = collect_ids(&batches);
+ ids.sort_unstable();
+ assert_eq!(ids, vec![1, 2, 3, 4]);
+ }
+
+ /// Row-group predicate pruning, deletion vectors and projection must
compose correctly:
+ /// the predicate keeps one row group, the DV deletes one of its rows,
projection keeps `id`.
+ #[tokio::test]
+ async fn test_mosaic_predicate_dv_projection_combination() {
+ let fields = pk_fields();
+ let data = write_multi_row_group_mosaic(
+ &[
+ pk_batch(vec![1, 2], vec!["a", "b"]),
+ pk_batch(vec![10, 11], vec!["c", "d"]),
+ pk_batch(vec![20, 21], vec!["e", "f"]),
+ ],
+ vec!["id".to_string()],
+ );
+
+ let file_io = FileIOBuilder::new("memory").build().unwrap();
+ let table_path = "memory:/mosaic_combo";
+ let bucket_path = format!("{table_path}/bucket-0");
+ let file_name = "part-0.mosaic";
+ file_io
+ .new_output(&format!("{bucket_path}/{file_name}"))
+ .unwrap()
+ .write(data.clone())
+ .await
+ .unwrap();
+ // Delete global row index 2 (id = 10, first row of the second row
group).
+ let dv = write_deletion_file(&file_io,
&format!("{table_path}/index/dv-0"), &[2]).await;
+
+ let table_schema_id = 1;
+ let split = DataSplitBuilder::new()
+ .with_snapshot(1)
+ .with_partition(crate::spec::BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path(bucket_path)
+ .with_total_buckets(1)
+ .with_data_files(vec![data_file(
+ file_name,
+ data.len() as i64,
+ 6,
+ table_schema_id,
+ )])
+ .with_data_deletion_files(vec![Some(dv)])
+ .build()
+ .unwrap();
+
+ let predicate: Predicate = PredicateBuilder::new(&fields)
+ .equal("id", Datum::Int(10))
+ .unwrap();
+ let read_type = vec![fields[0].clone()];
+ let schema_manager = SchemaManager::new(file_io.clone(),
table_path.to_string());
+ let reader = DataFileReader::new(
+ file_io,
+ schema_manager,
+ table_schema_id,
+ fields.clone(),
+ read_type,
+ vec![predicate],
+ );
+ let batches = reader
+ .read(&[split])
+ .unwrap()
+ .try_collect::<Vec<_>>()
+ .await
+ .unwrap();
+
+ assert_eq!(batches.iter().map(|b| b.num_columns()).max(), Some(1));
+ assert_eq!(collect_ids(&batches), vec![11]);
+ }
}
