query_planning/expr.rs
1//! Expression IR for native plan nodes.
2//!
3//! Expressions use Arrow datatypes and native column references but do not
4//! depend on DataFusion physical-expression objects. Unsupported expression
5//! forms should be rejected during capability analysis or validation before
6//! reaching cudf execution.
7
8use std::{fmt, ops::Not};
9
10use arrow::datatypes::DataType;
11
12mod column;
13mod literal;
14
15pub use self::{
16 column::{ColumnQualifier, ColumnRef},
17 literal::LiteralValue,
18};
19
20/// Resolves a projection expression that only renames one existing column.
21///
22/// Nested aliases do not evaluate or allocate a new device column. Callers
23/// that rely on this for a zero-copy path must still verify that the declared
24/// output type matches the referenced input field.
25#[doc(hidden)]
26pub fn alias_wrapped_column(expr: &Expr) -> Option<&ColumnRef> {
27 match expr {
28 Expr::Column(column) => Some(column),
29 Expr::Alias { expr, .. } => alias_wrapped_column(expr),
30 _ => None,
31 }
32}
33
34pub(crate) fn representation_cast_preserves_identity(source_type: &DataType, target_type: &DataType) -> bool {
35 source_type == target_type
36 || matches!(
37 (source_type, target_type),
38 (DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View, DataType::Utf8View)
39 | (DataType::Utf8View, DataType::Utf8 | DataType::LargeUtf8)
40 | (DataType::Utf8, DataType::LargeUtf8)
41 | (DataType::LargeUtf8, DataType::Utf8)
42 )
43}
44
45/// The column and target type of a cast, itself possibly aliased, whose operand
46/// is a bare or aliased column. Paired with
47/// [`representation_cast_preserves_identity`] it finds casts that change only
48/// the string representation and keep row identity.
49#[doc(hidden)]
50pub fn representation_cast_column(expr: &Expr) -> Option<(&ColumnRef, &DataType)> {
51 match expr {
52 Expr::Alias { expr, .. } => representation_cast_column(expr),
53 Expr::Cast { expr, target_type } => alias_wrapped_column(expr).map(|column| (column, target_type)),
54 _ => None,
55 }
56}
57
58/// The `(months, days, nanoseconds)` of a month-day-nano interval literal, when
59/// `expr` is one.
60#[doc(hidden)]
61pub fn interval_month_day_nano(expr: &Expr) -> Option<(i32, i32, i64)> {
62 let Expr::Literal(LiteralValue::IntervalMonthDayNano { months, days, nanoseconds }) = expr else {
63 return None;
64 };
65 Some((*months, *days, *nanoseconds))
66}
67
68/// Scalar expression IR used by native query-plan nodes.
69///
70/// `Expr` describes logical expression structure independently of DataFusion's
71/// physical-expression APIs. It is intentionally compact: it names input
72/// columns, literals, SQL-style predicates, scalar operations, and aliases that
73/// a native executor or validator can either lower directly or reject with a
74/// structured diagnostic. The `Display` implementation is part of the
75/// human-readable diagnostics contract; keep it stable unless the displayed IR
76/// is misleading.
77#[derive(Debug, Clone, PartialEq)]
78pub enum Expr {
79 /// A column reference resolved against the current node input schema.
80 Column(ColumnRef),
81 /// A typed literal value.
82 Literal(LiteralValue),
83 /// SQL `NOT`; a null operand yields null.
84 Not(Box<Expr>),
85 /// SQL `IS NULL`; the result is never null.
86 IsNull(Box<Expr>),
87 /// SQL `IS NOT NULL`; the result is never null.
88 IsNotNull(Box<Expr>),
89 /// Explicit cast to an Arrow data type.
90 Cast {
91 /// Expression whose result is cast.
92 expr: Box<Expr>,
93 /// Desired output type after the cast.
94 target_type: DataType,
95 },
96 /// SQL `TRY_CAST` to an Arrow data type. Rows whose value cannot be
97 /// represented in `target_type` evaluate to null and do not raise an error.
98 ///
99 /// Native execution currently admits only `Utf8` → integer targets;
100 /// other type pairs are rejected during validation/lowering.
101 TryCast {
102 /// Expression whose result is try-cast.
103 expr: Box<Expr>,
104 /// Desired output type after the try-cast.
105 target_type: DataType,
106 },
107 /// Binary infix operation such as arithmetic, comparison, or boolean logic.
108 Binary {
109 /// Left operand.
110 left: Box<Expr>,
111 /// Operation to apply.
112 op: BinaryOp,
113 /// Right operand.
114 right: Box<Expr>,
115 /// Optional declared output type for operators whose result type is a
116 /// frontend contract. The native backend cannot infer that type from
117 /// operands alone.
118 output_type: Option<DataType>,
119 },
120 /// SQL `LIKE` or `NOT LIKE` string predicate.
121 Like {
122 /// String expression being matched.
123 expr: Box<Expr>,
124 /// SQL LIKE pattern. `%` and `_` keep their SQL wildcard meaning.
125 pattern: String,
126 /// Optional SQL escape marker. `None` means no explicit escape clause.
127 escape: Option<String>,
128 /// `true` for `NOT LIKE`; `false` for `LIKE`.
129 negated: bool,
130 },
131 /// SQL substring expression.
132 ///
133 /// Current native support uses a one-based `start` with `start >= 1`.
134 /// When `length` is present, it must be nonnegative and
135 /// `(start - 1) + length` must fit in `i32`.
136 Substring {
137 /// String expression being sliced.
138 expr: Box<Expr>,
139 /// One-based start position. Current native execution requires `start >= 1`.
140 start: i32,
141 /// Optional character count. When present, it must be nonnegative and
142 /// `(start - 1) + length` must fit in `i32`.
143 length: Option<i32>,
144 },
145 /// SQL `LOWER`; requires a string input and keeps its type.
146 Lower(Box<Expr>),
147 /// SQL `UPPER`; requires a string input and keeps its type.
148 Upper(Box<Expr>),
149 /// SQL `TRIM(trim_chars FROM expr)`.
150 ///
151 /// `trim_chars` is a set: every leading or trailing character that appears
152 /// in it is removed. An empty set is the identity. The DataFusion frontend
153 /// lowers the one-argument SQL form with a single space and omits the node
154 /// when the set is empty.
155 Trim {
156 /// String expression trimmed at both ends.
157 expr: Box<Expr>,
158 /// Characters removed from both ends.
159 trim_chars: String,
160 },
161 /// Leading-only form of [`Self::Trim`], with the same set semantics.
162 LTrim {
163 /// String expression trimmed at its start.
164 expr: Box<Expr>,
165 /// Characters removed from the start.
166 trim_chars: String,
167 },
168 /// Trailing-only form of [`Self::Trim`], with the same set semantics.
169 RTrim {
170 /// String expression trimmed at its end.
171 expr: Box<Expr>,
172 /// Characters removed from the end.
173 trim_chars: String,
174 },
175 /// Literal substring containment test.
176 StringContains {
177 /// String expression searched.
178 expr: Box<Expr>,
179 /// Literal substring; wildcards have no meaning here.
180 pattern: String,
181 },
182 /// Literal prefix test.
183 StartsWith {
184 /// String expression tested.
185 expr: Box<Expr>,
186 /// Literal prefix.
187 prefix: String,
188 },
189 /// Literal suffix test.
190 EndsWith {
191 /// String expression tested.
192 expr: Box<Expr>,
193 /// Literal suffix.
194 suffix: String,
195 },
196 /// SQL `CHARACTER_LENGTH(expr)` or `OCTET_LENGTH(expr)` string expression.
197 StringLength {
198 /// String expression whose length is measured.
199 expr: Box<Expr>,
200 /// Whether the length counts characters or UTF-8 bytes.
201 unit: StringLengthUnit,
202 /// Declared frontend output type.
203 output_type: DataType,
204 },
205 /// SQL `REGEXP_REPLACE(expr, pattern, replacement)` string expression.
206 RegexpReplace {
207 /// String expression being transformed.
208 expr: Box<Expr>,
209 /// Regex pattern.
210 pattern: String,
211 /// Replacement template, including optional back-references.
212 replacement: String,
213 /// Declared frontend output type.
214 output_type: DataType,
215 },
216 /// SQL `CONCAT(expr, ...)` string expression.
217 ///
218 /// Native execution follows DataFusion semantics: null inputs contribute
219 /// empty strings and do not produce a null output row.
220 Concat(Vec<Expr>),
221 /// SQL `ABS` over a numeric input.
222 Abs(Box<Expr>),
223 /// SQL unary arithmetic or math function.
224 ///
225 /// The `output_type` is the frontend-declared result type. It is part of
226 /// the native contract because cudf's expression builder can infer a
227 /// physical type that still needs to be aligned with DataFusion's exposed
228 /// schema. Results follow IEEE 754: domain errors produce NaN or an
229 /// infinity without raising an error, and null inputs produce null.
230 UnaryMath {
231 /// Function applied to each row.
232 op: UnaryMathOp,
233 /// Numeric expression the function is applied to.
234 expr: Box<Expr>,
235 /// Declared output type after evaluation.
236 output_type: DataType,
237 },
238 /// SQL `ROUND(expr, decimal_places)` numeric expression.
239 ///
240 /// The `output_type` is the frontend-declared DataFusion result type. It
241 /// is part of the native contract because decimal `ROUND` may reduce scale
242 /// even when the lower-level cuDF operation preserves the input storage
243 /// type.
244 Round {
245 /// Numeric or decimal expression being rounded.
246 expr: Box<Expr>,
247 /// Number of decimal places. Half-way ties round away from zero.
248 decimal_places: i32,
249 /// Declared output type after rounding.
250 output_type: DataType,
251 },
252 /// SQL `IN` or `NOT IN` membership predicate.
253 ///
254 /// The IR can represent general SQL `IN` lists. Runnable native plans
255 /// currently require a non-empty list of literal, non-null candidate
256 /// values.
257 InList {
258 /// Expression whose value is compared with the list.
259 expr: Box<Expr>,
260 /// Candidate expressions in SQL list order. Runnable native plans
261 /// currently require this list to be non-empty and contain only
262 /// literal, non-null values.
263 list: Vec<Expr>,
264 /// `true` for `NOT IN`; `false` for `IN`.
265 negated: bool,
266 },
267 /// SQL `IS DISTINCT FROM` null-aware comparison.
268 ///
269 /// Unlike ordinary inequality this always produces a non-null boolean:
270 /// two nulls are not distinct, and exactly one null is distinct.
271 IsDistinctFrom {
272 /// Left comparison operand.
273 left: Box<Expr>,
274 /// Right comparison operand.
275 right: Box<Expr>,
276 },
277 /// SQL `IS NOT DISTINCT FROM` null-aware comparison.
278 ///
279 /// Unlike ordinary equality this always produces a non-null boolean:
280 /// two nulls are not distinct and therefore compare true.
281 IsNotDistinctFrom {
282 /// Left comparison operand.
283 left: Box<Expr>,
284 /// Right comparison operand.
285 right: Box<Expr>,
286 },
287 /// Searched SQL `CASE WHEN ... THEN ... ELSE ... END`.
288 Case {
289 /// Ordered branches. The first branch whose predicate evaluates true is selected.
290 branches: Vec<CaseWhen>,
291 /// Result expression used when no branch predicate evaluates true.
292 else_expr: Box<Expr>,
293 },
294 /// SQL date/time field extraction.
295 Extract {
296 /// Field to extract from the input expression.
297 field: DatePart,
298 /// Date, timestamp, or compatible temporal expression.
299 expr: Box<Expr>,
300 },
301 /// SQL `DATE_TRUNC(field, expr)` expression.
302 DateTrunc {
303 /// Field to truncate the input expression to.
304 field: DatePart,
305 /// Timestamp expression being truncated.
306 expr: Box<Expr>,
307 /// Declared frontend output type.
308 output_type: DataType,
309 },
310 /// Output alias for projection and aggregate argument naming.
311 Alias {
312 /// Expression being renamed.
313 expr: Box<Expr>,
314 /// Field name exposed to downstream schemas and diagnostics.
315 alias: String,
316 },
317}
318
319impl Expr {
320 /// Direct child expressions in evaluation order.
321 pub fn children(&self) -> impl Iterator<Item = &Expr> {
322 let children: Vec<&Expr> = match self {
323 Self::Column(_) | Self::Literal(_) => Vec::new(),
324 Self::Not(expr)
325 | Self::IsNull(expr)
326 | Self::IsNotNull(expr)
327 | Self::Cast { expr, .. }
328 | Self::TryCast { expr, .. }
329 | Self::Like { expr, .. }
330 | Self::Substring { expr, .. }
331 | Self::Lower(expr)
332 | Self::Upper(expr)
333 | Self::Trim { expr, .. }
334 | Self::LTrim { expr, .. }
335 | Self::RTrim { expr, .. }
336 | Self::StringContains { expr, .. }
337 | Self::StartsWith { expr, .. }
338 | Self::EndsWith { expr, .. }
339 | Self::StringLength { expr, .. }
340 | Self::RegexpReplace { expr, .. }
341 | Self::Abs(expr)
342 | Self::UnaryMath { expr, .. }
343 | Self::Round { expr, .. }
344 | Self::Extract { expr, .. }
345 | Self::DateTrunc { expr, .. }
346 | Self::Alias { expr, .. } => vec![expr.as_ref()],
347 Self::Binary { left, right, .. }
348 | Self::IsDistinctFrom { left, right }
349 | Self::IsNotDistinctFrom { left, right } => vec![left.as_ref(), right.as_ref()],
350 Self::Concat(exprs) => exprs.iter().collect(),
351 Self::InList { expr, list, .. } => std::iter::once(expr.as_ref()).chain(list).collect(),
352 Self::Case { branches, else_expr } => branches
353 .iter()
354 .flat_map(|branch| [&branch.when, &branch.then])
355 .chain(std::iter::once(else_expr.as_ref()))
356 .collect(),
357 };
358 children.into_iter()
359 }
360
361 /// Whether `predicate` holds for this expression or any expression below
362 /// it, visited depth first.
363 pub fn any(&self, mut predicate: impl FnMut(&Expr) -> bool) -> bool {
364 let mut pending = vec![self];
365 while let Some(expr) = pending.pop() {
366 if predicate(expr) {
367 return true;
368 }
369 pending.extend(expr.children());
370 }
371 false
372 }
373
374 /// Every column this expression references, in depth-first order.
375 pub fn columns(&self) -> impl Iterator<Item = &ColumnRef> {
376 let mut pending = vec![self];
377 std::iter::from_fn(move || {
378 while let Some(expr) = pending.pop() {
379 if let Self::Column(column) = expr {
380 return Some(column);
381 }
382 let first = pending.len();
383 pending.extend(expr.children());
384 pending[first..].reverse();
385 }
386 None
387 })
388 }
389
390 /// Unqualified reference resolved by name against the node's single input.
391 /// Multi-input nodes use [`Self::input_column`], since the same leaf name
392 /// may appear on both sides.
393 pub fn column(name: impl Into<String>) -> Self {
394 Self::Column(ColumnRef::unqualified(name))
395 }
396
397 /// Build a column expression qualified by a positional input.
398 ///
399 /// Use this for multi-input nodes such as joins, where equal leaf column
400 /// names may appear on both sides. `name` is the source/input field name;
401 /// `display_name` is the projected or user-visible name.
402 pub fn input_column(input: usize, name: impl Into<String>, display_name: impl Into<String>) -> Self {
403 Self::Column(ColumnRef::input(input, name, display_name))
404 }
405
406 /// Typed literal; the Arrow type travels with the value.
407 pub fn literal(value: LiteralValue) -> Self {
408 Self::Literal(value)
409 }
410
411 /// Binary operation whose result type the backend infers from the
412 /// operands. Use [`Self::binary_with_output_type`] when the frontend
413 /// declares a type the backend cannot infer, such as a decimal division's
414 /// result scale.
415 pub fn binary(left: Self, op: BinaryOp, right: Self) -> Self {
416 Self::Binary { left: Box::new(left), op, right: Box::new(right), output_type: None }
417 }
418
419 /// Build a binary operation expression with a declared result type.
420 pub fn binary_with_output_type(left: Self, op: BinaryOp, right: Self, output_type: DataType) -> Self {
421 Self::Binary { left: Box::new(left), op, right: Box::new(right), output_type: Some(output_type) }
422 }
423
424 /// SQL `NOT`; the [`Not`] operator impl builds the same node.
425 pub fn logical_not(expr: Self) -> Self {
426 Self::Not(Box::new(expr))
427 }
428
429 /// SQL `IS NULL`; the result is never null.
430 pub fn is_null(expr: Self) -> Self {
431 Self::IsNull(Box::new(expr))
432 }
433
434 /// SQL `IS NOT NULL`; the result is never null.
435 pub fn is_not_null(expr: Self) -> Self {
436 Self::IsNotNull(Box::new(expr))
437 }
438
439 /// Explicit cast to `target_type`. See [`Self::try_cast`] for the form
440 /// that yields null for values the target cannot represent.
441 pub fn cast(expr: Self, target_type: DataType) -> Self {
442 Self::Cast { expr: Box::new(expr), target_type }
443 }
444
445 /// SQL `TRY_CAST`; native execution admits only `Utf8` to integer targets.
446 pub fn try_cast(expr: Self, target_type: DataType) -> Self {
447 Self::TryCast { expr: Box::new(expr), target_type }
448 }
449
450 /// SQL `LIKE` without an `ESCAPE` clause; `%` and `_` are wildcards.
451 pub fn like(expr: Self, pattern: impl Into<String>) -> Self {
452 Self::Like { expr: Box::new(expr), pattern: pattern.into(), escape: None, negated: false }
453 }
454
455 /// SQL `NOT LIKE` without an `ESCAPE` clause.
456 pub fn not_like(expr: Self, pattern: impl Into<String>) -> Self {
457 Self::Like { expr: Box::new(expr), pattern: pattern.into(), escape: None, negated: true }
458 }
459
460 /// SQL `LIKE ... ESCAPE escape`; validation checks the escape clause
461 /// before execution.
462 pub fn like_escape(expr: Self, pattern: impl Into<String>, escape: impl Into<String>) -> Self {
463 Self::Like { expr: Box::new(expr), pattern: pattern.into(), escape: Some(escape.into()), negated: false }
464 }
465
466 /// SQL `NOT LIKE ... ESCAPE escape`.
467 pub fn not_like_escape(expr: Self, pattern: impl Into<String>, escape: impl Into<String>) -> Self {
468 Self::Like { expr: Box::new(expr), pattern: pattern.into(), escape: Some(escape.into()), negated: true }
469 }
470
471 /// Build a SQL substring expression.
472 ///
473 /// Current native support requires `start >= 1`, nonnegative `length` when
474 /// present, and `(start - 1) + length` to fit in `i32`.
475 pub fn substring(expr: Self, start: i32, length: Option<i32>) -> Self {
476 Self::Substring { expr: Box::new(expr), start, length }
477 }
478
479 /// SQL `LOWER`; requires a string input and keeps its type.
480 pub fn lower(expr: Self) -> Self {
481 Self::Lower(Box::new(expr))
482 }
483
484 /// SQL `UPPER`; requires a string input and keeps its type.
485 pub fn upper(expr: Self) -> Self {
486 Self::Upper(Box::new(expr))
487 }
488
489 /// SQL `TRIM(trim_chars FROM expr)`; see [`Self::Trim`] for the set
490 /// semantics of `trim_chars`.
491 pub fn trim(expr: Self, trim_chars: impl Into<String>) -> Self {
492 Self::Trim { expr: Box::new(expr), trim_chars: trim_chars.into() }
493 }
494
495 /// Leading-only trim; see [`Self::Trim`].
496 pub fn ltrim(expr: Self, trim_chars: impl Into<String>) -> Self {
497 Self::LTrim { expr: Box::new(expr), trim_chars: trim_chars.into() }
498 }
499
500 /// Trailing-only trim; see [`Self::Trim`].
501 pub fn rtrim(expr: Self, trim_chars: impl Into<String>) -> Self {
502 Self::RTrim { expr: Box::new(expr), trim_chars: trim_chars.into() }
503 }
504
505 /// Boolean test that `expr` contains the literal `pattern`; no wildcard
506 /// interpretation.
507 pub fn string_contains(expr: Self, pattern: impl Into<String>) -> Self {
508 Self::StringContains { expr: Box::new(expr), pattern: pattern.into() }
509 }
510
511 /// Boolean test that `expr` starts with the literal `prefix`.
512 pub fn starts_with(expr: Self, prefix: impl Into<String>) -> Self {
513 Self::StartsWith { expr: Box::new(expr), prefix: prefix.into() }
514 }
515
516 /// Boolean test that `expr` ends with the literal `suffix`.
517 pub fn ends_with(expr: Self, suffix: impl Into<String>) -> Self {
518 Self::EndsWith { expr: Box::new(expr), suffix: suffix.into() }
519 }
520
521 /// SQL `CHARACTER_LENGTH` with the frontend-declared `output_type`.
522 pub fn string_length(expr: Self, output_type: DataType) -> Self {
523 Self::StringLength { expr: Box::new(expr), unit: StringLengthUnit::Characters, output_type }
524 }
525
526 /// SQL `OCTET_LENGTH` with the frontend-declared `output_type`.
527 pub fn octet_length(expr: Self, output_type: DataType) -> Self {
528 Self::StringLength { expr: Box::new(expr), unit: StringLengthUnit::Bytes, output_type }
529 }
530
531 /// SQL `REGEXP_REPLACE`; `replacement` may contain back-references.
532 pub fn regexp_replace(
533 expr: Self,
534 pattern: impl Into<String>,
535 replacement: impl Into<String>,
536 output_type: DataType,
537 ) -> Self {
538 Self::RegexpReplace {
539 expr: Box::new(expr),
540 pattern: pattern.into(),
541 replacement: replacement.into(),
542 output_type,
543 }
544 }
545
546 /// SQL `CONCAT`; null inputs contribute empty strings, as in DataFusion.
547 pub fn concat(exprs: Vec<Self>) -> Self {
548 Self::Concat(exprs)
549 }
550
551 /// SQL `ABS` over a numeric input.
552 pub fn abs(expr: Self) -> Self {
553 Self::Abs(Box::new(expr))
554 }
555
556 /// Unary math function with the frontend-declared `output_type`; see
557 /// [`UnaryMathOp`] for the IEEE 754 domain behavior of each function.
558 pub fn unary_math(op: UnaryMathOp, expr: Self, output_type: DataType) -> Self {
559 Self::UnaryMath { op, expr: Box::new(expr), output_type }
560 }
561
562 /// SQL `ROUND`; ties round away from zero, and `output_type` carries the
563 /// decimal scale the frontend declared.
564 pub fn round(expr: Self, decimal_places: i32, output_type: DataType) -> Self {
565 Self::Round { expr: Box::new(expr), decimal_places, output_type }
566 }
567
568 /// Build a SQL `IN` predicate.
569 ///
570 /// The IR accepts expression candidates, but runnable native plans require
571 /// a non-empty list of literal, non-null candidate values.
572 pub fn in_list(expr: Self, list: Vec<Self>) -> Self {
573 Self::InList { expr: Box::new(expr), list, negated: false }
574 }
575
576 /// Build a SQL `NOT IN` predicate.
577 ///
578 /// The IR accepts expression candidates, but runnable native plans require
579 /// a non-empty list of literal, non-null candidate values.
580 pub fn not_in_list(expr: Self, list: Vec<Self>) -> Self {
581 Self::InList { expr: Box::new(expr), list, negated: true }
582 }
583
584 /// Null-aware inequality; the result is never null.
585 pub fn is_distinct_from(left: Self, right: Self) -> Self {
586 Self::IsDistinctFrom { left: Box::new(left), right: Box::new(right) }
587 }
588
589 /// Null-aware equality; the result is never null.
590 pub fn is_not_distinct_from(left: Self, right: Self) -> Self {
591 Self::IsNotDistinctFrom { left: Box::new(left), right: Box::new(right) }
592 }
593
594 /// Searched `CASE`. The IR requires an `ELSE` expression; a frontend
595 /// supplies one for a SQL `CASE` that omits it.
596 pub fn case(branches: Vec<CaseWhen>, else_expr: Self) -> Self {
597 Self::Case { branches, else_expr: Box::new(else_expr) }
598 }
599
600 /// SQL `EXTRACT(field FROM expr)`; native execution supports the fields in
601 /// [`DatePart::EXTRACT_FIELDS`].
602 pub fn extract(field: DatePart, expr: Self) -> Self {
603 Self::Extract { field, expr: Box::new(expr) }
604 }
605
606 /// SQL `DATE_TRUNC`; native execution supports the fields in
607 /// [`DatePart::DATE_TRUNC_FIELDS`].
608 pub fn date_trunc(field: DatePart, expr: Self, output_type: DataType) -> Self {
609 Self::DateTrunc { field, expr: Box::new(expr), output_type }
610 }
611
612 /// Renames the output field without changing evaluation. See
613 /// [`alias_wrapped_column`] for the rename that stays zero-copy.
614 pub fn alias(self, alias: impl Into<String>) -> Self {
615 Self::Alias { expr: Box::new(self), alias: alias.into() }
616 }
617
618 /// Returns true when this expression tree contains a `REGEXP_REPLACE`.
619 ///
620 /// Used by capability labeling and runtime preflight paths that size
621 /// replace-with-backrefs via the owning cudf output-size contract.
622 pub fn contains_regexp_replace(&self) -> bool {
623 self.any(|expr| matches!(expr, Self::RegexpReplace { .. }))
624 }
625}
626
627impl Not for Expr {
628 type Output = Self;
629
630 fn not(self) -> Self::Output {
631 Self::Not(Box::new(self))
632 }
633}
634
635impl fmt::Display for Expr {
636 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
637 match self {
638 Self::Column(column) => write!(f, "{column}"),
639 Self::Literal(value) => write!(f, "{value}"),
640 Self::Not(expr) => write!(f, "NOT {expr}"),
641 Self::IsNull(expr) => write!(f, "{expr} IS NULL"),
642 Self::IsNotNull(expr) => write!(f, "{expr} IS NOT NULL"),
643 Self::Cast { expr, target_type } => write!(f, "CAST({expr} AS {target_type})"),
644 Self::TryCast { expr, target_type } => write!(f, "TRY_CAST({expr} AS {target_type})"),
645 Self::Binary { left, op, right, output_type } => {
646 write!(f, "({left} {op} {right})")?;
647 if let Some(output_type) = output_type {
648 write!(f, "::{output_type}")?;
649 }
650 Ok(())
651 },
652 Self::Like { expr, pattern, escape, negated } => {
653 let operator = if *negated { "NOT LIKE" } else { "LIKE" };
654 write!(f, "{expr} {operator} {pattern:?}")?;
655 if let Some(escape) = escape {
656 write!(f, " ESCAPE {escape:?}")?;
657 }
658 Ok(())
659 },
660 Self::Substring { expr, start, length } => match length {
661 Some(length) => write!(f, "SUBSTRING({expr} FROM {start} FOR {length})"),
662 None => write!(f, "SUBSTRING({expr} FROM {start})"),
663 },
664 Self::Lower(expr) => write!(f, "LOWER({expr})"),
665 Self::Upper(expr) => write!(f, "UPPER({expr})"),
666 Self::Trim { expr, trim_chars } => write!(f, "TRIM({trim_chars:?} FROM {expr})"),
667 Self::LTrim { expr, trim_chars } => write!(f, "LTRIM({expr}, {trim_chars:?})"),
668 Self::RTrim { expr, trim_chars } => write!(f, "RTRIM({expr}, {trim_chars:?})"),
669 Self::StringContains { expr, pattern } => write!(f, "CONTAINS({expr}, {pattern:?})"),
670 Self::StartsWith { expr, prefix } => write!(f, "STARTS_WITH({expr}, {prefix:?})"),
671 Self::EndsWith { expr, suffix } => write!(f, "ENDS_WITH({expr}, {suffix:?})"),
672 Self::StringLength { expr, unit: StringLengthUnit::Characters, output_type } => {
673 write!(f, "CHARACTER_LENGTH({expr})::{output_type}")
674 },
675 Self::StringLength { expr, unit: StringLengthUnit::Bytes, output_type } => {
676 write!(f, "OCTET_LENGTH({expr})::{output_type}")
677 },
678 Self::RegexpReplace { expr, pattern, replacement, output_type } => {
679 write!(f, "REGEXP_REPLACE({expr}, {pattern:?}, {replacement:?})::{output_type}")
680 },
681 Self::Concat(exprs) => {
682 f.write_str("CONCAT(")?;
683 for (index, expr) in exprs.iter().enumerate() {
684 if index > 0 {
685 f.write_str(", ")?;
686 }
687 write!(f, "{expr}")?;
688 }
689 f.write_str(")")
690 },
691 Self::Abs(expr) => write!(f, "ABS({expr})"),
692 Self::UnaryMath { op, expr, output_type } => write!(f, "{op}({expr})::{output_type}"),
693 Self::Round { expr, decimal_places, output_type } => {
694 write!(f, "ROUND({expr}, {decimal_places})::{output_type}")
695 },
696 Self::InList { expr, list, negated } => {
697 let operator = if *negated { "NOT IN" } else { "IN" };
698 write!(f, "{expr} {operator} (")?;
699 for (index, item) in list.iter().enumerate() {
700 if index > 0 {
701 f.write_str(", ")?;
702 }
703 write!(f, "{item}")?;
704 }
705 f.write_str(")")
706 },
707 Self::IsDistinctFrom { left, right } => write!(f, "({left} IS DISTINCT FROM {right})"),
708 Self::IsNotDistinctFrom { left, right } => write!(f, "({left} IS NOT DISTINCT FROM {right})"),
709 Self::Case { branches, else_expr } => {
710 f.write_str("CASE")?;
711 for branch in branches {
712 write!(f, " WHEN {} THEN {}", branch.when, branch.then)?;
713 }
714 write!(f, " ELSE {else_expr} END")
715 },
716 Self::Extract { field, expr } => write!(f, "EXTRACT({field} FROM {expr})"),
717 Self::DateTrunc { field, expr, output_type } => {
718 write!(f, "DATE_TRUNC({field}, {expr})::{output_type}")
719 },
720 Self::Alias { expr, alias } => write!(f, "{expr} AS {alias}"),
721 }
722 }
723}
724
725/// One `WHEN predicate THEN value` branch of a searched SQL `CASE` expression.
726#[derive(Debug, Clone, PartialEq)]
727pub struct CaseWhen {
728 /// Boolean predicate for this branch.
729 pub when: Expr,
730 /// Result expression chosen when `when` evaluates true.
731 pub then: Expr,
732}
733
734impl CaseWhen {
735 /// Build one searched-case branch.
736 pub fn new(when: Expr, then: Expr) -> Self {
737 Self { when, then }
738 }
739}
740
741/// Unit counted by [`Expr::StringLength`].
742#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
743pub enum StringLengthUnit {
744 /// Unicode scalar values, as SQL `CHARACTER_LENGTH` counts them.
745 Characters,
746 /// UTF-8 bytes, as SQL `OCTET_LENGTH` counts them.
747 Bytes,
748}
749
750/// Date or timestamp field named by `EXTRACT` or `DATE_TRUNC`.
751#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
752pub enum DatePart {
753 /// Calendar year field.
754 Year,
755 /// Calendar month field.
756 Month,
757 /// Day of month for extraction, or day boundary for truncation.
758 Day,
759 /// Hour field for extraction, or hour boundary for truncation.
760 Hour,
761 /// Minute field for extraction, or minute boundary for truncation.
762 Minute,
763 /// Second field for extraction, or second boundary for truncation.
764 Second,
765 /// Millisecond field for extraction, or millisecond boundary for truncation.
766 Millisecond,
767 /// Microsecond field for extraction, or microsecond boundary for truncation.
768 Microsecond,
769 /// Nanosecond boundary for timestamp truncation.
770 Nanosecond,
771}
772
773impl DatePart {
774 /// Fields native `EXTRACT` reads from a `Date32` or a timezone-free
775 /// timestamp.
776 pub const EXTRACT_DATE_FIELDS: [Self; 3] = [Self::Year, Self::Month, Self::Day];
777 /// Fields native `EXTRACT` reads from a timezone-free timestamp only. As in
778 /// DataFusion, `MILLISECOND` and `MICROSECOND` include the whole seconds
779 /// of the minute, and sub-unit remainders are dropped.
780 pub const EXTRACT_TIME_FIELDS: [Self; 5] =
781 [Self::Hour, Self::Minute, Self::Second, Self::Millisecond, Self::Microsecond];
782 /// Fields native `EXTRACT` returns: [`Self::EXTRACT_DATE_FIELDS`] and
783 /// [`Self::EXTRACT_TIME_FIELDS`]. `NANOSECOND` is not native.
784 pub const EXTRACT_FIELDS: [Self; 8] = [
785 Self::Year,
786 Self::Month,
787 Self::Day,
788 Self::Hour,
789 Self::Minute,
790 Self::Second,
791 Self::Millisecond,
792 Self::Microsecond,
793 ];
794 /// Fields `EXTRACT` returns inside a post-join residual filter.
795 pub(crate) const POST_JOIN_EXTRACT_FIELDS: [Self; 3] = [Self::Year, Self::Month, Self::Day];
796 /// Fields native `DATE_TRUNC` floors a timestamp to: the fixed-duration
797 /// ones. Calendar months and years are not.
798 pub const DATE_TRUNC_FIELDS: [Self; 7] =
799 [Self::Day, Self::Hour, Self::Minute, Self::Second, Self::Millisecond, Self::Microsecond, Self::Nanosecond];
800
801 /// The field a SQL field name names, ignoring case.
802 pub fn from_sql_name(name: &str) -> Option<Self> {
803 [
804 Self::Year,
805 Self::Month,
806 Self::Day,
807 Self::Hour,
808 Self::Minute,
809 Self::Second,
810 Self::Millisecond,
811 Self::Microsecond,
812 Self::Nanosecond,
813 ]
814 .into_iter()
815 .find(|field| field.to_string().eq_ignore_ascii_case(name))
816 }
817}
818
819impl fmt::Display for DatePart {
820 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
821 f.write_str(match self {
822 Self::Year => "YEAR",
823 Self::Month => "MONTH",
824 Self::Day => "DAY",
825 Self::Hour => "HOUR",
826 Self::Minute => "MINUTE",
827 Self::Second => "SECOND",
828 Self::Millisecond => "MILLISECOND",
829 Self::Microsecond => "MICROSECOND",
830 Self::Nanosecond => "NANOSECOND",
831 })
832 }
833}
834
835/// Unary arithmetic or math function in native expression IR.
836#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
837pub enum UnaryMathOp {
838 /// Arithmetic negation of signed integers or floating-point values.
839 /// Integer arrays wrap at the minimum value; floating-point zero changes sign.
840 Negate,
841 /// Square root; negative inputs produce NaN.
842 Sqrt,
843 /// Natural logarithm; zero produces negative infinity and negative inputs
844 /// produce NaN.
845 Ln,
846 /// Natural exponential `e ^ x`.
847 Exp,
848 /// Largest integral floating-point value not above the input. NaN and
849 /// infinities pass through; the output type equals the floating-point input type.
850 Floor,
851 /// Smallest integral floating-point value not below the input. NaN and
852 /// infinities pass through; the output type equals the floating-point input type.
853 Ceil,
854}
855
856impl fmt::Display for UnaryMathOp {
857 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
858 f.write_str(match self {
859 Self::Negate => "NEGATE",
860 Self::Sqrt => "SQRT",
861 Self::Ln => "LN",
862 Self::Exp => "EXP",
863 Self::Floor => "FLOOR",
864 Self::Ceil => "CEIL",
865 })
866 }
867}
868
869/// Binary operator in native expression IR.
870#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
871pub enum BinaryOp {
872 /// Arithmetic addition. Temporal handling is context/type-dependent.
873 Add,
874 /// Arithmetic subtraction. Temporal handling is context/type-dependent.
875 Subtract,
876 /// Arithmetic multiplication under the operands' admitted dtype contract.
877 Multiply,
878 /// Arithmetic division under the operands' admitted dtype contract.
879 Divide,
880 /// Arithmetic remainder.
881 Modulo,
882 /// Integer bitwise AND.
883 BitwiseAnd,
884 /// Integer bitwise OR.
885 BitwiseOr,
886 /// Integer bitwise exclusive OR.
887 BitwiseXor,
888 /// Integer shift left by the right operand.
889 ShiftLeft,
890 /// Integer shift right by the right operand.
891 ShiftRight,
892 /// SQL equality; a null operand produces null.
893 Eq,
894 /// SQL inequality; a null operand produces null.
895 NotEq,
896 /// Strict less-than comparison.
897 Lt,
898 /// Less-than-or-equal comparison.
899 LtEq,
900 /// Strict greater-than comparison.
901 Gt,
902 /// Greater-than-or-equal comparison.
903 GtEq,
904 /// SQL three-valued boolean conjunction.
905 And,
906 /// SQL three-valued boolean disjunction.
907 Or,
908}
909
910impl fmt::Display for BinaryOp {
911 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
912 f.write_str(match self {
913 Self::Add => "+",
914 Self::Subtract => "-",
915 Self::Multiply => "*",
916 Self::Divide => "/",
917 Self::Modulo => "%",
918 Self::BitwiseAnd => "&",
919 Self::BitwiseOr => "|",
920 Self::BitwiseXor => "^",
921 Self::ShiftLeft => "<<",
922 Self::ShiftRight => ">>",
923 Self::Eq => "=",
924 Self::NotEq => "!=",
925 Self::Lt => "<",
926 Self::LtEq => "<=",
927 Self::Gt => ">",
928 Self::GtEq => ">=",
929 Self::And => "AND",
930 Self::Or => "OR",
931 })
932 }
933}
934
935/// Aggregate operation descriptor for a native aggregate node.
936///
937/// `group_exprs` names the grouping keys. An empty key list represents a
938/// global aggregate. `aggr_exprs` lists logical aggregate outputs in output
939/// schema order, using each expression's `alias` as the externally visible
940/// field name.
941///
942/// Logical `AVG` is intentionally represented here as one aggregate expression.
943/// Execution code must expand it into explicit `SUM` plus non-null `COUNT`
944/// state whenever partial/final aggregation is involved, and finalize with
945/// `SUM / COUNT`. It must never merge AVG states by averaging partial averages.
946///
947/// Logical `StddevSamp` defaults to DataFusion sample standard-deviation
948/// semantics: null inputs are ignored, inputs are accumulated as Float64
949/// Welford state (`count`, `mean`, `m2`), counts 0 and 1 produce null, and
950/// finalization is `sqrt(m2 / (count - 1))`. Execution may explicitly select
951/// the GPU-native tolerant contract, which keeps SQL null/ddof semantics but
952/// does not preserve DataFusion's order-sensitive final bits.
953#[derive(Debug, Clone, PartialEq)]
954pub struct AggregateSpec {
955 /// Grouping key expressions in output key order.
956 pub group_exprs: Vec<Expr>,
957 /// Aggregate expressions in output aggregate order.
958 pub aggr_exprs: Vec<AggregateExpr>,
959 /// Optional grouping-set masks appended to the group expression contract.
960 pub grouping_sets: Option<GroupingSetsSpec>,
961}
962
963impl AggregateSpec {
964 /// Creates a standard aggregate without grouping sets.
965 pub fn new(group_exprs: Vec<Expr>, aggr_exprs: Vec<AggregateExpr>) -> Self {
966 Self { group_exprs, aggr_exprs, grouping_sets: None }
967 }
968
969 /// Attaches GROUPING SETS / ROLLUP / CUBE masks to this aggregate.
970 pub fn with_grouping_sets(mut self, grouping_sets: GroupingSetsSpec) -> Self {
971 self.grouping_sets = Some(grouping_sets);
972 self
973 }
974}
975
976/// GROUPING SETS / ROLLUP / CUBE grouping masks.
977///
978/// `groups` follows DataFusion's physical `PhysicalGroupBy` convention: each
979/// mask has one entry per `group_exprs` item, and `true` means that grouping
980/// key is replaced by a typed NULL for this grouping set. `grouping_id_alias`
981/// names the generated grouping id key appended after the group expressions.
982#[derive(Debug, Clone, PartialEq, Eq)]
983pub struct GroupingSetsSpec {
984 /// Grouping-set masks; `true` means the corresponding key is nulled.
985 pub groups: Vec<Vec<bool>>,
986 /// Output alias for the generated grouping-id column.
987 pub grouping_id_alias: String,
988}
989
990impl GroupingSetsSpec {
991 /// Creates grouping-set metadata from masks and a grouping-id alias.
992 pub fn new(groups: Vec<Vec<bool>>, grouping_id_alias: impl Into<String>) -> Self {
993 Self { groups, grouping_id_alias: grouping_id_alias.into() }
994 }
995}
996
997/// One logical aggregate expression and its output alias.
998///
999/// `expr` is the aggregate argument. For `COUNT`, normal column expressions
1000/// count non-null values. Planner-produced count-star forms use a non-null
1001/// literal or the native zero-column count source representation.
1002///
1003/// When `filter` is present, it is evaluated for every input row before this
1004/// aggregate is updated. Only rows where the predicate evaluates true
1005/// contribute to this aggregate; false and null predicate values are ignored.
1006#[derive(Debug, Clone, PartialEq)]
1007pub struct AggregateExpr {
1008 /// Aggregate function to apply.
1009 pub function: AggregateFunction,
1010 /// Aggregate argument expression.
1011 pub expr: Expr,
1012 /// Output field alias for this aggregate.
1013 pub alias: String,
1014 /// Whether SQL DISTINCT semantics are requested for the argument.
1015 pub distinct: bool,
1016 /// Optional per-aggregate SQL `FILTER (WHERE ...)` predicate.
1017 pub filter: Option<Expr>,
1018}
1019
1020impl AggregateExpr {
1021 /// Creates a non-DISTINCT aggregate expression.
1022 pub fn new(function: AggregateFunction, expr: Expr, alias: impl Into<String>) -> Self {
1023 Self { function, expr, alias: alias.into(), distinct: false, filter: None }
1024 }
1025
1026 /// Sets the SQL DISTINCT flag.
1027 pub fn with_distinct(mut self, distinct: bool) -> Self {
1028 self.distinct = distinct;
1029 self
1030 }
1031
1032 /// Sets an optional SQL aggregate filter predicate.
1033 pub fn with_filter(mut self, filter: Option<Expr>) -> Self {
1034 self.filter = filter;
1035 self
1036 }
1037}
1038
1039/// Logical aggregate functions supported by the native IR.
1040///
1041/// `Avg` carries SQL AVG semantics: null inputs are ignored, an all-null or
1042/// empty input produces null, and distributed execution must track SUM and
1043/// non-null COUNT as state before final division.
1044#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1045pub enum AggregateFunction {
1046 /// SQL SUM over non-null inputs.
1047 Sum,
1048 /// SQL MIN over non-null inputs.
1049 Min,
1050 /// SQL MAX over non-null inputs.
1051 Max,
1052 /// SQL COUNT over non-null inputs or count-star representation.
1053 Count,
1054 /// Exact SQL count-distinct over non-null inputs.
1055 ///
1056 /// Nulls do not contribute to the count. The logical output is a count
1057 /// integer and must be explicitly cast to the frontend-declared output
1058 /// dtype when libcudf returns a different integer width. Partial/final
1059 /// execution must not merge per-partition `NUnique` counts by summing
1060 /// them; bounded execution must aggregate over the complete raw input for
1061 /// the distinct scope.
1062 NUnique,
1063 /// SQL AVG represented logically; execution expands it to sum/count state.
1064 Avg,
1065 /// SQL median over non-null inputs.
1066 ///
1067 /// Median is not mergeable through the existing partial/final aggregate
1068 /// state model and must remain single-stage unless an explicit state
1069 /// representation is added.
1070 Median,
1071 /// SQL sample standard deviation.
1072 StddevSamp,
1073 /// SQL population standard deviation.
1074 StddevPop,
1075 /// SQL sample variance.
1076 VarSamp,
1077 /// SQL population variance.
1078 VarPop,
1079}
1080
1081impl AggregateFunction {
1082 /// Returns true for floating statistical aggregates whose SQL result is
1083 /// governed by the statistical aggregate execution contract.
1084 pub const fn is_floating_statistical(self) -> bool {
1085 matches!(self, Self::StddevSamp | Self::StddevPop | Self::VarSamp | Self::VarPop)
1086 }
1087
1088 /// Degrees-of-freedom correction for floating statistical aggregates.
1089 pub const fn statistical_ddof(self) -> Option<u32> {
1090 match self {
1091 Self::StddevSamp | Self::VarSamp => Some(1),
1092 Self::StddevPop | Self::VarPop => Some(0),
1093 _ => None,
1094 }
1095 }
1096
1097 /// Stable lower-case function label used in native diagnostics.
1098 pub const fn as_str(self) -> &'static str {
1099 match self {
1100 Self::Sum => "sum",
1101 Self::Min => "min",
1102 Self::Max => "max",
1103 Self::Count => "count",
1104 Self::NUnique => "nunique",
1105 Self::Avg => "avg",
1106 Self::Median => "median",
1107 Self::StddevSamp => "stddev_samp",
1108 Self::StddevPop => "stddev_pop",
1109 Self::VarSamp => "var_samp",
1110 Self::VarPop => "var_pop",
1111 }
1112 }
1113}
1114
1115/// Ordered list of sort keys for a native sort node.
1116///
1117/// Sort expressions are evaluated in list order: later keys break ties from
1118/// earlier keys. A runnable sort node is expected to contain at least one key;
1119/// an empty list is useful only as an invalid/transient state before validation.
1120#[derive(Debug, Clone, PartialEq)]
1121pub struct SortSpec {
1122 /// Sort keys in precedence order.
1123 pub expressions: Vec<SortExpr>,
1124}
1125
1126impl SortSpec {
1127 /// Build a sort specification from ordered key expressions.
1128 pub fn new(expressions: Vec<SortExpr>) -> Self {
1129 Self { expressions }
1130 }
1131}
1132
1133/// One sort key and its ordering options.
1134#[derive(Debug, Clone, PartialEq)]
1135pub struct SortExpr {
1136 /// Expression that produces the key value for each row.
1137 pub expr: Expr,
1138 /// `true` for descending order, `false` for ascending order.
1139 pub descending: bool,
1140 /// `true` places nulls before non-nulls, `false` places nulls after them.
1141 ///
1142 /// Null placement is explicit and independent of sort direction.
1143 pub nulls_first: bool,
1144}
1145
1146impl SortExpr {
1147 /// Build one sort key.
1148 pub fn new(expr: Expr, descending: bool, nulls_first: bool) -> Self {
1149 Self { expr, descending, nulls_first }
1150 }
1151}
1152
1153/// Row-count limit with SQL offset/fetch semantics.
1154///
1155/// `skip` rows are discarded first. When `fetch` is `Some`, at most that many
1156/// rows are returned after `skip`; when `fetch` is `None`, all remaining rows
1157/// are returned. If `skip` is greater than the input row count, the result is
1158/// empty. Finite top-k planning is valid only when `fetch` is `Some` and
1159/// `skip + fetch` fits in `usize`.
1160#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1161pub struct LimitSpec {
1162 /// Number of leading rows to discard.
1163 pub skip: usize,
1164 /// Maximum number of rows to return after `skip`; `None` means no fetch bound.
1165 pub fetch: Option<usize>,
1166}
1167
1168impl LimitSpec {
1169 /// Build a limit with explicit `skip` and `fetch` values.
1170 pub fn new(skip: usize, fetch: impl Into<Option<usize>>) -> Self {
1171 Self { skip, fetch: fetch.into() }
1172 }
1173
1174 /// Build a limit that returns at most `fetch` rows from the start.
1175 pub fn fetch(fetch: usize) -> Self {
1176 Self { skip: 0, fetch: Some(fetch) }
1177 }
1178
1179 /// Build an offset-only limit that returns all rows after `skip`.
1180 pub fn offset(skip: usize) -> Self {
1181 Self { skip, fetch: None }
1182 }
1183}