jayzhan211 commented on code in PR #26094:
URL: https://github.com/apache/datafusion/pull/26094#discussion_r4225588646
##########
datafusion/physical-optimizer/src/aggregate_statistics.rs:
##########
@@ -46,20 +47,26 @@ impl AggregateStatistics {
}
impl PhysicalOptimizerRule for AggregateStatistics {
- #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
- #[expect(clippy::allow_attributes)] // See
https://github.com/apache/datafusion/issues/18881#issuecomment-3621545670
- #[allow(clippy::only_used_in_recursion)] // See
https://github.com/rust-lang/rust-clippy/issues/14566
fn optimize(
&self,
plan: Arc<dyn ExecutionPlan>,
config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.optimize_with_context(plan, &ConfigOnlyContext::new(config))
+ }
+
+ #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
+ fn optimize_with_context(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ context: &dyn PhysicalOptimizerContext,
) -> Result<Arc<dyn ExecutionPlan>> {
if let Some(partial_agg_exec) = take_optimizable(&plan) {
let partial_agg_exec = partial_agg_exec
.downcast_ref::<AggregateExec>()
.expect("take_optimizable() ensures that this is a
AggregateExec");
- let stats = StatisticsContext::new()
- .compute(partial_agg_exec.input().as_ref(),
&StatisticsArgs::new())?;
+ let stats = context
Review Comment:
This lets registered providers decide query results, not just plan choices:
`AggregateStatistics` folds `COUNT(*)`/`MIN`/`MAX` from a provider's `Exact`
values, and `LimitPushdown` (`limit_pushdown.rs:398`) drops a limit over a
provider's `Exact(0)`. No test covers it. The only provider in any test is the
`Inexact` `HashJoinExec` one in `statistics_registry.slt`. With a provider
claiming `Exact(42)` on a 3-row table, `SELECT COUNT(*)` returns 42 on this
branch. If this line goes back to main's `StatisticsContext::new()`, it returns
3, and every existing test still passes.
Test for `physical_planner.rs` tests (passes here, fails with the line
reverted):
```rs
/// `AggregateStatistics` trusts an `Exact` row count from a registered
/// provider: it answers `COUNT(*)` without scanning
#[tokio::test]
async fn aggregate_statistics_consults_statistics_providers() -> Result<()> {
use datafusion_common::stats::Precision;
use datafusion_datasource::source::DataSourceExec;
use datafusion_physical_plan::operator_statistics::{
ClosureStatisticsProvider, StatisticsResult,
};
let provider = ClosureStatisticsProvider::with_matches(
|plan| plan.is::<DataSourceExec>(),
|plan, child_stats| {
let child_stats = child_stats
.iter()
.map(|c| Arc::clone(c.base_arc()))
.collect::<Vec<_>>();
let mut stats = Arc::unwrap_or_clone(
plan.statistics_from_inputs(&child_stats,
&StatisticsArgs::new())?,
);
stats.num_rows = Precision::Exact(42);
Ok(StatisticsResult::Computed(stats.into()))
},
);
let state = SessionStateBuilder::new()
.with_default_features()
.with_statistics_registry(StatisticsRegistry::with_providers(vec![Arc::new(
provider,
)]))
.build();
let ctx = SessionContext::new_with_state(state);
ctx.sql("CREATE TABLE t AS VALUES (1), (2), (3)")
.await?
.collect()
.await?;
let batches = ctx.sql("SELECT COUNT(*) FROM t").await?.collect().await?;
assert_batches_eq!(
&[
"+----------+",
"| count(*) |",
"+----------+",
"| 42 |",
"+----------+"
],
&batches
);
Ok(())
}
```
State the contract where provider authors will read it
(`operator_statistics/mod.rs:257`):
```diff
/// Implementations can handle specific operator types or override default
/// estimation logic. The chain of providers is traversed until one returns
/// [`StatisticsResult::Computed`].
+///
+/// `Exact` statistics are trusted for correctness, not just cost: optimizer
+/// rules use them to answer `COUNT(*)`, `MIN` and `MAX` without scanning
and
+/// to remove limits. Return `Inexact` for anything that is an estimate.
```
And in `56.0.0.md`:
```diff
`JoinSelection` already did. Sessions without registered providers (the
-default) are not affected.
+default) are not affected. Provider statistics can now change query results:
+`AggregateStatistics` answers `COUNT(*)`, `MIN` and `MAX` from `Exact`
values,
+and `LimitPushdown` removes a limit over an `Exact(0)` input, so a provider
+must return `Exact` only for guaranteed values.
```
--
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]