1use std::borrow::Cow;
4
5use radixdb_core::{CompactArc, Schema, Value};
6use radixdb_functions::FunctionRegistry;
7use radixdb_sql::ast::*;
8use radixdb_storage::expression::Expression as StorageExpression;
9use rustc_hash::FxHashMap;
10
11use crate::context::ExecutionContext;
12use crate::optimizer::ExpressionSimplifier;
13use crate::pushdown;
14use crate::utils::substitute_outer_references;
15
16pub type DirectPushdown = (Option<Box<dyn StorageExpression>>, bool);
17
18pub fn prepare_bound_predicate(
21 predicate: Option<&Expression>,
22 schema: &Schema,
23 context: &ExecutionContext,
24) -> DirectPushdown {
25 predicate
26 .map(|predicate| pushdown::try_pushdown(predicate, schema, Some(context)))
27 .unwrap_or((None, false))
28}
29
30pub struct PreparedPredicate<'a> {
34 effective: Option<Cow<'a, Expression>>,
35 pub storage: Option<Box<dyn StorageExpression>>,
36 pub residual: Option<Expression>,
37}
38
39impl<'a> PreparedPredicate<'a> {
40 pub fn effective(&self) -> Option<&Expression> {
41 self.effective.as_deref()
42 }
43
44 pub fn memory_filter(&self) -> Option<&Expression> {
45 self.residual.as_ref().or_else(|| self.effective())
46 }
47
48 pub fn needs_memory_filter(&self) -> bool {
49 self.residual.is_some()
50 }
51}
52
53#[allow(clippy::too_many_arguments)]
56pub fn prepare_scan_predicate<'a>(
57 select_columns: &[Expression],
58 where_clause: Option<&'a Expression>,
59 all_columns: &[String],
60 schema: &Schema,
61 table_alias: Option<&str>,
62 context: &ExecutionContext,
63 functions: &FunctionRegistry,
64 where_has_subqueries: bool,
65) -> PreparedPredicate<'a> {
66 let aliases = build_alias_map_excluding(select_columns, Some(all_columns));
67 let aliased = (!aliases.is_empty())
68 .then(|| where_clause.map(|expression| substitute_aliases(expression, &aliases)))
69 .flatten();
70 let source = aliased.as_ref().or(where_clause);
71
72 let simplified = source.and_then(|expression| {
73 ExpressionSimplifier::with_registry(functions).try_simplify(expression)
74 });
75 let effective = if let Some(expression) = simplified {
76 Some(Cow::Owned(expression))
77 } else if let Some(expression) = aliased {
78 Some(Cow::Owned(expression))
79 } else {
80 where_clause.map(Cow::Borrowed)
81 };
82
83 let Some(predicate) = effective.as_deref() else {
84 return PreparedPredicate {
85 effective,
86 storage: None,
87 residual: None,
88 };
89 };
90
91 if where_has_subqueries {
92 let residual = predicate.clone();
93 return PreparedPredicate {
94 effective,
95 storage: None,
96 residual: Some(residual),
97 };
98 }
99
100 if let Some(outer_row) = context.outer_row() {
101 let scoped_outer_row: FxHashMap<CompactArc<str>, Value> = outer_row
102 .iter()
103 .filter(|(name, _)| {
104 let name = name.as_ref();
105 if let Some(dot) = name.rfind('.') {
106 !table_alias
107 .is_some_and(|qualifier| name[..dot].eq_ignore_ascii_case(qualifier))
108 } else {
109 schema.get_column_index(name).is_none()
110 }
111 })
112 .map(|(name, value)| (name.clone(), value.clone()))
113 .collect();
114 let substituted = substitute_outer_references(predicate, &scoped_outer_row);
115 let plan = pushdown::try_pushdown_plan(&substituted, schema, Some(context));
116 return if plan.storage_expr.is_some() {
117 PreparedPredicate {
118 effective,
119 storage: plan.storage_expr,
120 residual: plan.residual,
121 }
122 } else {
123 PreparedPredicate {
124 effective,
125 storage: None,
126 residual: Some(substituted),
127 }
128 };
129 }
130
131 let plan = pushdown::try_pushdown_plan(predicate, schema, Some(context));
132 PreparedPredicate {
133 effective,
134 storage: plan.storage_expr,
135 residual: plan.residual,
136 }
137}
138
139pub fn build_alias_map_excluding<'a>(
142 columns: &'a [Expression],
143 base_columns: Option<&[String]>,
144) -> FxHashMap<String, &'a Expression> {
145 let alias_count = columns
146 .iter()
147 .filter(|expression| matches!(expression, Expression::Aliased(_)))
148 .count();
149 if alias_count == 0 {
150 return FxHashMap::default();
151 }
152
153 let mut aliases = FxHashMap::with_capacity_and_hasher(alias_count, Default::default());
154 for expression in columns {
155 let Expression::Aliased(aliased) = expression else {
156 continue;
157 };
158 let name = aliased.alias.value_lower.to_string();
159 if base_columns.is_some_and(|base| {
160 base.iter().any(|column| {
161 column.eq_ignore_ascii_case(&name)
162 || column
163 .rsplit_once('.')
164 .is_some_and(|(_, tail)| tail.eq_ignore_ascii_case(&name))
165 })
166 }) {
167 continue;
168 }
169 aliases.insert(name, aliased.expression.as_ref());
170 }
171 aliases
172}
173
174pub fn substitute_aliases(
176 expression: &Expression,
177 aliases: &FxHashMap<String, &Expression>,
178) -> Expression {
179 match expression {
180 Expression::Identifier(identifier) => aliases
181 .get(identifier.value_lower.as_str())
182 .map_or_else(|| expression.clone(), |source| (*source).clone()),
183 Expression::Infix(infix) => Expression::Infix(InfixExpression {
184 token: infix.token.clone(),
185 left: Box::new(substitute_aliases(&infix.left, aliases)),
186 operator: infix.operator.clone(),
187 op_type: infix.op_type,
188 right: Box::new(substitute_aliases(&infix.right, aliases)),
189 }),
190 Expression::Prefix(prefix) => Expression::Prefix(PrefixExpression {
191 token: prefix.token.clone(),
192 operator: prefix.operator.clone(),
193 op_type: prefix.op_type,
194 right: Box::new(substitute_aliases(&prefix.right, aliases)),
195 }),
196 Expression::Between(between) => Expression::Between(BetweenExpression {
197 token: between.token.clone(),
198 expr: Box::new(substitute_aliases(&between.expr, aliases)),
199 lower: Box::new(substitute_aliases(&between.lower, aliases)),
200 upper: Box::new(substitute_aliases(&between.upper, aliases)),
201 not: between.not,
202 }),
203 Expression::In(input) => Expression::In(InExpression {
204 token: input.token.clone(),
205 left: Box::new(substitute_aliases(&input.left, aliases)),
206 right: Box::new(substitute_aliases(&input.right, aliases)),
207 not: input.not,
208 }),
209 Expression::FunctionCall(function) => Expression::FunctionCall(Box::new(FunctionCall {
210 token: function.token.clone(),
211 function: function.function.clone(),
212 arguments: function
213 .arguments
214 .iter()
215 .map(|argument| substitute_aliases(argument, aliases))
216 .collect(),
217 is_distinct: function.is_distinct,
218 order_by: function.order_by.clone(),
219 filter: function.filter.clone(),
220 })),
221 Expression::Case(case) => Expression::Case(Box::new(CaseExpression {
222 token: case.token.clone(),
223 value: case
224 .value
225 .as_ref()
226 .map(|value| Box::new(substitute_aliases(value, aliases))),
227 when_clauses: case
228 .when_clauses
229 .iter()
230 .map(|when| WhenClause {
231 token: when.token.clone(),
232 condition: substitute_aliases(&when.condition, aliases),
233 then_result: substitute_aliases(&when.then_result, aliases),
234 })
235 .collect(),
236 else_value: case
237 .else_value
238 .as_ref()
239 .map(|value| Box::new(substitute_aliases(value, aliases))),
240 })),
241 Expression::List(list) => Expression::List(Box::new(ListExpression {
242 token: list.token.clone(),
243 elements: list
244 .elements
245 .iter()
246 .map(|item| substitute_aliases(item, aliases))
247 .collect(),
248 })),
249 Expression::Like(like) => Expression::Like(LikeExpression {
250 token: like.token.clone(),
251 left: Box::new(substitute_aliases(&like.left, aliases)),
252 pattern: Box::new(substitute_aliases(&like.pattern, aliases)),
253 operator: like.operator.clone(),
254 escape: like
255 .escape
256 .as_ref()
257 .map(|value| Box::new(substitute_aliases(value, aliases))),
258 }),
259 _ => expression.clone(),
260 }
261}
262
263#[cfg(test)]
264mod tests {
265 use super::*;
266 use radixdb_core::{DataType, SchemaColumn};
267 use radixdb_sql::parse_sql;
268
269 #[test]
270 fn source_column_shadows_same_named_output_alias() {
271 let statements = parse_sql("SELECT a AS b FROM t WHERE b = 1").unwrap();
272 let [Statement::Select(statement)] = statements.as_slice() else {
273 panic!("expected SELECT");
274 };
275 let aliases =
276 build_alias_map_excluding(&statement.columns, Some(&["a".into(), "b".into()]));
277 assert!(aliases.is_empty());
278 }
279
280 #[test]
281 fn predicate_without_where_has_no_physical_parts() {
282 let schema = Schema::new(
283 "t",
284 vec![SchemaColumn::new(0, "id", DataType::Integer, false, true)],
285 );
286 let context = ExecutionContext::new();
287 let functions = FunctionRegistry::new();
288 let prepared = prepare_scan_predicate(
289 &[],
290 None,
291 &["id".into()],
292 &schema,
293 Some("t"),
294 &context,
295 &functions,
296 false,
297 );
298 assert!(prepared.effective().is_none());
299 assert!(prepared.storage.is_none());
300 assert!(!prepared.needs_memory_filter());
301 }
302}