This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new f38283b49b2 test: Add a test for RowFilter with nested type (#5600)
f38283b49b2 is described below
commit f38283b49b29f77e1bb2b0b2af07718724db3285
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Tue Apr 9 03:33:04 2024 -0700
test: Add a test for RowFilter with nested type (#5600)
---
parquet/src/arrow/async_reader/mod.rs | 88 +++++++++++++++++++++++++++++++++++
1 file changed, 88 insertions(+)
diff --git a/parquet/src/arrow/async_reader/mod.rs
b/parquet/src/arrow/async_reader/mod.rs
index 80a554026d9..c8a9f82c32c 100644
--- a/parquet/src/arrow/async_reader/mod.rs
+++ b/parquet/src/arrow/async_reader/mod.rs
@@ -1857,4 +1857,92 @@ mod tests {
assert_eq!(total_rows, expected);
}
}
+
+ #[tokio::test]
+ async fn test_row_filter_nested() {
+ let a = StringArray::from_iter_values(["a", "b", "b", "b", "c", "c"]);
+ let b = StructArray::from(vec![
+ (
+ Arc::new(Field::new("aa", DataType::Utf8, true)),
+ Arc::new(StringArray::from(vec!["a", "b", "b", "b", "c",
"c"])) as ArrayRef,
+ ),
+ (
+ Arc::new(Field::new("bb", DataType::Utf8, true)),
+ Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5",
"6"])) as ArrayRef,
+ ),
+ ]);
+ let c = Int32Array::from_iter(0..6);
+ let data = RecordBatch::try_from_iter([
+ ("a", Arc::new(a) as ArrayRef),
+ ("b", Arc::new(b) as ArrayRef),
+ ("c", Arc::new(c) as ArrayRef),
+ ])
+ .unwrap();
+
+ let mut buf = Vec::with_capacity(1024);
+ let mut writer = ArrowWriter::try_new(&mut buf, data.schema(),
None).unwrap();
+ writer.write(&data).unwrap();
+ writer.close().unwrap();
+
+ let data: Bytes = buf.into();
+ let metadata = parse_metadata(&data).unwrap();
+ let parquet_schema = metadata.file_metadata().schema_descr_ptr();
+
+ let test = TestReader {
+ data,
+ metadata: Arc::new(metadata),
+ requests: Default::default(),
+ };
+ let requests = test.requests.clone();
+
+ let a_scalar = StringArray::from_iter_values(["b"]);
+ let a_filter = ArrowPredicateFn::new(
+ ProjectionMask::leaves(&parquet_schema, vec![0]),
+ move |batch| eq(batch.column(0), &Scalar::new(&a_scalar)),
+ );
+
+ let b_scalar = StringArray::from_iter_values(["4"]);
+ let b_filter = ArrowPredicateFn::new(
+ ProjectionMask::leaves(&parquet_schema, vec![2]),
+ move |batch| {
+ // Filter on the second element of the struct.
+ let struct_array = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<StructArray>()
+ .unwrap();
+ eq(struct_array.column(0), &Scalar::new(&b_scalar))
+ },
+ );
+
+ let filter = RowFilter::new(vec![Box::new(a_filter),
Box::new(b_filter)]);
+
+ let mask = ProjectionMask::leaves(&parquet_schema, vec![0, 3]);
+ let stream = ParquetRecordBatchStreamBuilder::new(test)
+ .await
+ .unwrap()
+ .with_projection(mask.clone())
+ .with_batch_size(1024)
+ .with_row_filter(filter)
+ .build()
+ .unwrap();
+
+ let batches: Vec<_> = stream.try_collect().await.unwrap();
+ assert_eq!(batches.len(), 1);
+
+ let batch = &batches[0];
+ assert_eq!(batch.num_rows(), 1);
+ assert_eq!(batch.num_columns(), 2);
+
+ let col = batch.column(0);
+ let val = col.as_any().downcast_ref::<StringArray>().unwrap().value(0);
+ assert_eq!(val, "b");
+
+ let col = batch.column(1);
+ let val = col.as_any().downcast_ref::<Int32Array>().unwrap().value(0);
+ assert_eq!(val, 3);
+
+ // Should only have made 3 requests
+ assert_eq!(requests.lock().unwrap().len(), 3);
+ }
}
