Repository navigation
Conversation
github-actions
Bot
added
enhancement
pingzh
changed the title
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
Two new tests fail with the locked DataFusion dependencies.
|
|
||
| fn metric(plan: &dyn ExecutionPlan, name: &str) -> usize { | ||
| plan.metrics() | ||
| .and_then(|metrics| metrics.sum_by_name(name)) |
There was a problem hiding this comment.
[P2] Read the built-in output-row metric through its dedicated accessor. DataFusion 55.2's MetricsSet::sum_by_name deliberately excludes MetricValue::OutputRows, so this helper returns zero for output_rows even when the scan emitted 600 rows. The newly added live_bounds_prune_remapped_row_groups_and_keep_scan_metrics_on_the_scan test therefore fails at its 600-row assertion. I reproduced this with the exact locked DataFusion/Arrow/Parquet dependencies. Use the built-in output-row aggregation when this metric is requested so the new native suite can pass.
| .schema(scan.schema()) | ||
| .alias("sum") | ||
| .build() | ||
| .unwrap(); |
There was a problem hiding this comment.
[P2] Build the unsupported SUM fixture with a type accepted by the aggregate. This constructs SUM directly over the Int32 scan column, but DataFusion 55.2's physical AggregateExprBuilder does not apply SQL type coercion and Sum::return_type rejects Int32. Consequently build().unwrap() fails before the test reaches the wrapper's rejection check. I reproduced the failure against the locked dependencies; using a supported input type or explicitly applying the normal widening cast will let this test exercise the intended boundary.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
One additional component-level correctness issue.
| let Some(scan) = aggregate.input().downcast_ref::<DataSourceExec>() else { | ||
| return Ok(None); | ||
| }; | ||
| let Some((_, source)) = parquet_file_source(scan) else { |
There was a problem hiding this comment.
[P2] Reject computed expressions inside the scan projection before attaching this predicate. Checking that the aggregate argument is a Column does not prove it is a direct file column: DataSourceExec can project a computed expression and expose its result as a column to the aggregate. I reproduced this with a volatile integer UDF embedded in the Parquet scan projection and decoder filter pushdown enabled. Ordinary partial MIN returns 0 after two UDF calls; this wrapper is accepted, remaps the dynamic predicate through the projection, evaluates the UDF again during decoder filtering, and returns 10 after three calls. With decoder filtering disabled, both return 0. Please require the selected scan projection to map directly to the file column, or otherwise prove that re-evaluating it cannot change the result. This is an execution-level component reproduction using the locked DataFusion/Parquet versions; this module is not yet wired into Spark, so it does not assert an existing Spark query regression.
sunchao
left a comment
sunchao
left a comment
There was a problem hiding this comment.
The current fork CI also identifies a compilation lint failure.
| &session, | ||
| 2, | ||
| ); | ||
| let replaced = plan.with_new_children(vec![replacement]).unwrap(); |
There was a problem hiding this comment.
[P2] Use the non-deprecated child-replacement API in this test. The exact-head fork Rust job runs Clippy with warnings denied and fails here because ExecutionPlan::with_new_children is deprecated, before the new tests execute. The failure is visible in the fork CI job. Updating this call to the supported replacement helper will unblock that check.
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
I checked whether DataFusion's own MIN/MAX dynamic filter makes the wrapper unnecessary. AggregateExec::try_new already creates it, but it only reaches a scan through the FilterPushdown rule, which Comet does not run, and the wrapper reuses try_attach_parquet_reader_filter, so the wrapper makes sense to me. The comments are about the with_dynamic_filter_expr call that DataFusion main deprecates, code and test helpers copied from the TopK wrapper, and the module-wide dead_code allowance. I have not built or run any of this.
| let predicate = Arc::new(DynamicFilterPhysicalExpr::new( | ||
| self.template.aggr_expr()[0].expressions(), | ||
| lit(true), | ||
| )); | ||
| let reader = try_attach_parquet_reader_filter( | ||
| self.template.input(), | ||
| Arc::clone(&predicate), | ||
| &self.config, | ||
| )?; | ||
| let reader_filter_attached = reader.is_some(); | ||
| let input = reader.unwrap_or_else(|| Arc::clone(self.template.input())); | ||
| let aggregate = | ||
| Self::fresh_aggregate(&self.template, input)?.with_dynamic_filter_expr(predicate)?; |
There was a problem hiding this comment.
fresh_aggregate already goes through AggregateExec::try_new, which creates a fresh DynamicFilterPhysicalExpr and fresh accumulator bounds. Instead of building a second predicate and installing it with with_dynamic_filter_expr, could this build the fresh aggregate over the template input, take its filter from dynamic_expressions_produced() (the same accessor the tests use), attach the reader with it, and then call replace_children with the reader? In DataFusion 55.2 replace_children in Recompute mode copies the aggregate's dynamic_filter over, so the reader and the accumulators would share the same filter without the column-compatibility check in with_dynamic_filter_expr. On DataFusion main that method is #[deprecated(since = "56.0.0")] with no replacement, so this would otherwise be a clippy failure on the next DataFusion bump. I have not tried this, so please confirm the ordering works.
| fn execute_runtime_aggregate( | ||
| &self, | ||
| aggregate: AggregateExec, | ||
| partition: usize, | ||
| context: Arc<TaskContext>, | ||
| ) -> Result<SendableRecordBatchStream> { | ||
| let result = aggregate.execute(partition, context); | ||
| // The ungrouped stream registers its counters in execute(). Retain only | ||
| // their handles, including on error, so the runtime state can be dropped. | ||
| for metric in aggregate.metrics().unwrap_or_default().iter() { | ||
| self.metrics.register(Arc::clone(metric)); | ||
| } | ||
| drop(aggregate); | ||
| let input = result?; | ||
| // Also release state at EOF/error if the caller retains the spent stream. | ||
| let stream = futures::stream::unfold(Some(input), |input| async move { | ||
| let mut input = input?; | ||
| let batch = input.next().await?; | ||
| let remaining = if batch.is_ok() { Some(input) } else { None }; | ||
| Some((batch, remaining)) | ||
| }); | ||
| Ok(Box::pin(RecordBatchStreamAdapter::new( | ||
| self.schema(), | ||
| stream, | ||
| ))) | ||
| } |
There was a problem hiding this comment.
This function, build_runtime_aggregate, execute and most of the ExecutionPlan delegation have the same shape as TopKReaderFilterExec in topk.rs. execute_runtime_sort has the same metric registration, drop, and release-on-EOF unfold stream, line for line. Could the release stream and the metric hand-off move into a shared helper in dynamic_filter/mod.rs, so a lifecycle fix only has to be made in one place?
| fn parquet_input( | ||
| values: Vec<Option<i32>>, | ||
| key_type: &DataType, | ||
| session: &Arc<SessionContext>, | ||
| group_size: usize, | ||
| ) -> (tempfile::NamedTempFile, Arc<DataSourceExec>) { | ||
| let key = cast(&Int32Array::from(values), key_type).unwrap(); | ||
| let file_schema = Arc::new(Schema::new(vec![ | ||
| Field::new("payload", DataType::Int32, false), | ||
| Field::new("key", key_type.clone(), true), | ||
| ])); | ||
| let batch = RecordBatch::try_new( | ||
| Arc::clone(&file_schema), | ||
| vec![ | ||
| Arc::new(Int32Array::from_iter_values(0..key.len() as i32)) as ArrayRef, | ||
| key, | ||
| ], | ||
| ) | ||
| .unwrap(); | ||
| let file = tempfile::NamedTempFile::new().unwrap(); | ||
| let properties = WriterProperties::builder() | ||
| .set_max_row_group_row_count(Some(group_size)) | ||
| .set_statistics_enabled(EnabledStatistics::Chunk) | ||
| .set_dictionary_enabled(false) | ||
| .build(); | ||
| let mut writer = ArrowWriter::try_new( | ||
| file.reopen().unwrap(), | ||
| Arc::clone(&file_schema), | ||
| Some(properties), | ||
| ) | ||
| .unwrap(); | ||
| writer.write(&batch).unwrap(); | ||
| let metadata = writer.close().unwrap(); | ||
| assert_eq!( | ||
| metadata.num_row_groups(), | ||
| batch.num_rows().div_ceil(group_size) | ||
| ); | ||
| let scan = parquet_scan(&[&file], file_schema, vec![1], None, false, session); | ||
| (file, scan) | ||
| } | ||
|
|
||
| fn parquet_scan( | ||
| files: &[&tempfile::NamedTempFile], | ||
| file_schema: Arc<Schema>, | ||
| projection: Vec<usize>, | ||
| data_filters: Option<Vec<Arc<dyn PhysicalExpr>>>, | ||
| allow_type_promotion: bool, | ||
| session: &Arc<SessionContext>, | ||
| ) -> Arc<DataSourceExec> { | ||
| init_datasource_exec( | ||
| Arc::new(file_schema.project(&projection).unwrap()), | ||
| Some(file_schema), | ||
| None, | ||
| ObjectStoreUrl::local_filesystem(), | ||
| ObjectStoreBackend::Local, | ||
| vec![files | ||
| .iter() | ||
| .map(|file| { | ||
| PartitionedFile::from_path(file.path().to_string_lossy().into_owned()).unwrap() | ||
| }) | ||
| .collect()], | ||
| Some(projection), | ||
| data_filters, | ||
| None, | ||
| "UTC", | ||
| true, | ||
| false, | ||
| allow_type_promotion, | ||
| false, | ||
| session, | ||
| false, | ||
| false, | ||
| false, | ||
| false, | ||
| "CORRECTED", | ||
| "CORRECTED", | ||
| ) | ||
| .unwrap() | ||
| } |
There was a problem hiding this comment.
parquet_input and parquet_scan are exact copies of the helpers in topk/tests.rs, and session matches the TopK one apart from the batch-size argument. Here parquet_scan is always called with data_filters = None and allow_type_promotion = false. Could these live once in a #[cfg(test)] module under dynamic_filter shared by both wrappers?
| &DataType::Int32, | ||
| &session, | ||
| 2, | ||
| ); | ||
| let original = std::fs::read(file.path()).unwrap(); | ||
| let reader = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); | ||
| let mut builder = reader.metadata().as_ref().clone().into_builder(); | ||
| let groups = builder | ||
| .take_row_groups() | ||
| .into_iter() | ||
| .map(|group| { | ||
| let mut columns = group.columns().to_vec(); | ||
| let Some(ParquetStatistics::Int32(stats)) = columns[1].statistics() else { | ||
| panic!("INT32 statistics expected") | ||
| }; | ||
| let replacement = ParquetStatistics::int32( | ||
| stats.min_opt().copied(), | ||
| stats.max_opt().copied(), | ||
| None, | ||
| None, | ||
| false, | ||
| ); | ||
| columns[1] = columns[1] | ||
| .clone() | ||
| .into_builder() | ||
| .set_statistics(replacement) | ||
| .build() | ||
| .unwrap(); | ||
| group | ||
| .into_builder() | ||
| .set_column_metadata(columns) | ||
| .build() | ||
| .unwrap() | ||
| }) | ||
| .collect(); | ||
| let metadata = builder.set_row_groups(groups).build(); | ||
| let mut output = TrackedWrite::new(file.reopen().unwrap()); | ||
| output.write_all(&original).unwrap(); | ||
| ParquetMetaDataWriter::new_with_tracked(output, &metadata) | ||
| .finish() | ||
| .unwrap(); |
There was a problem hiding this comment.
This rewrites the footer statistics to drop null counts, which omit_last_null_count in topk/tests/statistics.rs already does for the last row group. Could that helper take the row groups to rewrite and be shared, instead of a second copy of the ParquetMetaDataWriter and TrackedWrite sequence?
| async fn signed_integer_extrema_preserve_empty_null_and_duplicate_results() { | ||
| let session = session(); | ||
| for key_type in [ | ||
| DataType::Int8, | ||
| DataType::Int16, | ||
| DataType::Int32, | ||
| DataType::Int64, | ||
| ] { | ||
| for values in [ | ||
| vec![], | ||
| vec![None; 4], | ||
| vec![None, Some(-8), Some(4), Some(-8), None], | ||
| ] { | ||
| for is_max in [false, true] { | ||
| let expected = if is_max { | ||
| values.iter().flatten().max() | ||
| } else { | ||
| values.iter().flatten().min() | ||
| } | ||
| .copied(); | ||
| let (_file, scan) = parquet_input(values.clone(), &key_type, &session, 2); | ||
| let plain = aggregate(scan, is_max); | ||
| let filtered = Arc::new(wrapper(&plain)); | ||
| let actual = collect(Arc::clone(&filtered) as _, session.task_ctx()) | ||
| .await | ||
| .unwrap(); | ||
| let baseline = collect(Arc::new(plain), session.task_ctx()).await.unwrap(); | ||
| assert_eq!(result(&actual), expected, "{key_type:?}, max={is_max}"); | ||
| assert_eq!(result(&actual), result(&baseline)); | ||
| assert_eq!( | ||
| metric(filtered.as_ref(), "dynamic_filter_minmax_filters_attached"), | ||
| 1 | ||
| ); | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
This matrix writes 24 Parquet files, and the baseline comparison repeats the expected assertion on the line above it. The result values do not depend on the wrapper, so the type, empty and all-NULL matrix looks like a better fit for a SQL file test once PR 2 wires the planner, since CometSqlFileTestSuite compares with Spark. Could this keep one type plus the attached-metric check, and drop the baseline comparison?
| // Spark planner integration is a separate step of apache/datafusion-comet#5774. | ||
| #[cfg_attr(not(test), allow(dead_code))] | ||
| mod min_max; |
There was a problem hiding this comment.
#[cfg_attr(not(test), allow(dead_code))] covers the whole module, so it will also hide any item that is still unused after PR 2. TopK landed with its planner hook in #6263. Could this PR include the planner step, or narrow the allowance to the try_new entry point PR 2 will call? About 590 of the 888 lines are tests for code nothing outside the module can reach yet, which makes it hard to judge the eligibility rules against what the planner will pass in.
Which issue does this PR close?
Part 1 of #5774; this PR does not close the issue. The three-PR plan is in this issue comment.
Rationale for this change
A partial
MIN(x)orMAX(x)can tighten a bound while consuming Parquet input and skip later row groups that cannot improve the result. DataFusion already updates that bound, but Comet needs to connect its producer to the native reader and scope the mutable state to one execution.What changes are included in this PR?
Add the internal
MinMaxReaderFilterExecwith the existing Parquet reader attachment. It accepts one unfiltered, non-distinct DataFusion MIN/MAX over a direct signed-integer column, with no grouping, in Partial mode, and one native input partition. Standalone projections, filters, limits, non-Parquet inputs, and pushed scan predicates remain ineligible.Each execution constructs a fresh aggregate, accumulator-bound state, and predicate. Aggregate metric handles and attachment counters belong to the wrapper; reader pruning and byte metrics stay with the scan. Unsupported reader attachment retains ordinary aggregation. Runtime state is released on completion, error, or cancellation.
This is the native building block only: no Spark configuration, serialization, or planner integration is included, and existing Spark execution is unchanged. Those changes, including handling identity casts in MIN/MAX arguments, belong to PR 2. The private module has a scoped dead-code allowance until that integration lands.
How are these changes tested?
Seven native tests cover:
rustfmt --checkandgit diff --checkpass. Two independent source reviews completed. The tests have not been executed locally; execution is pending CI.