hhhizzz commented on code in PR #10823:
URL: https://github.com/apache/arrow-rs/pull/10823#discussion_r3889617573
##########
parquet/src/arrow/async_reader/mod.rs:
##########
@@ -1260,59 +1262,146 @@ mod tests {
#[tokio::test]
async fn test_fuzz_async_reader_selection() {
+ const ROW_COUNT: usize = 7_300;
+ const CONSTRAINED_CACHE_SIZE: usize = 1024;
+ const AMPLE_CACHE_SIZE: usize = 1024 * 1024;
+ const CACHE_SIZES: [usize; 3] = [0, CONSTRAINED_CACHE_SIZE,
AMPLE_CACHE_SIZE];
+ const BATCH_SIZES: [usize; 5] = [1, 7, 32, 128, 1024];
+ const POLICIES: [RowSelectionPolicy; 3] = [
+ RowSelectionPolicy::Selectors,
+ RowSelectionPolicy::Mask,
+ RowSelectionPolicy::Auto { threshold: 32 },
+ ];
+ const ITERATIONS: usize = CACHE_SIZES.len() * POLICIES.len() * 2;
+
let testdata = arrow::util::test_util::parquet_test_data();
let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
let data = Bytes::from(std::fs::read(path).unwrap());
+ let metadata = ParquetMetaDataReader::new()
+ .with_page_index_policy(PageIndexPolicy::Required)
+ .parse_and_finish(&data)
+ .unwrap();
+ let options =
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
+ let metadata = ArrowReaderMetadata::try_new(metadata.into(),
options).unwrap();
+ assert_eq!(metadata.metadata().num_row_groups(), 1);
+ assert_eq!(
+ metadata.metadata().file_metadata().num_rows(),
+ ROW_COUNT as i64
+ );
- let mut rand = rng();
+ let mut rand = StdRng::seed_from_u64(42);
+
+ for iteration in 0..ITERATIONS {
+ let batch_size = BATCH_SIZES[iteration % BATCH_SIZES.len()];
+ let sparse = iteration % 2 == 0;
- for _ in 0..100 {
- let mut expected_rows = 0;
let mut total_rows = 0;
- let mut skip = false;
+ let mut skip = rand.random_bool(0.5);
let mut selectors = vec![];
- while total_rows < 7300 {
- let row_count: usize = rand.random_range(1..100);
-
- let row_count = row_count.min(7300 - total_rows);
+ while total_rows < ROW_COUNT {
+ let max_run = match (sparse, skip) {
+ (true, false) => 3,
+ (true, true) => 63,
+ (false, _) => 99,
+ };
+ let row_count = rand.random_range(1..=max_run).min(ROW_COUNT -
total_rows);
selectors.push(RowSelector { row_count, skip });
-
total_rows += row_count;
- if !skip {
- expected_rows += row_count;
- }
-
skip = !skip;
}
let selection = RowSelection::from(selectors);
+ let predicate_modulus = 2 + iteration as i32 % 7;
+ let predicate_remainder = iteration as i32 % predicate_modulus;
+ let output_column = rand.random_range(0..13);
+ let projection_columns = if output_column == 0 {
+ vec![0]
+ } else {
+ vec![0, output_column]
+ };
- let async_reader = TestReader::new(data.clone());
-
- let options =
-
ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
- let builder =
ParquetRecordBatchStreamBuilder::new_with_options(async_reader, options)
- .await
- .unwrap();
-
- assert_eq!(builder.metadata().num_row_groups(), 1);
-
- let col_idx: usize = rand.random_range(0..13);
- let mask = ProjectionMask::leaves(builder.parquet_schema(),
vec![col_idx]);
-
- let stream = builder
- .with_projection(mask.clone())
- .with_row_selection(selection.clone())
- .build()
- .expect("building stream");
+ let build_stream = |policy, cache_size, metrics| {
+ let builder =
ParquetRecordBatchStreamBuilder::new_with_metadata(
+ TestReader::new(data.clone()),
+ metadata.clone(),
+ );
+ let filter_projection =
ProjectionMask::leaves(builder.parquet_schema(), [0]);
+ let predicate =
+ ArrowPredicateFn::new(filter_projection, move |batch:
RecordBatch| {
+ let ids = batch.column(0).as_primitive::<Int32Type>();
+
Ok(BooleanArray::from_iter(ids.values().iter().map(|value| {
+ Some(value.rem_euclid(predicate_modulus) ==
predicate_remainder)
+ })))
+ });
+ let output_projection =
+ ProjectionMask::leaves(builder.parquet_schema(),
projection_columns.clone());
+
+ builder
+ .with_projection(output_projection)
+ .with_batch_size(batch_size)
+ .with_row_selection(selection.clone())
+ .with_row_filter(RowFilter::new(vec![Box::new(predicate)]))
+ .with_row_selection_policy(policy)
+ .with_max_predicate_cache_size(cache_size)
+ .with_metrics(metrics)
+ .build()
+ .unwrap()
+ };
- let async_batches: Vec<_> = stream.try_collect().await.unwrap();
+ let reference = build_stream(
+ RowSelectionPolicy::Selectors,
+ 0,
+ ArrowReaderMetrics::disabled(),
+ );
+ let schema = reference.schema().clone();
+ let expected: Vec<_> =
reference.try_collect().await.unwrap_or_else(|error| {
+ panic!(
+ "reference failed: iteration={iteration},
output_column={output_column}, \
+ batch_size={batch_size}: {error}"
+ )
+ });
+ let expected = concat_batches(&schema, &expected).unwrap();
+
+ for (policy_idx, policy) in POLICIES.iter().copied().enumerate() {
+ let cache_size = CACHE_SIZES[(iteration + policy_idx) %
CACHE_SIZES.len()];
Review Comment:
Could we make the constrained-cache mode observable here?
We assert the two endpoints: a zero-sized cache produces no cache reads, and
`AMPLE_CACHE_SIZE` produces some cache reads. However, nothing currently
verifies that `CONSTRAINED_CACHE_SIZE` actually fills the shared cache, rejects
later insertions, and causes the consumer to fall back to the inner reader.
Because the cache size is rotated across different generated selections, we
also cannot compare the 1 KiB and 1 MiB behavior on the same input. As written,
this test would still pass if the 1 KiB configuration behaved exactly like the
ample cache, or if it admitted no shared batches at all.
Since the constrained-cache/fallback interaction is one of the explicit axes
in #10747, could we run at least one identical selection/policy/batch-size case
with both constrained and ample limits, and assert that the constrained case
has fewer cache reads and/or more inner reads than the ample case? Ideally,
that case would admit some but not all batches, so both the cache-hit and
fallback paths are exercised.
##########
parquet/src/arrow/async_reader/mod.rs:
##########
@@ -1260,59 +1262,146 @@ mod tests {
#[tokio::test]
async fn test_fuzz_async_reader_selection() {
+ const ROW_COUNT: usize = 7_300;
+ const CONSTRAINED_CACHE_SIZE: usize = 1024;
+ const AMPLE_CACHE_SIZE: usize = 1024 * 1024;
+ const CACHE_SIZES: [usize; 3] = [0, CONSTRAINED_CACHE_SIZE,
AMPLE_CACHE_SIZE];
+ const BATCH_SIZES: [usize; 5] = [1, 7, 32, 128, 1024];
+ const POLICIES: [RowSelectionPolicy; 3] = [
+ RowSelectionPolicy::Selectors,
+ RowSelectionPolicy::Mask,
+ RowSelectionPolicy::Auto { threshold: 32 },
+ ];
+ const ITERATIONS: usize = CACHE_SIZES.len() * POLICIES.len() * 2;
Review Comment:
This replaces the previous 100 generated selections with 18 deterministic
selections. The stronger oracle and interaction matrix are clearly more
valuable, but they cover a different dimension and substantially reduce the
number of selection shapes being sampled.
It may be worth retaining the old cheap selection-only coverage separately,
or otherwise decoupling the number of generated selections from the
policy/cache schedule, so that adding the interaction matrix does not also
shrink the original selection-shape fuzzing.
--
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]
