adriangb commented on code in PR #16461:
URL: https://github.com/apache/datafusion/pull/16461#discussion_r2157528156
##########
datafusion/datasource-parquet/src/opener.rs:
##########
@@ -879,4 +972,107 @@ mod test {
assert_eq!(num_batches, 0);
assert_eq!(num_rows, 0);
}
+
+ #[tokio::test]
+ async fn test_prune_on_partition_value_and_data_value() {
+ let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+
+ // Note: number 3 is missing!
+ let batch = record_batch!(("a", Int32, vec![Some(1), Some(2),
Some(4)])).unwrap();
+ let data_size =
+ write_parquet(Arc::clone(&store), "part=1/file.parquet",
batch.clone()).await;
+
+ let file_schema = batch.schema();
+ let mut file = PartitionedFile::new(
+ "part=1/file.parquet".to_string(),
+ u64::try_from(data_size).unwrap(),
+ );
+ file.partition_values = vec![ScalarValue::Int32(Some(1))];
+
+ let table_schema = Arc::new(Schema::new(vec![
+ Field::new("part", DataType::Int32, false),
+ Field::new("a", DataType::Int32, false),
+ ]));
+
+ let make_opener = |predicate| {
+ ParquetOpener {
+ partition_index: 0,
+ projection: Arc::new([0]),
+ batch_size: 1024,
+ limit: None,
+ predicate: Some(predicate),
+ logical_file_schema: file_schema.clone(),
+ metadata_size_hint: None,
+ metrics: ExecutionPlanMetricsSet::new(),
+ parquet_file_reader_factory: Arc::new(
+ DefaultParquetFileReaderFactory::new(Arc::clone(&store)),
+ ),
+ partition_fields: vec![Arc::new(Field::new(
+ "part",
+ DataType::Int32,
+ false,
+ ))],
+ pushdown_filters: true, // note that this is true!
+ reorder_filters: true,
+ enable_page_index: false,
+ enable_bloom_filter: false,
+ schema_adapter_factory: Arc::new(DefaultSchemaAdapterFactory),
+ enable_row_group_stats_pruning: false, // note that this is
false!
+ coerce_int96: None,
+ }
+ };
+
+ let make_meta = || FileMeta {
+ object_meta: ObjectMeta {
+ location: Path::from("part=1/file.parquet"),
+ last_modified: Utc::now(),
+ size: u64::try_from(data_size).unwrap(),
+ e_tag: None,
+ version: None,
+ },
+ range: None,
+ extensions: None,
+ metadata_size_hint: None,
+ };
+
+ // Filter should match the partition value and data value
+ let expr = col("part").eq(lit(1)).or(col("a").eq(lit(1)));
+ let predicate = logical2physical(&expr, &table_schema);
+ let opener = make_opener(predicate);
+ let stream = opener
+ .open(make_meta(), file.clone())
+ .unwrap()
+ .await
+ .unwrap();
+ let (num_batches, num_rows) = count_batches_and_rows(stream).await;
+ assert_eq!(num_batches, 1);
+ assert_eq!(num_rows, 3);
+
+ // Filter should match the partition value but not the data value
+ let expr = col("part").eq(lit(1)).or(col("a").eq(lit(3)));
+ let predicate = logical2physical(&expr, &table_schema);
+ let opener = make_opener(predicate);
+ let stream = opener.open(make_meta(),
file.clone()).unwrap().await.unwrap();
+ let (num_batches, num_rows) = count_batches_and_rows(stream).await;
+ assert_eq!(num_batches, 1);
+ assert_eq!(num_rows, 3);
+
+ // Filter should not match the partition value but match the data value
+ let expr = col("part").eq(lit(2)).or(col("a").eq(lit(1)));
+ let predicate = logical2physical(&expr, &table_schema);
+ let opener = make_opener(predicate);
+ let stream = opener.open(make_meta(),
file.clone()).unwrap().await.unwrap();
+ let (num_batches, num_rows) = count_batches_and_rows(stream).await;
+ assert_eq!(num_batches, 1);
+ assert_eq!(num_rows, 1);
Review Comment:
This assertion fails on `main`: all 3 rows are passed because the row filter
cannot handle the partition columns. This PR somewhat coincidentally happens to
fix this!
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
