Skip to main content

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}