Skip to content

Commit fea0264

Browse files
committed
Address Copilot review: deduplicate read-plan machinery, restore input_file_name handling lost in rebase
Three fixes on top of the rebase onto latest main: 1. row_filter.rs carried a stale pre-#23638 copy of the column-resolution machinery (ParquetReadPlan, PushdownChecker, StructFieldAccess, build_filter_schema, etc.) that main has since moved into projection_read_plan.rs. Rebuilt the file on main's shared-module version, keeping only this PR's real additions (PrebuiltRowFilterCandidate, prebuild_row_filter_candidates, row_filter_from_prebuilt, and the RowFilterGenerator removal): 2194 → 1483 lines, -798 net. The stray "strinrg" typo lived in the deleted copy. 2. The rebase silently dropped main's input_file_name() handling in ParquetMorselizer (import, the per-file literal rewrite in open(), and the test_input_file_name_projection test) because this PR's restructure of open() overlapped the region #23638 later touched. Restored all three; the test passes again. 3. metrics: the row_filter_skipped_fully_matched doc claimed to count row groups, but the counter increments once per suppression event (a run of consecutive fully-matched RGs shares a single toggle). Doc now states the actual semantics. Verified: 158 datasource-parquet tests pass (with PARQUET_TEST_DATA), cargo doc clean, clippy clean on touched files.
1 parent 04d04a5 commit fea0264

3 files changed

Lines changed: 148 additions & 798 deletions

File tree

‎datafusion/datasource-parquet/src/metrics.rs‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,11 +61,16 @@ pub struct ParquetFileMetrics {
6161
/// the initial pruning but were proved unreachable mid-scan after the
6262
/// dynamic filter tightened.
6363
pub row_groups_pruned_dynamic_filter: Count,
64-
/// Number of row groups for which the per-row
65-
/// [`RowFilter`](parquet::arrow::arrow_reader::RowFilter) was skipped
66-
/// because the static stats proved every row of the RG satisfies the
67-
/// predicate. The decoder is rebuilt at the boundary with an empty
68-
/// row filter so the upcoming RG decodes without per-row evaluation.
64+
/// Number of times the per-row
65+
/// [`RowFilter`](parquet::arrow::arrow_reader::RowFilter) was toggled
66+
/// off at a row-group boundary because static stats proved every row
67+
/// of the upcoming row group(s) satisfies the predicate. The decoder
68+
/// is rebuilt at that boundary with an empty row filter so the
69+
/// fully-matched run decodes without per-row evaluation.
70+
///
71+
/// Note this counts *suppression events*, not row groups: a run of
72+
/// consecutive fully-matched row groups shares a single toggle (the
73+
/// filter stays off across the run with no further rebuilds).
6974
pub row_filter_skipped_fully_matched: Count,
7075
/// Total number of bytes scanned
7176
pub bytes_scanned: Count,

‎datafusion/datasource-parquet/src/opener/mod.rs‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ use arrow::datatypes::DataType;
4040
use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselizer};
4141
use datafusion_physical_expr::projection::ProjectionExprs;
4242
use datafusion_physical_expr_adapter::replace_columns_with_literals;
43+
use datafusion_physical_expr_adapter::rewrite::rewrite_input_file_name_in_projection;
4344
use std::collections::{HashMap, VecDeque};
4445
use std::fmt;
4546
use std::future::Future;
@@ -795,6 +796,9 @@ impl ParquetMorselizer {
795796
.transpose()?;
796797
}
797798

799+
// Replace any `input_file_name()` UDFs in the projection with a literal for this file.
800+
projection = rewrite_input_file_name_in_projection(projection, &file_name)?;
801+
798802
let predicate_creation_errors = MetricBuilder::new(&self.metrics)
799803
.with_category(MetricCategory::Rows)
800804
.global_counter("num_predicate_creation_errors");
@@ -3336,8 +3340,12 @@ mod test {
33363340
/// (e.g. `row_number`) plumbed through `TableSchema`/`ParquetOpener`.
33373341
mod virtual_columns {
33383342
use super::*;
3339-
use arrow::array::{Array, Int64Array};
3343+
use arrow::array::{Array, Int64Array, StringArray};
33403344
use arrow::datatypes::FieldRef;
3345+
use datafusion_common::config::ConfigOptions;
3346+
use datafusion_expr::ScalarUDF;
3347+
use datafusion_functions::core::input_file_name::InputFileNameFunc;
3348+
use datafusion_physical_expr::{ScalarFunctionExpr, projection::ProjectionExpr};
33413349
use parquet::arrow::RowNumber;
33423350

33433351
/// Build a parquet `row_number` virtual column field. Spark's
@@ -3351,6 +3359,54 @@ mod test {
33513359
)
33523360
}
33533361

3362+
fn input_file_name_expr() -> Arc<dyn PhysicalExpr> {
3363+
Arc::new(ScalarFunctionExpr::new(
3364+
"input_file_name",
3365+
Arc::new(ScalarUDF::from(InputFileNameFunc::new())),
3366+
vec![],
3367+
Arc::new(Field::new("input_file_name", DataType::Utf8, true)),
3368+
Arc::new(ConfigOptions::default()),
3369+
))
3370+
}
3371+
3372+
#[tokio::test]
3373+
async fn test_input_file_name_projection() {
3374+
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
3375+
let path = "dir/input_file_name.parquet";
3376+
let (file_schema, data_size) = write_grouped_file(&store, path, 1, 3).await;
3377+
3378+
let projection = ProjectionExprs::new([
3379+
ProjectionExpr::new(Arc::new(Column::new("value", 0)), "value"),
3380+
ProjectionExpr::new(input_file_name_expr(), "file_name"),
3381+
]);
3382+
3383+
let morselizer = ParquetMorselizerBuilder::new()
3384+
.with_store(Arc::clone(&store))
3385+
.with_schema(file_schema)
3386+
.with_projection(projection)
3387+
.build();
3388+
3389+
let file =
3390+
PartitionedFile::new(path.to_string(), u64::try_from(data_size).unwrap());
3391+
let mut stream = open_file(&morselizer, file).await.unwrap();
3392+
let batch = stream.next().await.unwrap().unwrap();
3393+
assert!(stream.next().await.is_none());
3394+
3395+
assert_eq!(batch.num_columns(), 2);
3396+
assert_eq!(batch.schema().field(0).name(), "value");
3397+
assert_eq!(batch.schema().field(1).name(), "file_name");
3398+
3399+
let file_names = batch
3400+
.column(1)
3401+
.as_any()
3402+
.downcast_ref::<StringArray>()
3403+
.expect("file_name column should be Utf8");
3404+
assert_eq!(file_names.len(), 3);
3405+
for i in 0..file_names.len() {
3406+
assert_eq!(file_names.value(i), path);
3407+
}
3408+
}
3409+
33543410
/// Collect every `Int64` value from the given column in every batch
33553411
/// of a stream. Used to verify the `row_number` column end to end.
33563412
async fn collect_int64_values(

0 commit comments

Comments
 (0)