spark_connect/
column.rs

1//! Column API mirroring PySpark's `pyspark.sql.connect.column.Column`.
2//!
3//! Provides the Column type and builder functions for constructing expressions.
4
5use std::ops::{Add, BitAnd, BitOr, Div, Mul, Not, Rem, Sub};
6
7use spark_connect_proto as proto;
8
9use crate::expression::{
10    Alias, CaseWhen, Cast, CastEvalMode, ColumnReference, Expression, ExtractValue, FrameBoundary,
11    LiteralExpression, SortOrder, UnresolvedFunction, UpdateFieldsExpr, WindowExpressionWrapper,
12};
13use crate::types::DataType;
14use crate::window::WindowSpec;
15
16/// `pyspark.sql.connect.column.Column`
17///
18/// Represents a column in a DataFrame, built from an Expression. All operations
19/// return new Columns by composing expressions.
20#[derive(Debug, Clone, PartialEq)]
21pub struct Column {
22    expr: Expression,
23}
24
25impl Column {
26    /// Create a Column from an Expression.
27    pub fn new(expr: Expression) -> Self {
28        Column { expr }
29    }
30
31    /// Get the underlying Expression.
32    pub fn expression(&self) -> &Expression {
33        &self.expr
34    }
35
36    /// Mirrors `pyspark.sql.column.Column.alias`.
37    pub fn alias(self, name: &str) -> Column {
38        Column {
39            expr: Expression::Alias(Box::new(Alias::new(self.expr, name))),
40        }
41    }
42
43    /// `Column.alias(name, metadata=...)` - attaches column metadata, serialized to
44    /// a JSON map (matching the reference client's `json.dumps(metadata)`).
45    pub fn alias_with_metadata(
46        self,
47        name: &str,
48        metadata: std::collections::BTreeMap<String, String>,
49    ) -> Column {
50        let json = serde_json::to_string(&metadata).unwrap_or_else(|_| "{}".to_string());
51        Column {
52            expr: Expression::Alias(Box::new(Alias::new(self.expr, name).with_metadata(json))),
53        }
54    }
55
56    /// Mirrors `pyspark.sql.column.Column.name` (alias for `alias`).
57    pub fn name(self, name: &str) -> Column {
58        self.alias(name)
59    }
60
61    /// Mirrors `pyspark.sql.column.Column.cast`.
62    pub fn cast(self, to_type: DataType) -> Column {
63        Column {
64            expr: Expression::Cast(Box::new(Cast::new(self.expr, to_type))),
65        }
66    }
67
68    /// Mirrors `pyspark.sql.column.Column.astype` (alias for `cast`).
69    pub fn astype(self, to_type: DataType) -> Column {
70        self.cast(to_type)
71    }
72
73    /// Mirrors `pyspark.sql.column.Column.cast` with a DDL type string, e.g.
74    /// `col("x").cast("string")` → `cast { type_str: "string" }`.
75    pub fn cast_str(self, type_name: &str) -> Column {
76        Column {
77            expr: Expression::Cast(Box::new(Cast::new_str(self.expr, type_name))),
78        }
79    }
80
81    /// Mirrors `pyspark.sql.column.Column.try_cast`.
82    pub fn try_cast(self, to_type: DataType) -> Column {
83        Column {
84            expr: Expression::Cast(Box::new(
85                Cast::new(self.expr, to_type).with_eval_mode(CastEvalMode::Try),
86            )),
87        }
88    }
89
90    /// `Column.try_cast` with a DDL type string.
91    pub fn try_cast_str(self, type_name: &str) -> Column {
92        Column {
93            expr: Expression::Cast(Box::new(
94                Cast::new_str(self.expr, type_name).with_eval_mode(CastEvalMode::Try),
95            )),
96        }
97    }
98
99    /// Mirrors `pyspark.sql.column.Column.isNull`.
100    pub fn is_null(self) -> Column {
101        Column {
102            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
103                "isNull",
104                vec![self.expr],
105            )),
106        }
107    }
108
109    /// Mirrors `pyspark.sql.column.Column.isNotNull`.
110    pub fn is_not_null(self) -> Column {
111        Column {
112            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
113                "isNotNull",
114                vec![self.expr],
115            )),
116        }
117    }
118
119    /// Mirrors `pyspark.sql.column.Column.substr`.
120    pub fn substr(self, start: Column, length: Column) -> Column {
121        Column {
122            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
123                "substr",
124                vec![self.expr, start.expr, length.expr],
125            )),
126        }
127    }
128
129    /// Mirrors `pyspark.sql.column.Column.like`.
130    pub fn like(self, pattern: &str) -> Column {
131        Column {
132            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
133                "like",
134                vec![
135                    self.expr,
136                    Expression::Literal(LiteralExpression::string(pattern)),
137                ],
138            )),
139        }
140    }
141
142    /// Mirrors `pyspark.sql.column.Column.rlike`.
143    pub fn rlike(self, pattern: &str) -> Column {
144        Column {
145            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
146                "rlike",
147                vec![
148                    self.expr,
149                    Expression::Literal(LiteralExpression::string(pattern)),
150                ],
151            )),
152        }
153    }
154
155    /// Mirrors `pyspark.sql.column.Column.contains`.
156    pub fn contains(self, other: Column) -> Column {
157        Column {
158            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
159                "contains",
160                vec![self.expr, other.expr],
161            )),
162        }
163    }
164
165    /// Mirrors `pyspark.sql.column.Column.ilike` (case-insensitive `like`).
166    pub fn ilike(self, pattern: &str) -> Column {
167        Column {
168            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
169                "ilike",
170                vec![
171                    self.expr,
172                    Expression::Literal(LiteralExpression::string(pattern)),
173                ],
174            )),
175        }
176    }
177
178    /// Mirrors `pyspark.sql.column.Column.isNaN`.
179    pub fn is_nan(self) -> Column {
180        Column {
181            expr: Expression::UnresolvedFunction(UnresolvedFunction::new("isNaN", vec![self.expr])),
182        }
183    }
184
185    /// Mirrors `pyspark.sql.column.Column.eqNullSafe` (null-safe equality `<=>`).
186    pub fn eq_null_safe(self, other: Column) -> Column {
187        Column {
188            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
189                "<=>",
190                vec![self.expr, other.expr],
191            )),
192        }
193    }
194
195    /// Mirrors `pyspark.sql.column.Column.bitwiseAND`.
196    pub fn bitwise_and(self, other: Column) -> Column {
197        Column {
198            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
199                "&",
200                vec![self.expr, other.expr],
201            )),
202        }
203    }
204
205    /// Mirrors `pyspark.sql.column.Column.bitwiseOR`.
206    pub fn bitwise_or(self, other: Column) -> Column {
207        Column {
208            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
209                "|",
210                vec![self.expr, other.expr],
211            )),
212        }
213    }
214
215    /// Mirrors `pyspark.sql.column.Column.bitwiseXOR`.
216    pub fn bitwise_xor(self, other: Column) -> Column {
217        Column {
218            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
219                "^",
220                vec![self.expr, other.expr],
221            )),
222        }
223    }
224
225    /// Mirrors `pyspark.sql.column.Column.between`:
226    /// `(self >= lower) & (self <= upper)`.
227    pub fn between(self, lower: Column, upper: Column) -> Column {
228        let lo = Expression::UnresolvedFunction(UnresolvedFunction::new(
229            ">=",
230            vec![self.expr.clone(), lower.expr],
231        ));
232        let hi = Expression::UnresolvedFunction(UnresolvedFunction::new(
233            "<=",
234            vec![self.expr, upper.expr],
235        ));
236        Column {
237            expr: Expression::UnresolvedFunction(UnresolvedFunction::new("and", vec![lo, hi])),
238        }
239    }
240
241    /// Mirrors `pyspark.sql.column.Column.isin`: membership test against the
242    /// given values (pass literals via `lit(..)` or other columns).
243    pub fn isin<C: Into<Column>>(self, values: impl IntoIterator<Item = C>) -> Column {
244        let values: Vec<Column> = values.into_iter().map(Into::into).collect();
245        let mut args = Vec::with_capacity(values.len() + 1);
246        args.push(self.expr);
247        args.extend(values.into_iter().map(|c| c.expr));
248        Column {
249            expr: Expression::UnresolvedFunction(UnresolvedFunction::new("in", args)),
250        }
251    }
252
253    /// Mirrors `pyspark.sql.column.Column.withField`: add or replace a field in
254    /// a struct-typed column.
255    pub fn with_field(self, field_name: &str, value: Column) -> Column {
256        Column {
257            expr: Expression::UpdateFields(Box::new(UpdateFieldsExpr::new(
258                self.expr,
259                field_name,
260                Some(value.expr),
261            ))),
262        }
263    }
264
265    /// Mirrors `pyspark.sql.column.Column.dropFields`: drop one or more fields
266    /// from a struct-typed column (chained `UpdateFields`, one per field).
267    pub fn drop_fields(self, field_names: Vec<&str>) -> Column {
268        let mut expr = self.expr;
269        for name in field_names {
270            expr = Expression::UpdateFields(Box::new(UpdateFieldsExpr::new(expr, name, None)));
271        }
272        Column { expr }
273    }
274
275    /// Mirrors `pyspark.sql.column.Column.startswith`.
276    pub fn startswith(self, other: Column) -> Column {
277        Column {
278            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
279                "startsWith",
280                vec![self.expr, other.expr],
281            )),
282        }
283    }
284
285    /// Mirrors `pyspark.sql.column.Column.endswith`.
286    pub fn endswith(self, other: Column) -> Column {
287        Column {
288            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
289                "endsWith",
290                vec![self.expr, other.expr],
291            )),
292        }
293    }
294
295    /// Mirrors `pyspark.sql.column.Column.asc`.
296    pub fn asc(self) -> Column {
297        self.asc_nulls_first()
298    }
299
300    /// Mirrors `pyspark.sql.column.Column.asc_nulls_first`.
301    pub fn asc_nulls_first(self) -> Column {
302        Column {
303            expr: Expression::SortOrder(Box::new(SortOrder::asc_nulls_first(self.expr))),
304        }
305    }
306
307    /// Mirrors `pyspark.sql.column.Column.asc_nulls_last`.
308    pub fn asc_nulls_last(self) -> Column {
309        Column {
310            expr: Expression::SortOrder(Box::new(SortOrder::asc_nulls_last(self.expr))),
311        }
312    }
313
314    /// Mirrors `pyspark.sql.column.Column.desc`.
315    pub fn desc(self) -> Column {
316        self.desc_nulls_last()
317    }
318
319    /// Mirrors `pyspark.sql.column.Column.desc_nulls_first`.
320    pub fn desc_nulls_first(self) -> Column {
321        Column {
322            expr: Expression::SortOrder(Box::new(SortOrder::desc_nulls_first(self.expr))),
323        }
324    }
325
326    /// Mirrors `pyspark.sql.column.Column.desc_nulls_last`.
327    pub fn desc_nulls_last(self) -> Column {
328        Column {
329            expr: Expression::SortOrder(Box::new(SortOrder::desc_nulls_last(self.expr))),
330        }
331    }
332
333    /// Mirrors `pyspark.sql.column.Column.when`.
334    pub fn when(self, condition: Column, value: Column) -> Column {
335        // If we already have a CaseWhen, add a branch
336        if let Expression::CaseWhen(case_when) = self.expr {
337            let mut branches = case_when.branches.clone();
338            branches.push((condition.expr, value.expr));
339            Column {
340                expr: Expression::CaseWhen(Box::new(CaseWhen {
341                    branches,
342                    else_expr: case_when.else_expr.clone(),
343                })),
344            }
345        } else {
346            // Start a new CaseWhen
347            Column {
348                expr: Expression::CaseWhen(Box::new(CaseWhen {
349                    branches: vec![(condition.expr, value.expr)],
350                    else_expr: None,
351                })),
352            }
353        }
354    }
355
356    /// Mirrors `pyspark.sql.column.Column.otherwise`.
357    pub fn otherwise(self, value: Column) -> Column {
358        if let Expression::CaseWhen(case_when) = self.expr {
359            Column {
360                expr: Expression::CaseWhen(Box::new(CaseWhen {
361                    branches: case_when.branches.clone(),
362                    else_expr: Some(Box::new(value.expr)),
363                })),
364            }
365        } else {
366            Column { expr: self.expr }
367        }
368    }
369
370    /// Mirrors `pyspark.sql.column.Column.getField` - struct field access.
371    /// Builds an `UnresolvedExtractValue` with the field name as a string literal.
372    pub fn get_field(self, name: &str) -> Column {
373        let extraction = Expression::Literal(LiteralExpression::string(name));
374        Column {
375            expr: Expression::UnresolvedExtractValue(Box::new(ExtractValue::new(
376                self.expr, extraction,
377            ))),
378        }
379    }
380
381    /// Mirrors `pyspark.sql.column.Column.__getitem__` / `getItem` - map/array/struct
382    /// extraction. Builds an `UnresolvedExtractValue`.
383    pub fn get_item(self, key: Column) -> Column {
384        Column {
385            expr: Expression::UnresolvedExtractValue(Box::new(ExtractValue::new(
386                self.expr, key.expr,
387            ))),
388        }
389    }
390
391    /// Convert to proto expression.
392    pub fn to_proto(&self) -> proto::Expression {
393        self.expr.to_proto()
394    }
395
396    /// Mirrors `pyspark.sql.column.Column.over` - window function application.
397    pub fn over(self, window_spec: WindowSpec) -> Column {
398        // Convert window frame spec if present
399        let frame_spec = window_spec.frame_spec.map(|(frame_type, lower, upper)| {
400            let frame_type_val = match frame_type {
401                crate::window::FrameType::Row => 1u32,
402                crate::window::FrameType::Range => 2u32,
403            };
404            let lower_boundary = match lower {
405                crate::window::FrameBound::UnboundedPreceding => FrameBoundary::UnboundedPreceding,
406                crate::window::FrameBound::Preceding(n) => FrameBoundary::Preceding(n),
407                crate::window::FrameBound::CurrentRow => FrameBoundary::CurrentRow,
408                crate::window::FrameBound::Following(n) => FrameBoundary::Following(n),
409                crate::window::FrameBound::UnboundedFollowing => FrameBoundary::UnboundedFollowing,
410            };
411            let upper_boundary = match upper {
412                crate::window::FrameBound::UnboundedPreceding => FrameBoundary::UnboundedPreceding,
413                crate::window::FrameBound::Preceding(n) => FrameBoundary::Preceding(n),
414                crate::window::FrameBound::CurrentRow => FrameBoundary::CurrentRow,
415                crate::window::FrameBound::Following(n) => FrameBoundary::Following(n),
416                crate::window::FrameBound::UnboundedFollowing => FrameBoundary::UnboundedFollowing,
417            };
418            (frame_type_val, lower_boundary, upper_boundary)
419        });
420
421        let window_expr = WindowExpressionWrapper::new(
422            self.expr,
423            window_spec.partition_spec,
424            window_spec.order_spec,
425            frame_spec,
426        );
427
428        Column {
429            expr: Expression::WindowExpression(Box::new(window_expr)),
430        }
431    }
432
433    // Comparison operators
434
435    /// Mirrors `pyspark.sql.column.Column.__eq__`.
436    pub fn eq(self, other: Column) -> Column {
437        Column {
438            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
439                "==",
440                vec![self.expr, other.expr],
441            )),
442        }
443    }
444
445    /// Mirrors `pyspark.sql.column.Column.__ne__`.
446    pub fn ne(self, other: Column) -> Column {
447        Column {
448            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
449                "not",
450                vec![Expression::UnresolvedFunction(UnresolvedFunction::new(
451                    "==",
452                    vec![self.expr, other.expr],
453                ))],
454            )),
455        }
456    }
457
458    /// Mirrors `pyspark.sql.column.Column.__gt__`.
459    pub fn gt(self, other: Column) -> Column {
460        Column {
461            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
462                ">",
463                vec![self.expr, other.expr],
464            )),
465        }
466    }
467
468    /// Mirrors `pyspark.sql.column.Column.__lt__`.
469    pub fn lt(self, other: Column) -> Column {
470        Column {
471            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
472                "<",
473                vec![self.expr, other.expr],
474            )),
475        }
476    }
477
478    /// Mirrors `pyspark.sql.column.Column.__ge__`.
479    pub fn ge(self, other: Column) -> Column {
480        Column {
481            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
482                ">=",
483                vec![self.expr, other.expr],
484            )),
485        }
486    }
487
488    /// Mirrors `pyspark.sql.column.Column.__le__`.
489    pub fn le(self, other: Column) -> Column {
490        Column {
491            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
492                "<=",
493                vec![self.expr, other.expr],
494            )),
495        }
496    }
497
498    // Arithmetic operators
499
500    /// Mirrors `pyspark.sql.column.Column.__add__`.
501    pub fn add(self, other: Column) -> Column {
502        Column {
503            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
504                "+",
505                vec![self.expr, other.expr],
506            )),
507        }
508    }
509
510    /// Mirrors `pyspark.sql.column.Column.__sub__`.
511    pub fn sub(self, other: Column) -> Column {
512        Column {
513            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
514                "-",
515                vec![self.expr, other.expr],
516            )),
517        }
518    }
519
520    /// Mirrors `pyspark.sql.column.Column.__mul__`.
521    pub fn mul(self, other: Column) -> Column {
522        Column {
523            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
524                "*",
525                vec![self.expr, other.expr],
526            )),
527        }
528    }
529
530    /// Mirrors `pyspark.sql.column.Column.__truediv__`.
531    pub fn div(self, other: Column) -> Column {
532        Column {
533            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
534                "/",
535                vec![self.expr, other.expr],
536            )),
537        }
538    }
539
540    /// Mirrors `pyspark.sql.column.Column.__mod__`.
541    pub fn modulo(self, other: Column) -> Column {
542        Column {
543            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
544                "%",
545                vec![self.expr, other.expr],
546            )),
547        }
548    }
549
550    // Logical operators
551
552    /// Mirrors `pyspark.sql.column.Column.__and__`.
553    pub fn and(self, other: Column) -> Column {
554        Column {
555            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
556                "and",
557                vec![self.expr, other.expr],
558            )),
559        }
560    }
561
562    /// Mirrors `pyspark.sql.column.Column.__or__`.
563    pub fn or(self, other: Column) -> Column {
564        Column {
565            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
566                "or",
567                vec![self.expr, other.expr],
568            )),
569        }
570    }
571
572    /// Mirrors `pyspark.sql.column.Column.__invert__`.
573    pub fn not(self) -> Column {
574        Column {
575            expr: Expression::UnresolvedFunction(UnresolvedFunction::new("not", vec![self.expr])),
576        }
577    }
578
579    /// Mirrors `pyspark.sql.column.Column.__neg__`.
580    pub fn neg(self) -> Column {
581        Column {
582            expr: Expression::UnresolvedFunction(UnresolvedFunction::new(
583                "negative",
584                vec![self.expr],
585            )),
586        }
587    }
588}
589
590// Operator trait implementations
591impl Add for Column {
592    type Output = Column;
593    fn add(self, other: Column) -> Column {
594        self.add(other)
595    }
596}
597
598impl Sub for Column {
599    type Output = Column;
600    fn sub(self, other: Column) -> Column {
601        self.sub(other)
602    }
603}
604
605impl Mul for Column {
606    type Output = Column;
607    fn mul(self, other: Column) -> Column {
608        self.mul(other)
609    }
610}
611
612impl Div for Column {
613    type Output = Column;
614    fn div(self, other: Column) -> Column {
615        self.div(other)
616    }
617}
618
619impl Rem for Column {
620    type Output = Column;
621    fn rem(self, other: Column) -> Column {
622        self.modulo(other)
623    }
624}
625
626impl BitAnd for Column {
627    type Output = Column;
628    fn bitand(self, other: Column) -> Column {
629        self.and(other)
630    }
631}
632
633impl BitOr for Column {
634    type Output = Column;
635    fn bitor(self, other: Column) -> Column {
636        self.or(other)
637    }
638}
639
640impl Not for Column {
641    type Output = Column;
642    fn not(self) -> Column {
643        Column {
644            expr: Expression::UnresolvedFunction(UnresolvedFunction::new("not", vec![self.expr])),
645        }
646    }
647}
648
649// Builder functions
650
651/// Mirrors `pyspark.sql.functions.col`.
652pub fn col(name: &str) -> Column {
653    Column {
654        expr: Expression::ColumnReference(ColumnReference::new(name)),
655    }
656}
657
658// A string is accepted anywhere a `Column` is (resolved as a column reference), so the
659// ergonomic `impl IntoIterator<Item = impl Into<Column>>` APIs accept `["a", "b"]` too.
660impl From<&str> for Column {
661    fn from(name: &str) -> Column {
662        col(name)
663    }
664}
665impl From<String> for Column {
666    fn from(name: String) -> Column {
667        col(&name)
668    }
669}
670impl From<&String> for Column {
671    fn from(name: &String) -> Column {
672        col(name)
673    }
674}
675
676/// Mirrors `pyspark.sql.functions.lit` for an integer value.
677///
678/// PySpark infers `IntegerType` for a Python int that fits in i32 and `LongType`
679/// otherwise (see `LiteralExpression._infer_type`), so `lit(1)` is `Integer(1)`,
680/// not `Long(1)`.
681pub fn lit(value: i64) -> Column {
682    let lit = if i32::try_from(value).is_ok() {
683        LiteralExpression::int(value as i32)
684    } else {
685        LiteralExpression::long(value)
686    };
687    Column {
688        expr: Expression::Literal(lit),
689    }
690}
691
692/// Create a literal from a string.
693pub fn lit_string(value: &str) -> Column {
694    Column {
695        expr: Expression::Literal(LiteralExpression::string(value)),
696    }
697}
698
699/// Create a literal from a double.
700pub fn lit_double(value: f64) -> Column {
701    Column {
702        expr: Expression::Literal(LiteralExpression::double(value)),
703    }
704}
705
706/// Create a literal from a boolean.
707pub fn lit_boolean(value: bool) -> Column {
708    Column {
709        expr: Expression::Literal(LiteralExpression::boolean(value)),
710    }
711}
712
713/// Mirrors `pyspark.sql.functions.when` - starts a CASE WHEN chain. Chain more
714/// branches with `Column::when` and finish with `Column::otherwise`.
715pub fn when(condition: Column, value: Column) -> Column {
716    Column {
717        expr: Expression::CaseWhen(Box::new(CaseWhen::new(vec![(condition.expr, value.expr)]))),
718    }
719}
720
721#[cfg(test)]
722mod tests {
723    use super::*;
724
725    #[test]
726    fn test_col_creation() {
727        let c = col("x");
728        assert!(matches!(c.expr, Expression::ColumnReference(_)));
729    }
730
731    #[test]
732    fn test_lit_creation() {
733        let c = lit(42);
734        assert!(matches!(c.expr, Expression::Literal(_)));
735    }
736
737    #[test]
738    fn test_addition() {
739        let c1 = col("a");
740        let c2 = lit(1);
741        let result = c1.add(c2);
742        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
743    }
744
745    #[test]
746    fn test_alias() {
747        let c = col("x");
748        let aliased = c.alias("y");
749        assert!(matches!(aliased.expr, Expression::Alias(_)));
750    }
751
752    #[test]
753    fn test_cast() {
754        let c = col("x");
755        let casted = c.cast(DataType::String {
756            collation: "UTF8_BINARY".to_string(),
757        });
758        assert!(matches!(casted.expr, Expression::Cast(_)));
759    }
760
761    #[test]
762    fn test_comparison_operators() {
763        let c1 = col("a");
764        let c2 = col("b");
765
766        // Test eq
767        let result = c1.clone().eq(c2.clone());
768        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
769
770        // Test ne
771        let result = col("a").ne(col("b"));
772        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
773
774        // Test gt
775        let result = col("a").gt(col("b"));
776        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
777
778        // Test lt
779        let result = col("a").lt(col("b"));
780        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
781
782        // Test ge
783        let result = col("a").ge(col("b"));
784        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
785
786        // Test le
787        let result = col("a").le(col("b"));
788        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
789    }
790
791    #[test]
792    fn test_arithmetic_operators() {
793        // Test sub
794        let result = col("a").sub(col("b"));
795        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
796
797        // Test mul
798        let result = col("a").mul(col("b"));
799        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
800
801        // Test div
802        let result = col("a").div(col("b"));
803        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
804
805        // Test modulo
806        let result = col("a").modulo(col("b"));
807        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
808    }
809
810    #[test]
811    fn test_logical_operators() {
812        // Test and
813        let result = col("a").and(col("b"));
814        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
815
816        // Test or
817        let result = col("a").or(col("b"));
818        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
819
820        // Test not
821        let result = col("a").not();
822        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
823
824        // Test neg
825        let result = col("a").neg();
826        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
827    }
828
829    #[test]
830    fn test_operator_traits() {
831        let c1 = col("a");
832        let c2 = col("b");
833
834        // Test Add trait
835        let result = c1.clone() + c2.clone();
836        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
837
838        // Test Sub trait
839        let result = col("a") - col("b");
840        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
841
842        // Test Mul trait
843        let result = col("a") * col("b");
844        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
845
846        // Test Div trait
847        let result = col("a") / col("b");
848        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
849
850        // Test Rem trait
851        let result = col("a") % col("b");
852        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
853
854        // Test BitAnd trait
855        let result = col("a") & col("b");
856        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
857
858        // Test BitOr trait
859        let result = col("a") | col("b");
860        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
861
862        // Test Not trait
863        let result = !col("a");
864        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
865    }
866
867    #[test]
868    fn test_from_implementations() {
869        // Test From<&str>
870        let c1: Column = "x".into();
871        assert!(matches!(c1.expr, Expression::ColumnReference(_)));
872
873        // Test From<String>
874        let c2: Column = "y".to_string().into();
875        assert!(matches!(c2.expr, Expression::ColumnReference(_)));
876
877        // Test From<&String>
878        let s = "z".to_string();
879        let c3: Column = (&s).into();
880        assert!(matches!(c3.expr, Expression::ColumnReference(_)));
881    }
882
883    #[test]
884    fn test_lit_functions() {
885        // Test lit with small number (fits in i32)
886        let c = lit(42);
887        assert!(matches!(
888            c.expr,
889            Expression::Literal(LiteralExpression::Integer(_))
890        ));
891
892        // Test lit with large number (needs i64)
893        let c = lit(5_000_000_000i64);
894        assert!(matches!(
895            c.expr,
896            Expression::Literal(LiteralExpression::Long(_))
897        ));
898
899        // Test lit_string
900        let c = lit_string("hello");
901        assert!(matches!(
902            c.expr,
903            Expression::Literal(LiteralExpression::String(_))
904        ));
905
906        // Test lit_double
907        let c = lit_double(3.14);
908        assert!(matches!(
909            c.expr,
910            Expression::Literal(LiteralExpression::Double(_))
911        ));
912
913        // Test lit_boolean
914        let c = lit_boolean(true);
915        assert!(matches!(
916            c.expr,
917            Expression::Literal(LiteralExpression::Boolean(_))
918        ));
919    }
920
921    #[test]
922    fn test_when_otherwise() {
923        // Test when
924        let cond = col("x").gt(lit(5));
925        let result = when(cond, lit(1));
926        assert!(matches!(result.expr, Expression::CaseWhen(_)));
927
928        // Test when with otherwise
929        let cond2 = col("y").lt(lit(10));
930        let result2 = result.when(cond2, lit(2));
931        assert!(matches!(result2.expr, Expression::CaseWhen(_)));
932
933        // Test otherwise
934        let final_result = result2.otherwise(lit(99));
935        assert!(matches!(final_result.expr, Expression::CaseWhen(_)));
936    }
937
938    #[test]
939    fn test_is_null_and_is_not_null() {
940        // Test is_null
941        let result = col("a").is_null();
942        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
943
944        // Test is_not_null
945        let result = col("a").is_not_null();
946        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
947    }
948
949    #[test]
950    fn test_getitem() {
951        // Test get_item
952        let result = col("array").get_item(lit(0));
953        assert!(matches!(result.expr, Expression::UnresolvedExtractValue(_)));
954    }
955
956    #[test]
957    fn test_between() {
958        let result = col("a").between(lit(1), lit(10));
959        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
960    }
961
962    #[test]
963    fn test_substring() {
964        let result = col("a").substr(lit(1), lit(3));
965        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
966    }
967
968    #[test]
969    fn test_various_methods() {
970        // Test cast_str
971        let result = col("a").cast_str("int");
972        assert!(matches!(result.expr, Expression::Cast(_)));
973
974        // Test desc
975        let result = col("a").desc();
976        assert!(matches!(result.expr, Expression::SortOrder(_)));
977
978        // Test desc_nulls_first
979        let result = col("a").desc_nulls_first();
980        assert!(matches!(result.expr, Expression::SortOrder(_)));
981
982        // Test desc_nulls_last
983        let result = col("a").desc_nulls_last();
984        assert!(matches!(result.expr, Expression::SortOrder(_)));
985
986        // Test asc
987        let result = col("a").asc();
988        assert!(matches!(result.expr, Expression::SortOrder(_)));
989
990        // Test asc_nulls_first
991        let result = col("a").asc_nulls_first();
992        assert!(matches!(result.expr, Expression::SortOrder(_)));
993
994        // Test asc_nulls_last
995        let result = col("a").asc_nulls_last();
996        assert!(matches!(result.expr, Expression::SortOrder(_)));
997    }
998
999    #[test]
1000    fn test_get_field() {
1001        // Test get_field
1002        let result = col("struct").get_field("field_name");
1003        assert!(matches!(result.expr, Expression::UnresolvedExtractValue(_)));
1004    }
1005
1006    #[test]
1007    fn test_string_functions() {
1008        // Test contains
1009        let result = col("a").contains(col("b"));
1010        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
1011
1012        // Test startswith
1013        let result = col("a").startswith(col("b"));
1014        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
1015
1016        // Test endswith
1017        let result = col("a").endswith(col("b"));
1018        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
1019
1020        // Test like
1021        let result = col("a").like("%pattern%");
1022        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
1023
1024        // Test rlike
1025        let result = col("a").rlike("[0-9]+");
1026        assert!(matches!(result.expr, Expression::UnresolvedFunction(_)));
1027    }
1028}
spark_connect/column.rs

spark_connect/
column.rs
