query/optimizer/
count_wildcard.rs1use datafusion::datasource::DefaultTableSource;
16use datafusion_common::tree_node::{
17 Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeVisitor,
18};
19use datafusion_common::{Column, Result as DataFusionResult, ScalarValue};
20use datafusion_expr::expr::{AggregateFunction, WindowFunction};
21use datafusion_expr::utils::COUNT_STAR_EXPANSION;
22use datafusion_expr::{Expr, LogicalPlan, WindowFunctionDefinition, col, lit};
23use datafusion_optimizer::AnalyzerRule;
24use datafusion_optimizer::utils::NamePreserver;
25use datafusion_sql::TableReference;
26use table::table::adapter::DfTableProviderAdapter;
27
28#[derive(Debug)]
34pub struct CountWildcardToTimeIndexRule;
35
36impl AnalyzerRule for CountWildcardToTimeIndexRule {
37 fn name(&self) -> &str {
38 "count_wildcard_to_time_index_rule"
39 }
40
41 fn analyze(
42 &self,
43 plan: LogicalPlan,
44 _config: &datafusion::config::ConfigOptions,
45 ) -> DataFusionResult<LogicalPlan> {
46 plan.transform_down_with_subqueries(&Self::analyze_internal)
47 .data()
48 }
49}
50
51impl CountWildcardToTimeIndexRule {
52 fn analyze_internal(plan: LogicalPlan) -> DataFusionResult<Transformed<LogicalPlan>> {
53 let name_preserver = NamePreserver::new(&plan);
54 let new_arg = if let Some(time_index) = Self::try_find_time_index_col(&plan) {
55 vec![col(time_index)]
56 } else {
57 vec![lit(COUNT_STAR_EXPANSION)]
58 };
59 plan.map_expressions(|expr| {
60 let original_name = name_preserver.save(&expr);
61 let transformed_expr = expr.transform_up(|expr| match expr {
62 Expr::WindowFunction(mut window_function)
63 if Self::is_count_star_window_aggregate(&window_function) =>
64 {
65 window_function.params.args.clone_from(&new_arg);
66 Ok(Transformed::yes(Expr::WindowFunction(window_function)))
67 }
68 Expr::AggregateFunction(mut aggregate_function)
69 if Self::is_count_star_aggregate(&aggregate_function) =>
70 {
71 aggregate_function.params.args.clone_from(&new_arg);
72 Ok(Transformed::yes(Expr::AggregateFunction(
73 aggregate_function,
74 )))
75 }
76 _ => Ok(Transformed::no(expr)),
77 })?;
78 Ok(transformed_expr.update_data(|data| original_name.restore(data)))
79 })
80 }
81
82 fn try_find_time_index_col(plan: &LogicalPlan) -> Option<Column> {
83 let mut finder = TimeIndexFinder::default();
84 plan.visit(&mut finder).unwrap();
86 let col = finder.into_column();
87
88 if let Some(col) = &col {
90 let mut is_valid = false;
91 if plan.inputs().len() > 1 {
93 return None;
94 }
95 for input in plan.inputs() {
96 if input.schema().has_column(col) {
97 is_valid = true;
98 break;
99 }
100 }
101 if !is_valid {
102 return None;
103 }
104 }
105
106 col
107 }
108}
109
110impl CountWildcardToTimeIndexRule {
112 #[expect(deprecated)]
113 fn args_at_most_wildcard_or_literal_one(args: &[Expr]) -> bool {
114 match args {
115 [] => true,
116 [Expr::Literal(ScalarValue::Int64(Some(v)), _)] => *v == 1,
117 [Expr::Wildcard { .. }] => true,
118 _ => false,
119 }
120 }
121
122 fn is_count_star_aggregate(aggregate_function: &AggregateFunction) -> bool {
123 let args = &aggregate_function.params.args;
124 matches!(aggregate_function,
125 AggregateFunction {
126 func,
127 ..
128 } if func.name() == "count" && Self::args_at_most_wildcard_or_literal_one(args))
129 }
130
131 fn is_count_star_window_aggregate(window_function: &WindowFunction) -> bool {
132 let args = &window_function.params.args;
133 matches!(window_function.fun,
134 WindowFunctionDefinition::AggregateUDF(ref udaf)
135 if udaf.name() == "count" && Self::args_at_most_wildcard_or_literal_one(args))
136 }
137}
138
139#[derive(Default)]
140struct TimeIndexFinder {
141 time_index_col: Option<String>,
142 table_alias: Option<TableReference>,
143}
144
145impl TreeNodeVisitor<'_> for TimeIndexFinder {
146 type Node = LogicalPlan;
147
148 fn f_down(&mut self, node: &Self::Node) -> DataFusionResult<TreeNodeRecursion> {
149 if let LogicalPlan::SubqueryAlias(subquery_alias) = node {
150 self.table_alias = Some(subquery_alias.alias.clone());
151 }
152
153 if let LogicalPlan::TableScan(table_scan) = &node
154 && let Some(source) = table_scan
155 .source
156 .as_any()
157 .downcast_ref::<DefaultTableSource>()
158 && let Some(adapter) = source
159 .table_provider
160 .as_any()
161 .downcast_ref::<DfTableProviderAdapter>()
162 {
163 let table_info = adapter.table().table_info();
164 self.table_alias
165 .get_or_insert(TableReference::bare(table_info.name.clone()));
166 self.time_index_col = table_info
167 .meta
168 .schema
169 .timestamp_column()
170 .map(|c| c.name.clone());
171
172 return Ok(TreeNodeRecursion::Stop);
173 }
174
175 if node.inputs().len() > 1 {
176 return Ok(TreeNodeRecursion::Stop);
178 }
179
180 Ok(TreeNodeRecursion::Continue)
181 }
182
183 fn f_up(&mut self, _node: &Self::Node) -> DataFusionResult<TreeNodeRecursion> {
184 Ok(TreeNodeRecursion::Stop)
185 }
186}
187
188impl TimeIndexFinder {
189 fn into_column(self) -> Option<Column> {
190 self.time_index_col
191 .map(|c| Column::new(self.table_alias, c))
192 }
193}
194
195#[cfg(test)]
196mod test {
197 use std::sync::Arc;
198
199 use datafusion::functions_aggregate::count::count_all;
200 use datafusion_expr::LogicalPlanBuilder;
201 use table::table::numbers::NumbersTable;
202
203 use super::*;
204
205 #[test]
206 fn uppercase_table_name() {
207 let numbers_table = NumbersTable::table_with_name(0, "AbCdE".to_string());
208 let table_source = Arc::new(DefaultTableSource::new(Arc::new(
209 DfTableProviderAdapter::new(numbers_table),
210 )));
211
212 let plan = LogicalPlanBuilder::scan_with_filters("t", table_source, None, vec![])
213 .unwrap()
214 .aggregate(Vec::<Expr>::new(), vec![count_all()])
215 .unwrap()
216 .alias(r#""FgHiJ""#)
217 .unwrap()
218 .build()
219 .unwrap();
220
221 let mut finder = TimeIndexFinder::default();
222 plan.visit(&mut finder).unwrap();
223
224 assert_eq!(finder.table_alias, Some(TableReference::bare("FgHiJ")));
225 assert!(finder.time_index_col.is_none());
226 }
227}