1use super::types::parse_type_name;
6use crate::column::Column;
7use polars::prelude::*;
8
9pub fn count(col: &Column) -> Column {
14 use polars::prelude::len;
15
16 let (expr, name) = if col.name() == "*" {
17 (len().cast(DataType::Int64), "count(1)".to_string())
18 } else {
19 (
20 col.expr().clone().count().cast(DataType::Int64),
21 format!("count({})", col.name()),
22 )
23 };
24 let mut c = Column::from_expr(expr, Some(name));
25 c.source_for_running_count = Some(col.name().to_string());
26 c
27}
28
29pub fn sum(col: &Column) -> Column {
33 let name = format!("sum({})", col.name());
34 let expr = col.expr().clone().cast(DataType::Float64).sum();
35 let mut c = Column::from_expr(expr, Some(name));
36 c.source_for_running = Some(col.name().to_string());
37 c
38}
39
40pub fn avg(col: &Column) -> Column {
44 let name = format!("avg({})", col.name());
45 let mut c = Column::from_expr(
46 col.expr().clone().cast(DataType::Float64).mean(),
47 Some(name),
48 );
49 c.source_for_running_mean = Some(col.name().to_string());
50 c
51}
52
53pub fn mean(col: &Column) -> Column {
55 avg(col)
56}
57
58pub fn max(col: &Column) -> Column {
60 let name = format!("max({})", col.name());
61 Column::from_expr(col.expr().clone().max(), Some(name))
62}
63
64pub fn min(col: &Column) -> Column {
66 let name = format!("min({})", col.name());
67 Column::from_expr(col.expr().clone().min(), Some(name))
68}
69
70pub fn first(col: &Column, ignorenulls: bool) -> Column {
72 let expr = if ignorenulls {
73 col.expr().clone().first_non_null()
74 } else {
75 col.expr().clone().first()
76 };
77 Column::from_expr(expr, None)
78}
79
80pub fn last(col: &Column, ignorenulls: bool) -> Column {
83 let _ = ignorenulls;
84 Column::from_last_agg(col)
85}
86
87pub fn any_value(col: &Column, ignorenulls: bool) -> Column {
89 let _ = ignorenulls;
90 Column::from_expr(col.expr().clone().first(), None)
91}
92
93pub fn count_if(col: &Column) -> Column {
95 use polars::prelude::DataType;
96 Column::from_expr(
97 col.expr().clone().cast(DataType::Int64).sum(),
98 Some("count_if".to_string()),
99 )
100}
101
102pub fn try_sum(col: &Column) -> Column {
104 Column::from_expr(col.expr().clone().sum(), Some("try_sum".to_string()))
105}
106
107pub fn try_avg(col: &Column) -> Column {
109 Column::from_expr(col.expr().clone().mean(), Some("try_avg".to_string()))
110}
111
112pub fn max_by(value_col: &Column, ord_col: &Column) -> Column {
114 use polars::prelude::{SortOptions, as_struct};
115 let st = as_struct(vec![
116 ord_col.expr().clone().alias("_ord"),
117 value_col.expr().clone().alias("_val"),
118 ]);
119 let e = st
120 .sort(SortOptions::default().with_order_descending(true))
121 .first()
122 .struct_()
123 .field_by_name("_val");
124 Column::from_expr(e, None)
125}
126
127pub fn min_by(value_col: &Column, ord_col: &Column) -> Column {
129 use polars::prelude::{SortOptions, as_struct};
130 let st = as_struct(vec![
131 ord_col.expr().clone().alias("_ord"),
132 value_col.expr().clone().alias("_val"),
133 ]);
134 let e = st
135 .sort(SortOptions::default())
136 .first()
137 .struct_()
138 .field_by_name("_val");
139 Column::from_expr(e, None)
140}
141
142pub fn collect_list(col: &Column) -> Column {
144 Column::from_expr(
145 col.expr().clone().implode(),
146 Some("collect_list".to_string()),
147 )
148}
149
150pub fn collect_set(col: &Column) -> Column {
152 Column::from_expr(
153 col.expr().clone().unique().implode(),
154 Some("collect_set".to_string()),
155 )
156}
157
158pub fn bool_and(col: &Column) -> Column {
160 Column::from_expr(col.expr().clone().all(true), Some("bool_and".to_string()))
161}
162
163pub fn every(col: &Column) -> Column {
165 Column::from_expr(col.expr().clone().all(true), Some("every".to_string()))
166}
167
168pub fn stddev(col: &Column) -> Column {
170 Column::from_expr(col.expr().clone().std(1), Some("stddev".to_string()))
171}
172
173pub fn variance(col: &Column) -> Column {
175 Column::from_expr(col.expr().clone().var(1), Some("variance".to_string()))
176}
177
178pub fn stddev_pop(col: &Column) -> Column {
180 Column::from_expr(col.expr().clone().std(0), Some("stddev_pop".to_string()))
181}
182
183pub fn stddev_samp(col: &Column) -> Column {
185 stddev(col)
186}
187
188pub fn std(col: &Column) -> Column {
190 stddev(col)
191}
192
193pub fn var_pop(col: &Column) -> Column {
195 Column::from_expr(col.expr().clone().var(0), Some("var_pop".to_string()))
196}
197
198pub fn var_samp(col: &Column) -> Column {
200 variance(col)
201}
202
203pub fn median(col: &Column) -> Column {
205 use polars::prelude::QuantileMethod;
206 Column::from_expr(
207 col.expr()
208 .clone()
209 .quantile(lit(0.5), QuantileMethod::Linear),
210 Some("median".to_string()),
211 )
212}
213
214pub fn approx_percentile(col: &Column, percentage: f64, _accuracy: Option<i32>) -> Column {
216 use polars::prelude::QuantileMethod;
217 Column::from_expr(
218 col.expr()
219 .clone()
220 .quantile(lit(percentage), QuantileMethod::Linear),
221 Some(format!("approx_percentile({percentage})")),
222 )
223}
224
225pub fn percentile_approx(col: &Column, percentage: f64, accuracy: Option<i32>) -> Column {
227 approx_percentile(col, percentage, accuracy)
228}
229
230pub fn mode(col: &Column) -> Column {
232 col.clone().mode()
233}
234
235pub fn count_distinct(col: &Column) -> Column {
237 use polars::prelude::DataType;
238 let name = format!("count_distinct({})", col.name());
239 Column::from_expr(
240 col.expr().clone().n_unique().cast(DataType::Int64),
241 Some(name),
242 )
243}
244
245pub fn approx_count_distinct(col: &Column, _rsd: Option<f64>) -> Column {
248 use polars::prelude::DataType;
249 let name = format!("approx_count_distinct({})", col.name());
250 let expr = col
251 .expr()
252 .clone()
253 .n_unique()
254 .cast(DataType::Int64)
255 .alias(&name);
256 Column::from_expr(expr, Some(name))
257}
258
259pub fn kurtosis(col: &Column) -> Column {
261 let name = format!("kurtosis({})", col.name());
262 Column::from_expr(
263 col.expr()
264 .clone()
265 .cast(DataType::Float64)
266 .kurtosis(true, true),
267 Some(name),
268 )
269}
270
271pub fn skewness(col: &Column) -> Column {
273 let name = format!("skewness({})", col.name());
274 Column::from_expr(
275 col.expr().clone().cast(DataType::Float64).skew(true),
276 Some(name),
277 )
278}
279
280fn covar_pop_expr_impl(e1: Expr, e2: Expr) -> Expr {
281 use polars::prelude::len;
282 let c1 = e1.clone().cast(DataType::Float64);
283 let c2 = e2.clone().cast(DataType::Float64);
284 let n = len().cast(DataType::Float64);
285 let sum_ab = (c1.clone() * c2.clone()).sum();
286 let sum_a = e1.sum().cast(DataType::Float64);
287 let sum_b = e2.sum().cast(DataType::Float64);
288 (sum_ab - sum_a * sum_b / n.clone()) / n
289}
290
291fn corr_expr_impl(e1: Expr, e2: Expr) -> Expr {
292 use polars::prelude::{len, lit, when};
293 let c1 = e1.clone().cast(DataType::Float64);
294 let c2 = e2.clone().cast(DataType::Float64);
295 let n = len().cast(DataType::Float64);
296 let n1 = (len() - lit(1)).cast(DataType::Float64);
297 let sum_ab = (c1.clone() * c2.clone()).sum();
298 let sum_a = e1.sum().cast(DataType::Float64);
299 let sum_b = e2.sum().cast(DataType::Float64);
300 let sum_a2 = (c1.clone() * c1).sum();
301 let sum_b2 = (c2.clone() * c2).sum();
302 let cov_samp = (sum_ab - sum_a.clone() * sum_b.clone() / n.clone()) / n1.clone();
303 let var_a = (sum_a2 - sum_a.clone() * sum_a / n.clone()) / n1.clone();
304 let var_b = (sum_b2 - sum_b.clone() * sum_b / n.clone()) / n1.clone();
305 let std_a = var_a.sqrt();
306 let std_b = var_b.sqrt();
307 when(len().gt(lit(1)))
308 .then(cov_samp / (std_a * std_b))
309 .otherwise(lit(f64::NAN))
310}
311
312fn covar_samp_expr_impl(e1: Expr, e2: Expr) -> Expr {
313 use polars::prelude::{len, lit, when};
314 let c1 = e1.clone().cast(DataType::Float64);
315 let c2 = e2.clone().cast(DataType::Float64);
316 let n = len().cast(DataType::Float64);
317 let sum_ab = (c1.clone() * c2.clone()).sum();
318 let sum_a = e1.sum().cast(DataType::Float64);
319 let sum_b = e2.sum().cast(DataType::Float64);
320 when(len().gt(lit(1)))
321 .then((sum_ab - sum_a * sum_b / n.clone()) / (len() - lit(1)).cast(DataType::Float64))
322 .otherwise(lit(f64::NAN))
323}
324
325pub fn covar_pop_expr(col1: &str, col2: &str) -> Expr {
327 use polars::prelude::col as pl_col;
328 covar_pop_expr_impl(pl_col(col1), pl_col(col2))
329}
330
331pub fn covar_pop(col1: &Column, col2: &Column) -> Column {
333 Column::from_expr(
334 covar_pop_expr_impl(col1.expr().clone(), col2.expr().clone()),
335 Some("covar_pop".to_string()),
336 )
337}
338
339pub fn corr(col1: &Column, col2: &Column) -> Column {
341 Column::from_expr(
342 corr_expr_impl(col1.expr().clone(), col2.expr().clone()),
343 Some("corr".to_string()),
344 )
345}
346
347pub fn covar_samp_expr(col1: &str, col2: &str) -> Expr {
349 use polars::prelude::col as pl_col;
350 covar_samp_expr_impl(pl_col(col1), pl_col(col2))
351}
352
353pub fn corr_expr(col1: &str, col2: &str) -> Expr {
355 use polars::prelude::col as pl_col;
356 corr_expr_impl(pl_col(col1), pl_col(col2))
357}
358
359fn regr_cond_and_sums(y_col: &str, x_col: &str) -> (Expr, Expr, Expr, Expr, Expr, Expr) {
362 use polars::prelude::col as pl_col;
363 let y = pl_col(y_col).cast(DataType::Float64);
364 let x = pl_col(x_col).cast(DataType::Float64);
365 let cond = y.clone().is_not_null().and(x.clone().is_not_null());
366 let n = y
367 .clone()
368 .filter(cond.clone())
369 .count()
370 .cast(DataType::Float64);
371 let sum_x = x.clone().filter(cond.clone()).sum();
372 let sum_y = y.clone().filter(cond.clone()).sum();
373 let sum_xx = (x.clone() * x.clone()).filter(cond.clone()).sum();
374 let sum_yy = (y.clone() * y.clone()).filter(cond.clone()).sum();
375 let sum_xy = (x * y).filter(cond).sum();
376 (n, sum_x, sum_y, sum_xx, sum_yy, sum_xy)
377}
378
379pub fn regr_count_expr(y_col: &str, x_col: &str) -> Expr {
381 let (n, ..) = regr_cond_and_sums(y_col, x_col);
382 n
383}
384
385pub fn regr_avgx_expr(y_col: &str, x_col: &str) -> Expr {
387 use polars::prelude::{lit, when};
388 let (n, sum_x, ..) = regr_cond_and_sums(y_col, x_col);
389 when(n.clone().gt(lit(0.0)))
390 .then(sum_x / n)
391 .otherwise(lit(f64::NAN))
392}
393
394pub fn regr_avgy_expr(y_col: &str, x_col: &str) -> Expr {
396 use polars::prelude::{lit, when};
397 let (n, _, sum_y, ..) = regr_cond_and_sums(y_col, x_col);
398 when(n.clone().gt(lit(0.0)))
399 .then(sum_y / n)
400 .otherwise(lit(f64::NAN))
401}
402
403pub fn regr_sxx_expr(y_col: &str, x_col: &str) -> Expr {
405 use polars::prelude::{lit, when};
406 let (n, sum_x, _, sum_xx, ..) = regr_cond_and_sums(y_col, x_col);
407 when(n.clone().gt(lit(0.0)))
408 .then(sum_xx - sum_x.clone() * sum_x / n)
409 .otherwise(lit(f64::NAN))
410}
411
412pub fn regr_syy_expr(y_col: &str, x_col: &str) -> Expr {
414 use polars::prelude::{lit, when};
415 let (n, _, sum_y, _, sum_yy, _) = regr_cond_and_sums(y_col, x_col);
416 when(n.clone().gt(lit(0.0)))
417 .then(sum_yy - sum_y.clone() * sum_y / n)
418 .otherwise(lit(f64::NAN))
419}
420
421pub fn regr_sxy_expr(y_col: &str, x_col: &str) -> Expr {
423 use polars::prelude::{lit, when};
424 let (n, sum_x, sum_y, _, _, sum_xy) = regr_cond_and_sums(y_col, x_col);
425 when(n.clone().gt(lit(0.0)))
426 .then(sum_xy - sum_x * sum_y / n)
427 .otherwise(lit(f64::NAN))
428}
429
430pub fn regr_slope_expr(y_col: &str, x_col: &str) -> Expr {
432 use polars::prelude::{lit, when};
433 let (n, sum_x, sum_y, sum_xx, _sum_yy, sum_xy) = regr_cond_and_sums(y_col, x_col);
434 let regr_sxx = sum_xx.clone() - sum_x.clone() * sum_x.clone() / n.clone();
435 let regr_sxy = sum_xy - sum_x * sum_y / n.clone();
436 when(n.gt(lit(1.0)).and(regr_sxx.clone().gt(lit(0.0))))
437 .then(regr_sxy / regr_sxx)
438 .otherwise(lit(f64::NAN))
439}
440
441pub fn regr_intercept_expr(y_col: &str, x_col: &str) -> Expr {
443 use polars::prelude::{lit, when};
444 let (n, sum_x, sum_y, sum_xx, _, sum_xy) = regr_cond_and_sums(y_col, x_col);
445 let regr_sxx = sum_xx - sum_x.clone() * sum_x.clone() / n.clone();
446 let regr_sxy = sum_xy.clone() - sum_x.clone() * sum_y.clone() / n.clone();
447 let slope = regr_sxy.clone() / regr_sxx.clone();
448 let avg_y = sum_y / n.clone();
449 let avg_x = sum_x / n.clone();
450 when(n.gt(lit(1.0)).and(regr_sxx.clone().gt(lit(0.0))))
451 .then(avg_y - slope * avg_x)
452 .otherwise(lit(f64::NAN))
453}
454
455pub fn regr_r2_expr(y_col: &str, x_col: &str) -> Expr {
457 use polars::prelude::{lit, when};
458 let (n, sum_x, sum_y, sum_xx, sum_yy, sum_xy) = regr_cond_and_sums(y_col, x_col);
459 let regr_sxx = sum_xx - sum_x.clone() * sum_x.clone() / n.clone();
460 let regr_syy = sum_yy - sum_y.clone() * sum_y.clone() / n.clone();
461 let regr_sxy = sum_xy - sum_x * sum_y / n;
462 when(
463 regr_sxx
464 .clone()
465 .gt(lit(0.0))
466 .and(regr_syy.clone().gt(lit(0.0))),
467 )
468 .then(regr_sxy.clone() * regr_sxy / (regr_sxx * regr_syy))
469 .otherwise(lit(f64::NAN))
470}
471
472pub fn when(condition: &Column) -> WhenBuilder {
485 WhenBuilder::new(expr_coerce_to_boolean(condition.expr().clone()))
486}
487
488pub fn when_then_otherwise_null(condition: &Column, value: &Column) -> Column {
491 use polars::prelude::*;
492 let null_expr = lit(NULL);
493 let expr = polars::prelude::when(expr_coerce_to_boolean(condition.expr().clone()))
494 .then(value.expr().clone())
495 .otherwise(null_expr);
496 crate::column::Column::from_expr(expr, None)
497}
498
499pub struct WhenBuilder {
501 condition: Expr,
502}
503
504impl WhenBuilder {
505 fn new(condition: Expr) -> Self {
506 WhenBuilder { condition }
507 }
508
509 pub fn then(self, value: &Column) -> ThenBuilder {
511 use polars::prelude::*;
512 let when_then = when(self.condition).then(value.expr().clone());
513 ThenBuilder::new(when_then)
514 }
515
516 pub fn otherwise(self, _value: &Column) -> Column {
521 panic!(
524 "when().otherwise() requires .then() to be called first. Use when(cond).then(val1).otherwise(val2)"
525 );
526 }
527}
528
529pub struct ThenBuilder {
531 state: WhenThenState,
532}
533
534enum WhenThenState {
535 Single(Box<polars::prelude::Then>),
536 Chained(Box<polars::prelude::ChainedThen>),
537}
538
539pub struct ChainedWhenBuilder {
541 inner: polars::prelude::ChainedWhen,
542}
543
544impl ThenBuilder {
545 fn new(when_then: polars::prelude::Then) -> Self {
546 ThenBuilder {
547 state: WhenThenState::Single(Box::new(when_then)),
548 }
549 }
550
551 fn new_chained(chained: polars::prelude::ChainedThen) -> Self {
552 ThenBuilder {
553 state: WhenThenState::Chained(Box::new(chained)),
554 }
555 }
556
557 pub fn when(self, condition: &Column) -> ChainedWhenBuilder {
560 let cond = expr_coerce_to_boolean(condition.expr().clone());
561 let chained_when = match self.state {
562 WhenThenState::Single(t) => t.when(cond),
563 WhenThenState::Chained(ct) => ct.when(cond),
564 };
565 ChainedWhenBuilder {
566 inner: chained_when,
567 }
568 }
569
570 pub fn otherwise(self, value: &Column) -> Column {
572 let expr = match self.state {
573 WhenThenState::Single(t) => t.otherwise(value.expr().clone()),
574 WhenThenState::Chained(ct) => ct.otherwise(value.expr().clone()),
575 };
576 crate::column::Column::from_expr(expr, None)
577 }
578}
579
580impl ChainedWhenBuilder {
581 pub fn then(self, value: &Column) -> ThenBuilder {
583 ThenBuilder::new_chained(self.inner.then(value.expr().clone()))
584 }
585}
586
587pub fn upper(column: &Column) -> Column {
589 column.clone().upper()
590}
591
592pub fn lower(column: &Column) -> Column {
594 column.clone().lower()
595}
596
597pub fn substring(column: &Column, start: i64, length: Option<i64>) -> Column {
599 column.clone().substr(start, length)
600}
601
602pub fn length(column: &Column) -> Column {
604 column.clone().length()
605}
606
607pub fn trim(column: &Column) -> Column {
609 column.clone().trim()
610}
611
612pub fn ltrim(column: &Column) -> Column {
614 column.clone().ltrim()
615}
616
617pub fn rtrim(column: &Column) -> Column {
619 column.clone().rtrim()
620}
621
622pub fn btrim(column: &Column, trim_str: Option<&str>) -> Column {
624 column.clone().btrim(trim_str)
625}
626
627pub fn locate(substr: &str, column: &Column, pos: i64) -> Column {
629 column.clone().locate(substr, pos)
630}
631
632pub fn conv(column: &Column, from_base: i32, to_base: i32) -> Column {
634 column.clone().conv(from_base, to_base)
635}
636
637pub fn hex(column: &Column) -> Column {
639 column.clone().hex()
640}
641
642pub fn unhex(column: &Column) -> Column {
644 column.clone().unhex()
645}
646
647pub fn encode(column: &Column, charset: &str) -> Column {
649 column.clone().encode(charset)
650}
651
652pub fn decode(column: &Column, charset: &str) -> Column {
654 column.clone().decode(charset)
655}
656
657pub fn to_binary(column: &Column, fmt: &str) -> Column {
659 column.clone().to_binary(fmt)
660}
661
662pub fn try_to_binary(column: &Column, fmt: &str) -> Column {
664 column.clone().try_to_binary(fmt)
665}
666
667pub fn aes_encrypt(column: &Column, key: &str) -> Column {
669 column.clone().aes_encrypt(key)
670}
671
672pub fn aes_decrypt(column: &Column, key: &str) -> Column {
674 column.clone().aes_decrypt(key)
675}
676
677pub fn try_aes_decrypt(column: &Column, key: &str) -> Column {
679 column.clone().try_aes_decrypt(key)
680}
681
682pub fn bin(column: &Column) -> Column {
684 column.clone().bin()
685}
686
687pub fn getbit(column: &Column, pos: i64) -> Column {
689 column.clone().getbit(pos)
690}
691
692pub fn bit_and(left: &Column, right: &Column) -> Column {
694 left.clone().bit_and(right)
695}
696
697pub fn bit_or(left: &Column, right: &Column) -> Column {
699 left.clone().bit_or(right)
700}
701
702pub fn bit_xor(left: &Column, right: &Column) -> Column {
704 left.clone().bit_xor(right)
705}
706
707pub fn bit_count(column: &Column) -> Column {
709 column.clone().bit_count()
710}
711
712pub fn bitwise_not(column: &Column) -> Column {
714 column.clone().bitwise_not()
715}
716
717pub fn bitmap_bit_position(column: &Column) -> Column {
721 use polars::prelude::DataType;
722 let expr = column.expr().clone().cast(DataType::Int32);
723 Column::from_expr(expr, None)
724}
725
726pub fn bitmap_bucket_number(column: &Column) -> Column {
728 use polars::prelude::DataType;
729 let expr = column.expr().clone().cast(DataType::Int64) / lit(32768i64);
730 Column::from_expr(expr, None)
731}
732
733pub fn bitmap_count(column: &Column) -> Column {
735 use polars::prelude::{DataType, Field};
736 let expr = column.expr().clone().map(
737 |s| crate::column::expect_col(crate::udfs::apply_bitmap_count(s)),
738 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Int64)),
739 );
740 Column::from_expr(expr, None)
741}
742
743pub fn bitmap_construct_agg(column: &Column) -> polars::prelude::Expr {
746 use polars::prelude::{DataType, Field};
747 column.expr().clone().implode().map(
748 |s| crate::column::expect_col(crate::udfs::apply_bitmap_construct_agg(s)),
749 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Binary)),
750 )
751}
752
753pub fn bitmap_or_agg(column: &Column) -> polars::prelude::Expr {
755 use polars::prelude::{DataType, Field};
756 column.expr().clone().implode().map(
757 |s| crate::column::expect_col(crate::udfs::apply_bitmap_or_agg(s)),
758 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Binary)),
759 )
760}
761
762pub fn bit_get(column: &Column, pos: i64) -> Column {
764 getbit(column, pos)
765}
766
767pub fn assert_true(column: &Column, err_msg: Option<&str>) -> Column {
770 column.clone().assert_true(err_msg)
771}
772
773pub fn raise_error(message: &str) -> Column {
775 let msg = message.to_string();
776 let expr = lit(0i64).map(
777 move |_col| -> PolarsResult<polars::prelude::Column> {
778 Err(PolarsError::ComputeError(msg.clone().into()))
779 },
780 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Int64)),
781 );
782 Column::from_expr(expr, Some("raise_error".to_string()))
783}
784
785pub fn spark_partition_id() -> Column {
787 Column::from_expr(lit(0i32), Some("spark_partition_id".to_string()))
788}
789
790pub fn input_file_name() -> Column {
792 Column::from_expr(lit(""), Some("input_file_name".to_string()))
793}
794
795pub fn monotonically_increasing_id() -> Column {
798 Column::from_expr(lit(0i64), Some("monotonically_increasing_id".to_string()))
799}
800
801pub fn current_catalog() -> Column {
803 Column::from_expr(lit("spark_catalog"), Some("current_catalog".to_string()))
804}
805
806pub fn current_database() -> Column {
808 Column::from_expr(lit("default"), Some("current_database".to_string()))
809}
810
811pub fn current_schema() -> Column {
813 Column::from_expr(lit("default"), Some("current_schema".to_string()))
814}
815
816pub fn current_user() -> Column {
818 Column::from_expr(lit("unknown"), Some("current_user".to_string()))
819}
820
821pub fn user() -> Column {
823 Column::from_expr(lit("unknown"), Some("user".to_string()))
824}
825
826pub fn rand(seed: Option<u64>) -> Column {
829 Column::from_rand(seed)
830}
831
832pub fn randn(seed: Option<u64>) -> Column {
835 Column::from_randn(seed)
836}
837
838pub fn call_udf(name: &str, cols: &[Column]) -> Result<Column, PolarsError> {
841 use polars::prelude::Column as PlColumn;
842
843 let (registry, case_sensitive) =
844 crate::udf_context::get_thread_udf_context().ok_or_else(|| {
845 PolarsError::InvalidOperation(
846 "call_udf: no session. Use SparkSession.builder().get_or_create() first.".into(),
847 )
848 })?;
849
850 let udf = registry
852 .get_rust_udf(name, case_sensitive)?
853 .ok_or_else(|| {
854 PolarsError::InvalidOperation(format!("call_udf: UDF '{name}' not found").into())
855 })?;
856
857 let exprs: Vec<Expr> = cols.iter().map(|c| c.expr().clone()).collect();
858 let output_type = DataType::String; let expr = if exprs.len() == 1 {
861 let udf = udf.clone();
862 exprs.into_iter().next().unwrap().map(
863 move |c| {
864 let s = c.take_materialized_series();
865 udf.apply(&[s]).map(|out| PlColumn::new("_".into(), out))
866 },
867 move |_schema, field| Ok(Field::new(field.name().clone(), output_type.clone())),
868 )
869 } else {
870 let udf = udf.clone();
871 let first = exprs[0].clone();
872 let rest: Vec<Expr> = exprs[1..].to_vec();
873 first.map_many(
874 move |columns| {
875 let series: Vec<Series> = columns
876 .iter_mut()
877 .map(|c| std::mem::take(c).take_materialized_series())
878 .collect();
879 udf.apply(&series).map(|out| PlColumn::new("_".into(), out))
880 },
881 &rest,
882 move |_schema, fields| Ok(Field::new(fields[0].name().clone(), output_type.clone())),
883 )
884 };
885
886 Ok(Column::from_expr(expr, Some(format!("{name}()"))))
887}
888
889pub fn arrays_overlap(left: &Column, right: &Column) -> Column {
891 left.clone().arrays_overlap(right)
892}
893
894pub fn arrays_zip(left: &Column, right: &Column) -> Column {
896 left.clone().arrays_zip(right)
897}
898
899pub fn explode_outer(column: &Column) -> Column {
901 column.clone().explode_outer()
902}
903
904pub fn posexplode_outer(column: &Column) -> (Column, Column) {
906 column.clone().posexplode_outer()
907}
908
909pub fn array_agg(column: &Column) -> Column {
911 column.clone().array_agg()
912}
913
914pub fn transform_keys(column: &Column, key_expr: Expr) -> Column {
916 column.clone().transform_keys(key_expr)
917}
918
919pub fn transform_values(column: &Column, value_expr: Expr) -> Column {
921 column.clone().transform_values(value_expr)
922}
923
924pub fn str_to_map(
926 column: &Column,
927 pair_delim: Option<&str>,
928 key_value_delim: Option<&str>,
929) -> Column {
930 let pd = pair_delim.unwrap_or(",");
931 let kvd = key_value_delim.unwrap_or(":");
932 column.clone().str_to_map(pd, kvd)
933}
934
935pub fn regexp_extract(column: &Column, pattern: &str, group_index: usize) -> Column {
937 column.clone().regexp_extract(pattern, group_index)
938}
939
940pub fn regexp_replace(column: &Column, pattern: &str, replacement: &str) -> Column {
942 column.clone().regexp_replace(pattern, replacement)
943}
944
945pub fn split(column: &Column, delimiter: &str, limit: Option<i32>) -> Column {
947 column.clone().split(delimiter, limit)
948}
949
950pub fn initcap(column: &Column) -> Column {
952 column.clone().initcap()
953}
954
955pub fn regexp_extract_all(column: &Column, pattern: &str) -> Column {
957 column.clone().regexp_extract_all(pattern)
958}
959
960pub fn regexp_like(column: &Column, pattern: &str) -> Column {
962 column.clone().regexp_like(pattern)
963}
964
965pub fn regexp_count(column: &Column, pattern: &str) -> Column {
967 column.clone().regexp_count(pattern)
968}
969
970pub fn regexp_substr(column: &Column, pattern: &str) -> Column {
972 column.clone().regexp_substr(pattern)
973}
974
975pub fn split_part(column: &Column, delimiter: &str, part_num: i64) -> Column {
977 column.clone().split_part(delimiter, part_num)
978}
979
980pub fn regexp_instr(column: &Column, pattern: &str, group_idx: Option<usize>) -> Column {
982 column.clone().regexp_instr(pattern, group_idx)
983}
984
985pub fn find_in_set(str_column: &Column, set_column: &Column) -> Column {
987 str_column.clone().find_in_set(set_column)
988}
989
990pub fn format_string(format: &str, columns: &[&Column]) -> Column {
992 use polars::prelude::*;
993 if columns.is_empty() {
994 panic!("format_string needs at least one column");
995 }
996 let format_owned = format.to_string();
997 let args: Vec<Expr> = columns.iter().skip(1).map(|c| c.expr().clone()).collect();
998 let expr = columns[0].expr().clone().map_many(
999 move |cols| {
1000 crate::column::expect_col(crate::udfs::apply_format_string(cols, &format_owned))
1001 },
1002 &args,
1003 |_schema, fields| Ok(Field::new(fields[0].name().clone(), DataType::String)),
1004 );
1005 crate::column::Column::from_expr(expr, None)
1006}
1007
1008pub fn printf(format: &str, columns: &[&Column]) -> Column {
1010 format_string(format, columns)
1011}
1012
1013pub fn repeat(column: &Column, n: i32) -> Column {
1015 column.clone().repeat(n)
1016}
1017
1018pub fn reverse(column: &Column) -> Column {
1020 column.clone().reverse()
1021}
1022
1023pub fn instr(column: &Column, substr: &str) -> Column {
1025 column.clone().instr(substr)
1026}
1027
1028pub fn position(substr: &str, column: &Column) -> Column {
1030 column.clone().instr(substr)
1031}
1032
1033pub fn ascii(column: &Column) -> Column {
1035 column.clone().ascii()
1036}
1037
1038pub fn format_number(column: &Column, decimals: u32) -> Column {
1040 column.clone().format_number(decimals)
1041}
1042
1043pub fn overlay(column: &Column, replace: &str, pos: i64, length: i64) -> Column {
1045 column.clone().overlay(replace, pos, length)
1046}
1047
1048pub fn char(column: &Column) -> Column {
1050 column.clone().char()
1051}
1052
1053pub fn chr(column: &Column) -> Column {
1055 column.clone().chr()
1056}
1057
1058pub fn base64(column: &Column) -> Column {
1060 column.clone().base64()
1061}
1062
1063pub fn unbase64(column: &Column) -> Column {
1065 column.clone().unbase64()
1066}
1067
1068pub fn sha1(column: &Column) -> Column {
1070 column.clone().sha1()
1071}
1072
1073pub fn sha2(column: &Column, bit_length: i32) -> Column {
1075 column.clone().sha2(bit_length)
1076}
1077
1078pub fn md5(column: &Column) -> Column {
1080 column.clone().md5()
1081}
1082
1083pub fn lpad(column: &Column, length: i32, pad: &str) -> Column {
1085 column.clone().lpad(length, pad)
1086}
1087
1088pub fn rpad(column: &Column, length: i32, pad: &str) -> Column {
1090 column.clone().rpad(length, pad)
1091}
1092
1093pub fn translate(column: &Column, from_str: &str, to_str: &str) -> Column {
1095 column.clone().translate(from_str, to_str)
1096}
1097
1098pub fn mask(
1100 column: &Column,
1101 upper_char: Option<char>,
1102 lower_char: Option<char>,
1103 digit_char: Option<char>,
1104 other_char: Option<char>,
1105) -> Column {
1106 column
1107 .clone()
1108 .mask(upper_char, lower_char, digit_char, other_char)
1109}
1110
1111pub fn substring_index(column: &Column, delimiter: &str, count: i64) -> Column {
1113 column.clone().substring_index(delimiter, count)
1114}
1115
1116pub fn left(column: &Column, n: i64) -> Column {
1118 column.clone().left(n)
1119}
1120
1121pub fn right(column: &Column, n: i64) -> Column {
1123 column.clone().right(n)
1124}
1125
1126pub fn replace(column: &Column, search: &str, replacement: &str) -> Column {
1128 column.clone().replace(search, replacement)
1129}
1130
1131pub fn startswith(column: &Column, prefix: &str) -> Column {
1133 column.clone().startswith(prefix)
1134}
1135
1136pub fn endswith(column: &Column, suffix: &str) -> Column {
1138 column.clone().endswith(suffix)
1139}
1140
1141pub fn contains(column: &Column, substring: &str) -> Column {
1143 column.clone().contains(substring)
1144}
1145
1146pub fn like(column: &Column, pattern: &str, escape_char: Option<char>) -> Column {
1149 column.clone().like(pattern, escape_char)
1150}
1151
1152pub fn ilike(column: &Column, pattern: &str, escape_char: Option<char>) -> Column {
1155 column.clone().ilike(pattern, escape_char)
1156}
1157
1158pub fn rlike(column: &Column, pattern: &str) -> Column {
1160 column.clone().regexp_like(pattern)
1161}
1162
1163pub fn regexp(column: &Column, pattern: &str) -> Column {
1165 rlike(column, pattern)
1166}
1167
1168pub fn soundex(column: &Column) -> Column {
1170 column.clone().soundex()
1171}
1172
1173pub fn levenshtein(column: &Column, other: &Column) -> Column {
1175 column.clone().levenshtein(other)
1176}
1177
1178pub fn crc32(column: &Column) -> Column {
1180 column.clone().crc32()
1181}
1182
1183pub fn xxhash64(column: &Column) -> Column {
1185 column.clone().xxhash64()
1186}
1187
1188pub fn abs(column: &Column) -> Column {
1190 column.clone().abs()
1191}
1192
1193pub fn ceil(column: &Column) -> Column {
1195 column.clone().ceil()
1196}
1197
1198pub fn floor(column: &Column) -> Column {
1200 column.clone().floor()
1201}
1202
1203pub fn round(column: &Column, scale: i32) -> Column {
1205 column.clone().round(scale)
1206}
1207
1208pub fn bround(column: &Column, scale: i32) -> Column {
1210 column.clone().bround(scale)
1211}
1212
1213pub fn negate(column: &Column) -> Column {
1215 column.clone().negate()
1216}
1217
1218pub fn negative(column: &Column) -> Column {
1220 negate(column)
1221}
1222
1223pub fn positive(column: &Column) -> Column {
1225 column.clone()
1226}
1227
1228pub fn cot(column: &Column) -> Column {
1230 column.clone().cot()
1231}
1232
1233pub fn csc(column: &Column) -> Column {
1235 column.clone().csc()
1236}
1237
1238pub fn sec(column: &Column) -> Column {
1240 column.clone().sec()
1241}
1242
1243pub fn e() -> Column {
1245 Column::from_expr(lit(std::f64::consts::E), Some("e".to_string()))
1246}
1247
1248pub fn pi() -> Column {
1250 Column::from_expr(lit(std::f64::consts::PI), Some("pi".to_string()))
1251}
1252
1253pub fn sqrt(column: &Column) -> Column {
1255 column.clone().sqrt()
1256}
1257
1258pub fn pow(column: &Column, exp: i64) -> Column {
1260 column.clone().pow(exp)
1261}
1262
1263pub fn exp(column: &Column) -> Column {
1265 column.clone().exp()
1266}
1267
1268pub fn log(column: &Column) -> Column {
1270 column.clone().log()
1271}
1272
1273pub fn log_with_base(column: &Column, base: f64) -> Column {
1275 crate::column::Column::from_expr(column.expr().clone().log(lit(base)), None)
1276}
1277
1278pub fn sin(column: &Column) -> Column {
1280 column.clone().sin()
1281}
1282
1283pub fn cos(column: &Column) -> Column {
1285 column.clone().cos()
1286}
1287
1288pub fn tan(column: &Column) -> Column {
1290 column.clone().tan()
1291}
1292
1293pub fn asin(column: &Column) -> Column {
1295 column.clone().asin()
1296}
1297
1298pub fn acos(column: &Column) -> Column {
1300 column.clone().acos()
1301}
1302
1303pub fn atan(column: &Column) -> Column {
1305 column.clone().atan()
1306}
1307
1308pub fn atan2(y: &Column, x: &Column) -> Column {
1310 y.clone().atan2(x)
1311}
1312
1313pub fn degrees(column: &Column) -> Column {
1315 column.clone().degrees()
1316}
1317
1318pub fn radians(column: &Column) -> Column {
1320 column.clone().radians()
1321}
1322
1323pub fn signum(column: &Column) -> Column {
1325 column.clone().signum()
1326}
1327
1328pub fn sign(column: &Column) -> Column {
1330 signum(column)
1331}
1332
1333fn unwrap_alias_literal(expr: &Expr) -> Expr {
1344 let mut e = expr.clone();
1345 while let Expr::Alias(inner, _) = &e {
1346 if matches!(inner.as_ref(), Expr::Literal(_)) {
1347 e = (**inner).clone();
1348 } else {
1349 break;
1350 }
1351 }
1352 e
1353}
1354
1355pub fn expr_coerce_to_boolean(expr: Expr) -> Expr {
1356 use polars::prelude::Operator;
1357 use std::sync::Arc;
1358 match &expr {
1359 Expr::BinaryExpr { left, op, right }
1360 if matches!(
1361 op,
1362 Operator::Eq
1363 | Operator::NotEq
1364 | Operator::Lt
1365 | Operator::LtEq
1366 | Operator::Gt
1367 | Operator::GtEq
1368 ) =>
1369 {
1370 Expr::BinaryExpr {
1373 left: Arc::new(unwrap_alias_literal(left)),
1374 op: *op,
1375 right: Arc::new(unwrap_alias_literal(right)),
1376 }
1377 }
1378 Expr::BinaryExpr { left, op, right } if matches!(op, Operator::And | Operator::Or) => {
1379 let left_c = expr_coerce_to_boolean((**left).clone());
1380 let right_c = expr_coerce_to_boolean((**right).clone());
1381 Expr::BinaryExpr {
1382 left: std::sync::Arc::new(left_c),
1383 op: *op,
1384 right: std::sync::Arc::new(right_c),
1385 }
1386 }
1387 Expr::Alias(inner, _name) => {
1388 let coerced = expr_coerce_to_boolean((**inner).clone());
1389 if matches!(
1391 &coerced,
1392 Expr::BinaryExpr { op, .. } if matches!(
1393 op,
1394 Operator::Eq
1395 | Operator::NotEq
1396 | Operator::Lt
1397 | Operator::LtEq
1398 | Operator::Gt
1399 | Operator::GtEq
1400 )
1401 ) {
1402 coerced
1403 } else {
1404 Expr::Alias(Arc::new(coerced), _name.clone())
1405 }
1406 }
1407 _ => expr.map(
1408 move |col| crate::column::expect_col(crate::udfs::apply_string_to_boolean(col, false)),
1409 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Boolean)),
1410 ),
1411 }
1412}
1413
1414fn cast_impl(column: &Column, type_name: &str, strict: bool) -> Result<Column, String> {
1415 let dtype = parse_type_name(type_name)?;
1416 let base_name = column.name().to_string();
1417
1418 if dtype == DataType::Boolean {
1419 let out_name = base_name.clone();
1420 let expr = column.expr().clone().map(
1421 move |col| crate::column::expect_col(crate::udfs::apply_string_to_boolean(col, strict)),
1422 move |_schema, _field| Ok(Field::new(out_name.clone().into(), DataType::Boolean)),
1423 );
1424 return Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)));
1426 }
1427 if dtype == DataType::Date {
1428 let expr = column.expr().clone().map(
1429 move |col| crate::column::expect_col(crate::udfs::apply_string_to_date(col, strict)),
1430 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Date)),
1431 );
1432 return Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)));
1433 }
1434 if matches!(dtype, DataType::Datetime(_, _)) {
1435 use polars::datatypes::TimeUnit;
1436 let expr = column.expr().clone().map(
1437 move |col| {
1438 crate::column::expect_col(crate::udfs::apply_string_to_datetime(col, strict))
1439 },
1440 |_schema, field| {
1441 Ok(Field::new(
1442 field.name().clone(),
1443 DataType::Datetime(TimeUnit::Microseconds, None),
1444 ))
1445 },
1446 );
1447 return Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)));
1448 }
1449 if dtype == DataType::Int32 || dtype == DataType::Int64 {
1450 let is_plain_column = matches!(column.expr(), Expr::Column(_));
1456 if !is_plain_column {
1457 let expr = column.expr().clone().strict_cast(dtype);
1458 return Ok(Column::from_expr(expr, Some(base_name)));
1459 }
1460 let target = dtype.clone();
1461 let expr = column.expr().clone().map(
1462 move |col| {
1463 crate::column::expect_col(crate::udfs::apply_string_to_int(
1464 col,
1465 false, target.clone(),
1467 ))
1468 },
1469 move |_schema, field| Ok(Field::new(field.name().clone(), dtype.clone())),
1470 );
1471 return Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)));
1472 }
1473 if dtype == DataType::Float64 {
1474 let expr = column.expr().clone().map(
1476 move |col| {
1477 crate::column::expect_col(crate::udfs::apply_string_to_double(
1478 col, false, ))
1480 },
1481 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Float64)),
1482 );
1483 return Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)));
1484 }
1485 let expr = if strict {
1486 column.expr().clone().strict_cast(dtype)
1487 } else {
1488 column.expr().clone().cast(dtype)
1489 };
1490 Ok(Column::from_expr(expr.alias(&base_name), Some(base_name)))
1492}
1493
1494pub fn cast(column: &Column, type_name: &str) -> Result<Column, String> {
1499 cast_impl(column, type_name, true)
1500}
1501
1502pub fn try_cast(column: &Column, type_name: &str) -> Result<Column, String> {
1506 cast_impl(column, type_name, false)
1507}
1508
1509pub fn to_char(column: &Column, format: Option<&str>) -> Result<Column, String> {
1513 match format {
1514 Some(fmt) => Ok(column
1515 .clone()
1516 .date_format(&crate::udfs::pyspark_format_to_chrono(fmt))),
1517 None => cast(column, "string"),
1518 }
1519}
1520
1521pub fn to_varchar(column: &Column, format: Option<&str>) -> Result<Column, String> {
1523 to_char(column, format)
1524}
1525
1526pub fn to_number(column: &Column, _format: Option<&str>) -> Result<Column, String> {
1529 cast(column, "double")
1530}
1531
1532pub fn try_to_number(column: &Column, _format: Option<&str>) -> Result<Column, String> {
1535 try_cast(column, "double")
1536}
1537
1538pub fn to_timestamp_fused_strip_fraction(column: &Column, format: &str) -> Result<Column, String> {
1542 use polars::prelude::{DataType, Field, TimeUnit};
1543 let ref_ts = chrono::Utc::now();
1544 let out_name = format!("to_timestamp({}, {format})", column.name());
1545 let out_name2 = out_name.clone();
1546 let format_owned = format.to_string();
1547 let expr = column.expr().clone().map(
1548 move |s| {
1549 crate::column::expect_col(crate::udfs::apply_to_timestamp_strip_fraction_recent_null(
1550 s,
1551 &format_owned,
1552 ref_ts,
1553 ))
1554 },
1555 move |_schema, field| match field.dtype() {
1556 DataType::String => Ok(Field::new(
1557 out_name2.clone().into(),
1558 DataType::Datetime(TimeUnit::Microseconds, None),
1559 )),
1560 _ => Err(polars::prelude::PolarsError::ComputeError(
1561 "to_timestamp fused path requires StringType".into(),
1562 )),
1563 },
1564 );
1565 Ok(crate::column::Column::from_expr(
1566 expr.alias(&out_name),
1567 Some(out_name),
1568 ))
1569}
1570
1571pub fn to_timestamp(column: &Column, format: Option<&str>) -> Result<Column, String> {
1574 use polars::prelude::{DataType, Field, TimeUnit};
1575 let fmt_owned = format.map(|s| s.to_string());
1576 let out_name = match format {
1577 None => format!("to_timestamp({})", column.name()),
1578 Some(f) => format!("to_timestamp({}, {f})", column.name()),
1579 };
1580 let out_name2 = out_name.clone();
1581 let use_recent_null = format == Some("yyyy-MM-dd'T'HH:mm:ss")
1584 && !crate::udfs::is_simple_column_ref(column.expr());
1585 let base_expr = column.expr().clone();
1586 let expr = base_expr.map(
1587 move |s| {
1588 crate::column::expect_col(crate::udfs::apply_to_timestamp_format(
1589 s,
1590 fmt_owned.as_deref(),
1591 true,
1592 use_recent_null,
1593 ))
1594 },
1595 move |_schema, field| {
1596 match field.dtype() {
1598 DataType::String
1599 | DataType::Date
1600 | DataType::Datetime(_, _)
1601 | DataType::Int32
1602 | DataType::Int64
1603 | DataType::Float64
1604 | DataType::Boolean => Ok(Field::new(
1605 out_name2.clone().into(),
1606 DataType::Datetime(TimeUnit::Microseconds, None),
1607 )),
1608 _ => Err(polars::prelude::PolarsError::ComputeError(
1609 "to_timestamp requires StringType, TimestampType, IntegerType, LongType, DateType, or DoubleType"
1610 .into(),
1611 )),
1612 }
1613 },
1614 );
1615 Ok(crate::column::Column::from_expr(
1616 expr.alias(&out_name),
1617 Some(out_name),
1618 ))
1619}
1620
1621pub fn try_to_timestamp(column: &Column, format: Option<&str>) -> Result<Column, String> {
1624 use polars::prelude::*;
1625 let fmt_owned = format.map(|s| s.to_string());
1626 let expr = column.expr().clone().map(
1627 move |s| {
1628 crate::column::expect_col(crate::udfs::apply_to_timestamp_format(
1629 s,
1630 fmt_owned.as_deref(),
1631 false,
1632 false,
1633 ))
1634 },
1635 |_schema, field| {
1636 Ok(Field::new(
1637 field.name().clone(),
1638 DataType::Datetime(TimeUnit::Microseconds, None),
1639 ))
1640 },
1641 );
1642 Ok(crate::column::Column::from_expr(expr, None))
1643}
1644
1645pub fn to_timestamp_ltz(column: &Column, format: Option<&str>) -> Result<Column, String> {
1647 use polars::prelude::{DataType, Field, TimeUnit};
1648 match format {
1649 None => crate::cast(column, "timestamp"),
1650 Some(fmt) => {
1651 let fmt_owned = fmt.to_string();
1652 let expr = column.expr().clone().map(
1653 move |s| {
1654 crate::column::expect_col(crate::udfs::apply_to_timestamp_ltz_format(
1655 s,
1656 Some(&fmt_owned),
1657 true,
1658 ))
1659 },
1660 |_schema, field| {
1661 Ok(Field::new(
1662 field.name().clone(),
1663 DataType::Datetime(TimeUnit::Microseconds, None),
1664 ))
1665 },
1666 );
1667 Ok(crate::column::Column::from_expr(expr, None))
1668 }
1669 }
1670}
1671
1672pub fn to_timestamp_ntz(column: &Column, format: Option<&str>) -> Result<Column, String> {
1674 use polars::prelude::{DataType, Field, TimeUnit};
1675 match format {
1676 None => crate::cast(column, "timestamp"),
1677 Some(fmt) => {
1678 let fmt_owned = fmt.to_string();
1679 let expr = column.expr().clone().map(
1680 move |s| {
1681 crate::column::expect_col(crate::udfs::apply_to_timestamp_ntz_format(
1682 s,
1683 Some(&fmt_owned),
1684 true,
1685 ))
1686 },
1687 |_schema, field| {
1688 Ok(Field::new(
1689 field.name().clone(),
1690 DataType::Datetime(TimeUnit::Microseconds, None),
1691 ))
1692 },
1693 );
1694 Ok(crate::column::Column::from_expr(expr, None))
1695 }
1696 }
1697}
1698
1699pub fn try_divide(left: &Column, right: &Column) -> Column {
1701 use polars::prelude::*;
1702 let zero_cond = right.expr().clone().cast(DataType::Float64).eq(lit(0.0f64));
1703 let null_expr = lit(NULL);
1704 let div_expr =
1705 left.expr().clone().cast(DataType::Float64) / right.expr().clone().cast(DataType::Float64);
1706 let expr = polars::prelude::when(zero_cond)
1707 .then(null_expr)
1708 .otherwise(div_expr);
1709 crate::column::Column::from_expr(expr, None)
1710}
1711
1712pub fn try_add(left: &Column, right: &Column) -> Column {
1714 let args = [right.expr().clone()];
1715 let expr = left.expr().clone().map_many(
1716 |cols| crate::column::expect_col(crate::udfs::apply_try_add(cols)),
1717 &args,
1718 |_schema, fields| Ok(fields[0].clone()),
1719 );
1720 Column::from_expr(expr, None)
1721}
1722
1723pub fn try_subtract(left: &Column, right: &Column) -> Column {
1725 let args = [right.expr().clone()];
1726 let expr = left.expr().clone().map_many(
1727 |cols| crate::column::expect_col(crate::udfs::apply_try_subtract(cols)),
1728 &args,
1729 |_schema, fields| Ok(fields[0].clone()),
1730 );
1731 Column::from_expr(expr, None)
1732}
1733
1734pub fn try_multiply(left: &Column, right: &Column) -> Column {
1736 let args = [right.expr().clone()];
1737 let expr = left.expr().clone().map_many(
1738 |cols| crate::column::expect_col(crate::udfs::apply_try_multiply(cols)),
1739 &args,
1740 |_schema, fields| Ok(fields[0].clone()),
1741 );
1742 Column::from_expr(expr, None)
1743}
1744
1745pub fn try_element_at(column: &Column, index: i64) -> Column {
1747 column.clone().element_at(index)
1748}
1749
1750pub fn width_bucket(value: &Column, min_val: f64, max_val: f64, num_bucket: i64) -> Column {
1752 if num_bucket <= 0 {
1753 panic!(
1754 "width_bucket: num_bucket must be positive, got {}",
1755 num_bucket
1756 );
1757 }
1758 use polars::prelude::*;
1759 let v = value.expr().clone().cast(DataType::Float64);
1760 let min_expr = lit(min_val);
1761 let max_expr = lit(max_val);
1762 let nb = num_bucket as f64;
1763 let width = (max_val - min_val) / nb;
1764 let bucket_expr = (v.clone() - min_expr.clone()) / lit(width);
1765 let floor_bucket = bucket_expr.floor().cast(DataType::Int64) + lit(1i64);
1766 let bucket_clamped = floor_bucket.clip(lit(1i64), lit(num_bucket));
1767 let expr = polars::prelude::when(v.clone().lt(min_expr))
1768 .then(lit(0i64))
1769 .when(v.gt_eq(max_expr))
1770 .then(lit(num_bucket + 1))
1771 .otherwise(bucket_clamped)
1772 .cast(DataType::Int64);
1773 crate::column::Column::from_expr(expr, None)
1774}
1775
1776pub fn elt(index: &Column, columns: &[&Column]) -> Column {
1779 use polars::prelude::*;
1780 if columns.is_empty() {
1781 panic!("elt requires at least one column");
1782 }
1783 let idx_expr = index.expr().clone();
1784 let null_expr = lit(NULL);
1785 let mut expr = null_expr;
1786 for (i, c) in columns.iter().enumerate().rev() {
1787 let n = (i + 1) as i64;
1788 expr = polars::prelude::when(idx_expr.clone().eq(lit(n)))
1789 .then(c.expr().clone())
1790 .otherwise(expr);
1791 }
1792 crate::column::Column::from_expr(expr, None)
1793}
1794
1795pub fn bit_length(column: &Column) -> Column {
1797 column.clone().bit_length()
1798}
1799
1800pub fn octet_length(column: &Column) -> Column {
1802 column.clone().octet_length()
1803}
1804
1805pub fn char_length(column: &Column) -> Column {
1807 column.clone().char_length()
1808}
1809
1810pub fn character_length(column: &Column) -> Column {
1812 column.clone().character_length()
1813}
1814
1815pub fn typeof_(column: &Column) -> Column {
1817 column.clone().typeof_()
1818}
1819
1820pub fn isnan(column: &Column) -> Column {
1822 column.clone().is_nan()
1823}
1824
1825pub fn greatest(columns: &[&Column]) -> Result<Column, String> {
1827 if columns.is_empty() {
1828 return Err("greatest requires at least one column".to_string());
1829 }
1830 if columns.len() == 1 {
1831 return Ok((*columns[0]).clone());
1832 }
1833 let mut expr = columns[0].expr().clone();
1834 for c in columns.iter().skip(1) {
1835 let args = [c.expr().clone()];
1836 expr = expr.map_many(
1837 |cols| crate::column::expect_col(crate::udfs::apply_greatest2(cols)),
1838 &args,
1839 |_schema, fields| Ok(fields[0].clone()),
1840 );
1841 }
1842 Ok(Column::from_expr(expr, None))
1843}
1844
1845pub fn least(columns: &[&Column]) -> Result<Column, String> {
1847 if columns.is_empty() {
1848 return Err("least requires at least one column".to_string());
1849 }
1850 if columns.len() == 1 {
1851 return Ok((*columns[0]).clone());
1852 }
1853 let mut expr = columns[0].expr().clone();
1854 for c in columns.iter().skip(1) {
1855 let args = [c.expr().clone()];
1856 expr = expr.map_many(
1857 |cols| crate::column::expect_col(crate::udfs::apply_least2(cols)),
1858 &args,
1859 |_schema, fields| Ok(fields[0].clone()),
1860 );
1861 }
1862 Ok(Column::from_expr(expr, None))
1863}
1864
1865pub fn year(column: &Column) -> Column {
1867 column.clone().year()
1868}
1869
1870pub fn month(column: &Column) -> Column {
1872 column.clone().month()
1873}
1874
1875pub fn day(column: &Column) -> Column {
1877 column.clone().day()
1878}
1879
1880pub fn to_date(column: &Column, format: Option<&str>) -> Result<Column, String> {
1882 let fmt = format.map(|s| s.to_string());
1883 let out_name = match format {
1884 None => format!("to_date({})", column.name()),
1885 Some(f) => format!("to_date({}, {f})", column.name()),
1886 };
1887 let out_name2 = out_name.clone();
1888 let expr = column.expr().clone().map(
1889 move |col| {
1890 crate::column::expect_col(crate::udfs::apply_string_to_date_format(
1891 col,
1892 fmt.as_deref(),
1893 true, ))
1895 },
1896 move |_schema, _field| Ok(Field::new(out_name2.clone().into(), DataType::Date)),
1897 );
1898 Ok(Column::from_expr(expr.alias(&out_name), Some(out_name)))
1899}
1900
1901pub fn date_format(column: &Column, format: &str) -> Column {
1903 use polars::prelude::*;
1904 let out_name = format!("date_format({}, {format})", column.name());
1905 let chrono_fmt = crate::udfs::pyspark_format_to_chrono(format);
1906 let parsed = column.expr().clone().map(
1907 |s| crate::column::expect_col(crate::udfs::apply_string_to_date_format(s, None, false)),
1908 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Date)),
1909 );
1910 let expr = parsed.dt().strftime(&chrono_fmt).alias(&out_name);
1911 crate::column::Column::from_expr(expr, Some(out_name))
1912}
1913
1914pub fn current_date() -> Column {
1916 use polars::prelude::*;
1917 let today = chrono::Utc::now().date_naive();
1918 let days = (today - robin_sparkless_core::date_utils::epoch_naive_date()).num_days() as i32;
1919 crate::column::Column::from_expr(
1920 Expr::Literal(LiteralValue::Scalar(Scalar::new_date(days))),
1921 None,
1922 )
1923}
1924
1925pub fn current_timestamp() -> Column {
1927 use polars::prelude::*;
1928 let ts = chrono::Utc::now().timestamp_micros();
1929 crate::column::Column::from_expr(
1930 Expr::Literal(LiteralValue::Scalar(Scalar::new_datetime(
1931 ts,
1932 TimeUnit::Microseconds,
1933 None,
1934 ))),
1935 None,
1936 )
1937}
1938
1939pub fn curdate() -> Column {
1941 current_date()
1942}
1943
1944pub fn now() -> Column {
1946 current_timestamp()
1947}
1948
1949pub fn localtimestamp() -> Column {
1951 current_timestamp()
1952}
1953
1954pub fn date_diff(end: &Column, start: &Column) -> Column {
1956 datediff(end, start)
1957}
1958
1959pub fn dateadd(column: &Column, n: i32) -> Column {
1961 date_add(column, n)
1962}
1963
1964pub fn extract(column: &Column, field: &str) -> Column {
1966 column.clone().extract(field)
1967}
1968
1969pub fn date_part(column: &Column, field: &str) -> Column {
1971 extract(column, field)
1972}
1973
1974pub fn datepart(column: &Column, field: &str) -> Column {
1976 extract(column, field)
1977}
1978
1979pub fn unix_micros(column: &Column) -> Column {
1981 column.clone().unix_micros()
1982}
1983
1984pub fn unix_millis(column: &Column) -> Column {
1986 column.clone().unix_millis()
1987}
1988
1989pub fn unix_seconds(column: &Column) -> Column {
1991 column.clone().unix_seconds()
1992}
1993
1994pub fn dayname(column: &Column) -> Column {
1996 column.clone().dayname()
1997}
1998
1999pub fn weekday(column: &Column) -> Column {
2001 column.clone().weekday()
2002}
2003
2004pub fn hour(column: &Column) -> Column {
2006 column.clone().hour()
2007}
2008
2009pub fn minute(column: &Column) -> Column {
2011 column.clone().minute()
2012}
2013
2014pub fn second(column: &Column) -> Column {
2016 column.clone().second()
2017}
2018
2019pub fn date_add(column: &Column, n: i32) -> Column {
2021 column.clone().date_add(n)
2022}
2023
2024pub fn date_sub(column: &Column, n: i32) -> Column {
2026 column.clone().date_sub(n)
2027}
2028
2029pub fn datediff(end: &Column, start: &Column) -> Column {
2031 start.clone().datediff(end)
2032}
2033
2034pub fn last_day(column: &Column) -> Column {
2036 column.clone().last_day()
2037}
2038
2039pub fn trunc(column: &Column, format: &str) -> Column {
2041 column.clone().trunc(format)
2042}
2043
2044pub fn date_trunc(format: &str, column: &Column) -> Column {
2046 trunc(column, format)
2047}
2048
2049pub fn quarter(column: &Column) -> Column {
2051 column.clone().quarter()
2052}
2053
2054pub fn weekofyear(column: &Column) -> Column {
2056 column.clone().weekofyear()
2057}
2058
2059pub fn dayofweek(column: &Column) -> Column {
2061 column.clone().dayofweek()
2062}
2063
2064pub fn dayofyear(column: &Column) -> Column {
2066 column.clone().dayofyear()
2067}
2068
2069pub fn add_months(column: &Column, n: i32) -> Column {
2071 column.clone().add_months(n)
2072}
2073
2074pub fn months_between(end: &Column, start: &Column, round_off: bool) -> Column {
2077 end.clone().months_between(start, round_off)
2078}
2079
2080pub fn next_day(column: &Column, day_of_week: &str) -> Column {
2082 column.clone().next_day(day_of_week)
2083}
2084
2085pub fn unix_timestamp_now() -> Column {
2087 use polars::prelude::*;
2088 let secs = chrono::Utc::now().timestamp();
2089 crate::column::Column::from_expr(lit(secs), None)
2090}
2091
2092pub fn unix_timestamp(column: &Column, format: Option<&str>) -> Column {
2094 column.clone().unix_timestamp(format)
2095}
2096
2097pub fn to_unix_timestamp(column: &Column, format: Option<&str>) -> Column {
2099 unix_timestamp(column, format)
2100}
2101
2102pub fn from_unixtime(column: &Column, format: Option<&str>) -> Column {
2104 column.clone().from_unixtime(format)
2105}
2106
2107pub fn make_date(year: &Column, month: &Column, day: &Column) -> Column {
2109 use polars::prelude::*;
2110 let args = [month.expr().clone(), day.expr().clone()];
2111 let expr = year.expr().clone().map_many(
2112 |cols| crate::column::expect_col(crate::udfs::apply_make_date(cols)),
2113 &args,
2114 |_schema, fields| Ok(Field::new(fields[0].name().clone(), DataType::Date)),
2115 );
2116 crate::column::Column::from_expr(expr, None)
2117}
2118
2119pub fn make_timestamp(
2122 year: &Column,
2123 month: &Column,
2124 day: &Column,
2125 hour: &Column,
2126 minute: &Column,
2127 sec: &Column,
2128 timezone: Option<&str>,
2129) -> Column {
2130 use polars::prelude::*;
2131 let tz_owned = timezone.map(|s| s.to_string());
2132 let args = [
2133 month.expr().clone(),
2134 day.expr().clone(),
2135 hour.expr().clone(),
2136 minute.expr().clone(),
2137 sec.expr().clone(),
2138 ];
2139 let expr = year.expr().clone().map_many(
2140 move |cols| {
2141 crate::column::expect_col(crate::udfs::apply_make_timestamp(cols, tz_owned.as_deref()))
2142 },
2143 &args,
2144 |_schema, fields| {
2145 Ok(Field::new(
2146 fields[0].name().clone(),
2147 DataType::Datetime(TimeUnit::Microseconds, None),
2148 ))
2149 },
2150 );
2151 crate::column::Column::from_expr(expr, None)
2152}
2153
2154pub fn timestampadd(unit: &str, amount: &Column, ts: &Column) -> Column {
2156 ts.clone().timestampadd(unit, amount)
2157}
2158
2159pub fn timestampdiff(unit: &str, start: &Column, end: &Column) -> Column {
2161 start.clone().timestampdiff(unit, end)
2162}
2163
2164pub fn days(n: i64) -> Column {
2166 make_interval(0, 0, 0, n, 0, 0, 0)
2167}
2168
2169pub fn hours(n: i64) -> Column {
2171 make_interval(0, 0, 0, 0, n, 0, 0)
2172}
2173
2174pub fn minutes(n: i64) -> Column {
2176 make_interval(0, 0, 0, 0, 0, n, 0)
2177}
2178
2179pub fn months(n: i64) -> Column {
2181 make_interval(0, n, 0, 0, 0, 0, 0)
2182}
2183
2184pub fn years(n: i64) -> Column {
2186 make_interval(n, 0, 0, 0, 0, 0, 0)
2187}
2188
2189pub fn from_utc_timestamp(column: &Column, tz: &str) -> Column {
2191 column.clone().from_utc_timestamp(tz)
2192}
2193
2194pub fn to_utc_timestamp(column: &Column, tz: &str) -> Column {
2196 column.clone().to_utc_timestamp(tz)
2197}
2198
2199pub fn convert_timezone(source_tz: &str, target_tz: &str, column: &Column) -> Column {
2201 let source_tz = source_tz.to_string();
2202 let target_tz = target_tz.to_string();
2203 let expr = column.expr().clone().map(
2204 move |s| {
2205 crate::column::expect_col(crate::udfs::apply_convert_timezone(
2206 s, &source_tz, &target_tz,
2207 ))
2208 },
2209 |_schema, field| Ok(field.clone()),
2210 );
2211 crate::column::Column::from_expr(expr, None)
2212}
2213
2214pub fn current_timezone() -> Column {
2216 use polars::prelude::*;
2217 crate::column::Column::from_expr(lit("UTC"), None)
2218}
2219
2220pub fn make_interval(
2222 years: i64,
2223 months: i64,
2224 weeks: i64,
2225 days: i64,
2226 hours: i64,
2227 mins: i64,
2228 secs: i64,
2229) -> Column {
2230 use polars::prelude::*;
2231 let total_days = years * 365 + months * 30 + weeks * 7 + days;
2233 let args = DurationArgs::new()
2234 .with_days(lit(total_days))
2235 .with_hours(lit(hours))
2236 .with_minutes(lit(mins))
2237 .with_seconds(lit(secs));
2238 let dur = duration(args);
2239 crate::column::Column::from_expr(dur, None)
2240}
2241
2242pub fn make_dt_interval(days: i64, hours: i64, minutes: i64, seconds: i64) -> Column {
2244 use polars::prelude::*;
2245 let args = DurationArgs::new()
2246 .with_days(lit(days))
2247 .with_hours(lit(hours))
2248 .with_minutes(lit(minutes))
2249 .with_seconds(lit(seconds));
2250 let dur = duration(args);
2251 crate::column::Column::from_expr(dur, None)
2252}
2253
2254pub fn make_ym_interval(years: i32, months: i32) -> Column {
2256 use polars::prelude::*;
2257 let total_months = years * 12 + months;
2258 crate::column::Column::from_expr(lit(total_months), None)
2259}
2260
2261pub fn make_timestamp_ntz(
2263 year: &Column,
2264 month: &Column,
2265 day: &Column,
2266 hour: &Column,
2267 minute: &Column,
2268 sec: &Column,
2269) -> Column {
2270 make_timestamp(year, month, day, hour, minute, sec, None)
2271}
2272
2273pub fn timestamp_seconds(column: &Column) -> Column {
2275 column.clone().timestamp_seconds()
2276}
2277
2278pub fn timestamp_millis(column: &Column) -> Column {
2280 column.clone().timestamp_millis()
2281}
2282
2283pub fn timestamp_micros(column: &Column) -> Column {
2285 column.clone().timestamp_micros()
2286}
2287
2288pub fn unix_date(column: &Column) -> Column {
2290 column.clone().unix_date()
2291}
2292
2293pub fn date_from_unix_date(column: &Column) -> Column {
2295 column.clone().date_from_unix_date()
2296}
2297
2298pub fn pmod(dividend: &Column, divisor: &Column) -> Column {
2300 dividend.clone().pmod(divisor)
2301}
2302
2303pub fn factorial(column: &Column) -> Column {
2305 column.clone().factorial()
2306}
2307
2308pub fn concat(columns: &[&Column]) -> Column {
2311 use polars::prelude::*;
2312 if columns.is_empty() {
2313 panic!("concat requires at least one column");
2314 }
2315 let exprs: Vec<Expr> = columns
2316 .iter()
2317 .map(|c| c.expr().clone().cast(DataType::String))
2318 .collect();
2319 crate::column::Column::from_expr(concat_str(&exprs, "", false), None)
2320}
2321
2322pub fn concat_ws(separator: &str, columns: &[&Column]) -> Column {
2325 use polars::prelude::*;
2326 if columns.is_empty() {
2327 panic!("concat_ws requires at least one column");
2328 }
2329 let exprs: Vec<Expr> = columns
2330 .iter()
2331 .map(|c| c.expr().clone().cast(DataType::String))
2332 .collect();
2333 crate::column::Column::from_expr(concat_str(&exprs, separator, true), None)
2337}
2338
2339pub fn row_number(column: &Column) -> Column {
2349 column.clone().row_number(false)
2350}
2351
2352pub fn rank(column: &Column, descending: bool) -> Column {
2354 column.clone().rank(descending)
2355}
2356
2357pub fn dense_rank(column: &Column, descending: bool) -> Column {
2359 column.clone().dense_rank(descending)
2360}
2361
2362pub fn lag(column: &Column, n: i64) -> Column {
2364 column.clone().lag(n)
2365}
2366
2367pub fn lead(column: &Column, n: i64) -> Column {
2369 column.clone().lead(n)
2370}
2371
2372pub fn first_value(column: &Column) -> Column {
2374 column.clone().first_value()
2375}
2376
2377pub fn last_value(column: &Column) -> Column {
2379 column.clone().last_value()
2380}
2381
2382pub fn percent_rank(column: &Column, partition_by: &[&str], descending: bool) -> Column {
2384 column.clone().percent_rank(partition_by, descending)
2385}
2386
2387pub fn cume_dist(column: &Column, partition_by: &[&str], descending: bool) -> Column {
2389 column.clone().cume_dist(partition_by, descending)
2390}
2391
2392pub fn ntile(column: &Column, n: u32, partition_by: &[&str], descending: bool) -> Column {
2394 column.clone().ntile(n, partition_by, descending)
2395}
2396
2397pub fn nth_value(column: &Column, n: i64, partition_by: &[&str], descending: bool) -> Column {
2399 column.clone().nth_value(n, partition_by, descending)
2400}
2401
2402pub fn coalesce(columns: &[&Column]) -> Column {
2413 use polars::prelude::*;
2414 if columns.is_empty() {
2415 panic!("coalesce requires at least one column");
2416 }
2417 let exprs: Vec<Expr> = columns.iter().map(|c| c.expr().clone()).collect();
2418 let expr = coalesce(&exprs);
2419 crate::column::Column::from_expr(expr, None)
2420}
2421
2422pub fn nvl(column: &Column, value: &Column) -> Column {
2424 coalesce(&[column, value])
2425}
2426
2427pub fn ifnull(column: &Column, value: &Column) -> Column {
2429 nvl(column, value)
2430}
2431
2432pub fn nullif(column: &Column, value: &Column) -> Column {
2434 use polars::prelude::*;
2435 let cond = column.expr().clone().eq(value.expr().clone());
2436 let null_lit = lit(NULL);
2437 let expr = when(cond).then(null_lit).otherwise(column.expr().clone());
2438 crate::column::Column::from_expr(expr, None)
2439}
2440
2441pub fn nanvl(column: &Column, value: &Column) -> Column {
2443 use polars::prelude::*;
2444 let cond = column.expr().clone().is_nan();
2445 let expr = when(cond)
2446 .then(value.expr().clone())
2447 .otherwise(column.expr().clone());
2448 crate::column::Column::from_expr(expr, None)
2449}
2450
2451pub fn nvl2(col1: &Column, col2: &Column, col3: &Column) -> Column {
2453 use polars::prelude::*;
2454 let cond = col1.expr().clone().is_not_null();
2455 let expr = when(cond)
2456 .then(col2.expr().clone())
2457 .otherwise(col3.expr().clone());
2458 crate::column::Column::from_expr(expr, None)
2459}
2460
2461pub fn substr(column: &Column, start: i64, length: Option<i64>) -> Column {
2463 substring(column, start, length)
2464}
2465
2466pub fn power(column: &Column, exp: i64) -> Column {
2468 pow(column, exp)
2469}
2470
2471pub fn ln(column: &Column) -> Column {
2473 log(column)
2474}
2475
2476pub fn ceiling(column: &Column) -> Column {
2478 ceil(column)
2479}
2480
2481pub fn lcase(column: &Column) -> Column {
2483 lower(column)
2484}
2485
2486pub fn ucase(column: &Column) -> Column {
2488 upper(column)
2489}
2490
2491pub fn dayofmonth(column: &Column) -> Column {
2493 column.clone().dayofmonth()
2494}
2495
2496pub fn to_degrees(column: &Column) -> Column {
2498 degrees(column)
2499}
2500
2501pub fn to_radians(column: &Column) -> Column {
2503 radians(column)
2504}
2505
2506pub fn cosh(column: &Column) -> Column {
2508 column.clone().cosh()
2509}
2510pub fn sinh(column: &Column) -> Column {
2512 column.clone().sinh()
2513}
2514pub fn tanh(column: &Column) -> Column {
2516 column.clone().tanh()
2517}
2518pub fn acosh(column: &Column) -> Column {
2520 column.clone().acosh()
2521}
2522pub fn asinh(column: &Column) -> Column {
2524 column.clone().asinh()
2525}
2526pub fn atanh(column: &Column) -> Column {
2528 column.clone().atanh()
2529}
2530pub fn cbrt(column: &Column) -> Column {
2532 column.clone().cbrt()
2533}
2534pub fn expm1(column: &Column) -> Column {
2536 column.clone().expm1()
2537}
2538pub fn log1p(column: &Column) -> Column {
2540 column.clone().log1p()
2541}
2542pub fn log10(column: &Column) -> Column {
2544 column.clone().log10()
2545}
2546pub fn log2(column: &Column) -> Column {
2548 column.clone().log2()
2549}
2550pub fn rint(column: &Column) -> Column {
2552 column.clone().rint()
2553}
2554pub fn hypot(x: &Column, y: &Column) -> Column {
2556 let xx = x.expr().clone() * x.expr().clone();
2557 let yy = y.expr().clone() * y.expr().clone();
2558 crate::column::Column::from_expr((xx + yy).sqrt(), None)
2559}
2560
2561pub fn isnull(column: &Column) -> Column {
2563 column.clone().is_null()
2564}
2565
2566pub fn isnotnull(column: &Column) -> Column {
2568 column.clone().is_not_null()
2569}
2570
2571pub fn array(columns: &[&Column]) -> Result<crate::column::Column, PolarsError> {
2574 use polars::prelude::*;
2575 if columns.is_empty() {
2576 let empty_inner = Series::new("".into(), Vec::<i64>::new());
2579 let list_series = ListChunked::from_iter([Some(empty_inner)])
2580 .with_name("array".into())
2581 .into_series();
2582 let expr = lit(list_series).first();
2583 return Ok(crate::column::Column::from_expr(expr, None));
2584 }
2585 let exprs: Vec<Expr> = columns.iter().map(|c| c.expr().clone()).collect();
2586 let expr = concat_list(exprs)
2587 .map_err(|e| PolarsError::ComputeError(format!("array concat_list: {e}").into()))?;
2588 let mut col = crate::column::Column::from_expr(expr, None);
2589 col.is_array_expr = true;
2590 Ok(col)
2591}
2592
2593pub fn array_size(column: &Column) -> Column {
2595 column.clone().array_size()
2596}
2597
2598pub fn size(column: &Column) -> Column {
2600 column.clone().array_size()
2601}
2602
2603pub fn cardinality(column: &Column) -> Column {
2605 column.clone().cardinality()
2606}
2607
2608pub fn array_contains(column: &Column, value: &Column) -> Column {
2610 column.clone().array_contains(value.expr().clone())
2611}
2612
2613pub fn array_join(column: &Column, separator: &str) -> Column {
2615 column.clone().array_join(separator)
2616}
2617
2618pub fn array_max(column: &Column) -> Column {
2620 column.clone().array_max()
2621}
2622
2623pub fn array_min(column: &Column) -> Column {
2625 column.clone().array_min()
2626}
2627
2628pub fn element_at(column: &Column, index: i64) -> Column {
2630 column.clone().element_at(index)
2631}
2632
2633pub fn array_sort(column: &Column) -> Column {
2635 column.clone().array_sort()
2636}
2637
2638pub fn array_distinct(column: &Column) -> Column {
2640 column.clone().array_distinct()
2641}
2642
2643pub fn array_slice(column: &Column, start: i64, length: Option<i64>) -> Column {
2645 column.clone().array_slice(start, length)
2646}
2647
2648pub fn sequence(start: &Column, stop: &Column, step: Option<&Column>) -> Column {
2651 use polars::prelude::{DataType, Field, as_struct, lit};
2652 let step_expr = step
2653 .map(|c| c.expr().clone().alias("2"))
2654 .unwrap_or_else(|| lit(1i64).alias("2"));
2655 let struct_expr = as_struct(vec![
2656 start.expr().clone().alias("0"),
2657 stop.expr().clone().alias("1"),
2658 step_expr,
2659 ]);
2660 let out_dtype = DataType::List(Box::new(DataType::Int64));
2661 let expr = struct_expr.map(
2662 |s| crate::column::expect_col(crate::udfs::apply_sequence(s)),
2663 move |_schema, field| Ok(Field::new(field.name().clone(), out_dtype.clone())),
2664 );
2665 crate::column::Column::from_expr(expr, None)
2666}
2667
2668pub fn shuffle(column: &Column) -> Column {
2670 let expr = column.expr().clone().map(
2671 |s| crate::column::expect_col(crate::udfs::apply_shuffle(s)),
2672 |_schema, field| Ok(field.clone()),
2673 );
2674 crate::column::Column::from_expr(expr, None)
2675}
2676
2677pub fn inline(column: &Column) -> Column {
2680 column.clone().explode()
2681}
2682
2683pub fn inline_outer(column: &Column) -> Column {
2685 column.clone().explode_outer()
2686}
2687
2688pub fn explode(column: &Column) -> Column {
2690 column.clone().explode()
2691}
2692
2693pub fn array_position(column: &Column, value: &Column) -> Column {
2696 column.clone().array_position(value.expr().clone())
2697}
2698
2699pub fn array_compact(column: &Column) -> Column {
2701 column.clone().array_compact()
2702}
2703
2704pub fn array_remove(column: &Column, value: &Column) -> Column {
2707 column.clone().array_remove(value.expr().clone())
2708}
2709
2710pub fn array_repeat(column: &Column, n: i64) -> Column {
2712 column.clone().array_repeat(n)
2713}
2714
2715pub fn array_flatten(column: &Column) -> Column {
2717 column.clone().array_flatten()
2718}
2719
2720pub fn array_exists(column: &Column, predicate: Expr) -> Column {
2722 column.clone().array_exists(predicate)
2723}
2724
2725pub fn array_forall(column: &Column, predicate: Expr) -> Column {
2727 column.clone().array_forall(predicate)
2728}
2729
2730pub fn array_filter(column: &Column, predicate: Expr) -> Column {
2732 column.clone().array_filter(predicate)
2733}
2734
2735pub fn array_transform(column: &Column, f: Expr) -> Column {
2737 column.clone().array_transform(f)
2738}
2739
2740pub fn array_sum(column: &Column) -> Column {
2742 column.clone().array_sum()
2743}
2744
2745pub fn aggregate(column: &Column, zero: &Column) -> Column {
2747 column.clone().array_aggregate(zero)
2748}
2749
2750pub fn array_mean(column: &Column) -> Column {
2752 column.clone().array_mean()
2753}
2754
2755pub fn posexplode(column: &Column) -> (Column, Column) {
2758 column.clone().posexplode()
2759}
2760
2761pub fn create_map(key_values: &[&Column]) -> Result<Column, PolarsError> {
2765 use polars::chunked_array::StructChunked;
2766 use polars::prelude::{IntoSeries, ListChunked, as_struct, concat_list, lit};
2767 if key_values.is_empty() {
2768 let key_s = Series::new("key".into(), Vec::<String>::new());
2770 let value_s = Series::new("value".into(), Vec::<String>::new());
2771 let fields: [&Series; 2] = [&key_s, &value_s];
2772 let empty_struct = StructChunked::from_series(
2773 polars::prelude::PlSmallStr::EMPTY,
2774 0,
2775 fields.iter().copied(),
2776 )
2777 .map_err(|e| PolarsError::ComputeError(format!("create_map empty struct: {e}").into()))?
2778 .into_series();
2779 let list_series = ListChunked::from_iter([Some(empty_struct)])
2780 .with_name("create_map".into())
2781 .into_series();
2782 let expr = lit(list_series).first();
2783 return Ok(crate::column::Column::from_expr(expr, None));
2784 }
2785 let mut struct_exprs: Vec<Expr> = Vec::new();
2786 for i in (0..key_values.len()).step_by(2) {
2787 if i + 1 < key_values.len() {
2788 let k = key_values[i].expr().clone().alias("key");
2789 let v = key_values[i + 1].expr().clone().alias("value");
2790 struct_exprs.push(as_struct(vec![k, v]));
2791 }
2792 }
2793 let expr = concat_list(struct_exprs)
2794 .map_err(|e| PolarsError::ComputeError(format!("create_map concat_list: {e}").into()))?;
2795 Ok(crate::column::Column::from_expr(expr, None))
2796}
2797
2798pub fn map_keys(column: &Column) -> Column {
2800 column.clone().map_keys()
2801}
2802
2803pub fn map_values(column: &Column) -> Column {
2805 column.clone().map_values()
2806}
2807
2808pub fn map_entries(column: &Column) -> Column {
2810 column.clone().map_entries()
2811}
2812
2813pub fn map_from_arrays(keys: &Column, values: &Column) -> Column {
2815 keys.clone().map_from_arrays(values)
2816}
2817
2818pub fn map_concat(a: &Column, b: &Column) -> Column {
2820 a.clone().map_concat(b)
2821}
2822
2823pub fn map_from_entries(column: &Column) -> Column {
2825 column.clone().map_from_entries()
2826}
2827
2828pub fn map_contains_key(map_col: &Column, key: &Column) -> Column {
2830 map_col.clone().map_contains_key(key)
2831}
2832
2833pub fn get(map_col: &Column, key: &Column) -> Column {
2835 map_col.clone().get(key)
2836}
2837
2838pub fn map_filter(map_col: &Column, predicate: Expr) -> Column {
2840 map_col.clone().map_filter(predicate)
2841}
2842
2843pub fn map_zip_with(map1: &Column, map2: &Column, merge: Expr) -> Column {
2845 map1.clone().map_zip_with(map2, merge)
2846}
2847
2848pub fn zip_with_coalesce(left: &Column, right: &Column) -> Column {
2850 use polars::prelude::col;
2851 let left_field = col("").struct_().field_by_name("left");
2852 let right_field = col("").struct_().field_by_name("right");
2853 let merge = crate::column::Column::from_expr(
2854 coalesce(&[
2855 &crate::column::Column::from_expr(left_field, None),
2856 &crate::column::Column::from_expr(right_field, None),
2857 ])
2858 .into_expr(),
2859 None,
2860 );
2861 left.clone().zip_with(right, merge.into_expr())
2862}
2863
2864pub fn map_zip_with_coalesce(map1: &Column, map2: &Column) -> Column {
2866 use polars::prelude::col;
2867 let v1 = col("").struct_().field_by_name("value1");
2868 let v2 = col("").struct_().field_by_name("value2");
2869 let merge = coalesce(&[
2870 &crate::column::Column::from_expr(v1, None),
2871 &crate::column::Column::from_expr(v2, None),
2872 ])
2873 .into_expr();
2874 map1.clone().map_zip_with(map2, merge)
2875}
2876
2877pub fn map_filter_value_gt(map_col: &Column, threshold: f64) -> Column {
2879 use polars::prelude::{col, lit};
2880 let pred = col("").struct_().field_by_name("value").gt(lit(threshold));
2881 map_col.clone().map_filter(pred)
2882}
2883
2884pub fn struct_(columns: &[&Column]) -> Column {
2887 use polars::prelude::as_struct;
2888 if columns.is_empty() {
2889 panic!("struct requires at least one column");
2890 }
2891 let exprs: Vec<Expr> = columns.iter().map(|c| c.expr().clone()).collect();
2892 crate::column::Column::from_expr(as_struct(exprs), None)
2893}
2894
2895pub fn named_struct(pairs: &[(&str, &Column)]) -> Column {
2898 use polars::prelude::as_struct;
2899 if pairs.is_empty() {
2900 panic!("named_struct requires at least one (name, column) pair");
2901 }
2902 let exprs: Vec<Expr> = pairs
2903 .iter()
2904 .map(|(name, col)| col.expr().clone().alias(*name))
2905 .collect();
2906 crate::column::Column::from_expr(as_struct(exprs), None)
2907}
2908
2909pub fn array_append(array: &Column, elem: &Column) -> Column {
2911 array.clone().array_append(elem)
2912}
2913
2914pub fn array_prepend(array: &Column, elem: &Column) -> Column {
2916 array.clone().array_prepend(elem)
2917}
2918
2919pub fn array_insert(array: &Column, pos: &Column, elem: &Column) -> Column {
2921 array.clone().array_insert(pos, elem)
2922}
2923
2924pub fn array_except(a: &Column, b: &Column) -> Column {
2926 a.clone().array_except(b)
2927}
2928
2929pub fn array_intersect(a: &Column, b: &Column) -> Column {
2931 a.clone().array_intersect(b)
2932}
2933
2934pub fn array_union(a: &Column, b: &Column) -> Column {
2936 a.clone().array_union(b)
2937}
2938
2939pub fn zip_with(left: &Column, right: &Column, merge: Expr) -> Column {
2941 left.clone().zip_with(right, merge)
2942}
2943
2944pub fn get_json_object(column: &Column, path: &str) -> Column {
2946 column.clone().get_json_object(path)
2947}
2948
2949pub fn json_object_keys(column: &Column) -> Column {
2951 column.clone().json_object_keys()
2952}
2953
2954pub fn json_tuple(column: &Column, keys: &[&str]) -> Column {
2956 column.clone().json_tuple(keys)
2957}
2958
2959pub fn from_csv(column: &Column) -> Column {
2961 column.clone().from_csv()
2962}
2963
2964pub fn to_csv(column: &Column) -> Column {
2966 column.clone().to_csv()
2967}
2968
2969pub fn schema_of_csv(_column: &Column) -> Column {
2971 Column::from_expr(
2972 lit("STRUCT<_c0: STRING, _c1: STRING>".to_string()),
2973 Some("schema_of_csv".to_string()),
2974 )
2975}
2976
2977pub fn schema_of_json(_column: &Column) -> Column {
2979 Column::from_expr(
2980 lit("STRUCT<>".to_string()),
2981 Some("schema_of_json".to_string()),
2982 )
2983}
2984
2985pub fn from_json(column: &Column, schema: Option<polars::datatypes::DataType>) -> Column {
2987 column.clone().from_json(schema)
2988}
2989
2990pub fn to_json(column: &Column) -> Column {
2992 column.clone().to_json()
2993}
2994
2995pub fn isin(column: &Column, other: &Column) -> Column {
2997 column.clone().isin(other)
2998}
2999
3000pub fn isin_i64(column: &Column, values: &[i64]) -> Column {
3003 let str_vals: Vec<String> = values.iter().map(|n| n.to_string()).collect();
3004 let s: Series = Series::from_iter(str_vals.iter().map(String::as_str));
3005 Column::from_expr(
3006 column
3007 .expr()
3008 .clone()
3009 .cast(DataType::String)
3010 .is_in(lit(s).implode(), false),
3011 None,
3012 )
3013}
3014
3015pub fn isin_str(column: &Column, values: &[&str]) -> Column {
3017 let s: Series = Series::from_iter(values.iter().copied());
3018 Column::from_expr(column.expr().clone().is_in(lit(s).implode(), false), None)
3019}
3020
3021pub fn url_decode(column: &Column) -> Column {
3023 column.clone().url_decode()
3024}
3025
3026pub fn url_encode(column: &Column) -> Column {
3028 column.clone().url_encode()
3029}
3030
3031pub fn shift_left(column: &Column, n: i32) -> Column {
3033 column.clone().shift_left(n)
3034}
3035
3036pub fn shift_right(column: &Column, n: i32) -> Column {
3038 column.clone().shift_right(n)
3039}
3040
3041pub fn shift_right_unsigned(column: &Column, n: i32) -> Column {
3043 column.clone().shift_right_unsigned(n)
3044}
3045
3046pub fn version() -> Column {
3048 Column::from_expr(
3049 lit(concat!("robin-sparkless-", env!("CARGO_PKG_VERSION"))),
3050 None,
3051 )
3052}
3053
3054pub fn equal_null(left: &Column, right: &Column) -> Column {
3056 left.clone().eq_null_safe(right)
3057}
3058
3059pub fn json_array_length(column: &Column, path: &str) -> Column {
3061 column.clone().json_array_length(path)
3062}
3063
3064pub fn parse_url(column: &Column, part: &str, key: Option<&str>) -> Column {
3067 column.clone().parse_url(part, key)
3068}
3069
3070pub fn hash(columns: &[&Column]) -> Column {
3072 use polars::prelude::*;
3073 if columns.is_empty() {
3074 return crate::column::Column::from_expr(lit(0i64), None);
3075 }
3076 if columns.len() == 1 {
3077 return columns[0].clone().hash();
3078 }
3079 let exprs: Vec<Expr> = columns.iter().map(|c| c.expr().clone()).collect();
3080 let struct_expr = polars::prelude::as_struct(exprs);
3081 let name = columns[0].name().to_string();
3082 let expr = struct_expr.map(
3083 |s| crate::column::expect_col(crate::udfs::apply_hash_struct(s)),
3084 |_schema, field| Ok(Field::new(field.name().clone(), DataType::Int64)),
3085 );
3086 crate::column::Column::from_expr(expr, Some(name))
3087}
3088
3089pub fn stack(columns: &[&Column]) -> Column {
3091 struct_(columns)
3092}
3093
3094#[cfg(test)]
3095mod tests {
3096 use super::*;
3097 use crate::functions::{col, lit_bool, lit_f64, lit_i32, lit_i64, lit_str};
3098 use polars::prelude::{DataType, IntoLazy, Series, df};
3099
3100 #[test]
3101 fn test_col_creates_column() {
3102 let column = col("test");
3103 assert_eq!(column.name(), "test");
3104 }
3105
3106 #[test]
3107 fn test_lit_i32() {
3108 let column = lit_i32(42);
3109 assert_eq!(column.name(), "<expr>");
3111 }
3112
3113 #[test]
3114 fn test_lit_i64() {
3115 let column = lit_i64(123456789012345i64);
3116 assert_eq!(column.name(), "<expr>");
3117 }
3118
3119 #[test]
3120 fn test_lit_f64() {
3121 let column = lit_f64(std::f64::consts::PI);
3122 assert_eq!(column.name(), "<expr>");
3123 }
3124
3125 #[test]
3126 fn test_lit_bool() {
3127 let column = lit_bool(true);
3128 assert_eq!(column.name(), "<expr>");
3129 }
3130
3131 #[test]
3132 fn test_lit_str() {
3133 let column = lit_str("hello");
3134 assert_eq!(column.name(), "<expr>");
3135 }
3136
3137 #[test]
3138 fn cast_null_literal_to_int_and_long() {
3139 let s = Series::new_null("x".into(), 3);
3142
3143 let col_i32 =
3145 crate::udfs::apply_string_to_int(s.clone().into_column(), true, DataType::Int32)
3146 .expect("cast to Int32 should succeed")
3147 .expect("column should be present");
3148 let col_i64 = crate::udfs::apply_string_to_int(s.into_column(), true, DataType::Int64)
3149 .expect("cast to Int64 should succeed")
3150 .expect("column should be present");
3151
3152 assert_eq!(col_i32.dtype(), &DataType::Int32);
3153 assert_eq!(col_i64.dtype(), &DataType::Int64);
3154 assert!(col_i32.is_null().all());
3155 assert!(col_i64.is_null().all());
3156 }
3157
3158 #[test]
3159 fn test_create_map_empty() {
3160 let empty_col = create_map(&[]).unwrap();
3162 let df = df!("id" => &[1i64, 2i64]).unwrap();
3163 let out = df
3164 .lazy()
3165 .with_columns([empty_col.into_expr().alias("m")])
3166 .collect()
3167 .unwrap();
3168 assert_eq!(out.height(), 2);
3169 let m = out.column("m").unwrap();
3170 assert_eq!(m.len(), 2);
3171 let list = m.list().unwrap();
3172 for i in 0..2 {
3173 let row = list.get(i).unwrap();
3174 assert_eq!(row.len(), 0);
3175 }
3176 }
3177
3178 #[test]
3179 fn test_count_aggregation() {
3180 let column = col("value");
3181 let result = count(&column);
3182 assert_eq!(result.name(), "count(value)");
3183 }
3184
3185 #[test]
3186 fn test_sum_aggregation() {
3187 let column = col("value");
3188 let result = sum(&column);
3189 assert_eq!(result.name(), "sum(value)");
3190 }
3191
3192 #[test]
3193 fn test_avg_aggregation() {
3194 let column = col("value");
3195 let result = avg(&column);
3196 assert_eq!(result.name(), "avg(value)");
3197 }
3198
3199 #[test]
3200 fn test_max_aggregation() {
3201 let column = col("value");
3202 let result = max(&column);
3203 assert_eq!(result.name(), "max(value)");
3204 }
3205
3206 #[test]
3207 fn test_min_aggregation() {
3208 let column = col("value");
3209 let result = min(&column);
3210 assert_eq!(result.name(), "min(value)");
3211 }
3212
3213 #[test]
3214 fn test_when_then_otherwise() {
3215 let df = df!(
3217 "age" => &[15, 25, 35]
3218 )
3219 .unwrap();
3220
3221 let age_col = col("age");
3223 let condition = age_col.gt(polars::prelude::lit(18));
3224 let result = when(&condition)
3225 .then(&lit_str("adult"))
3226 .otherwise(&lit_str("minor"));
3227
3228 let result_df = df
3230 .lazy()
3231 .with_column(result.into_expr().alias("status"))
3232 .collect()
3233 .unwrap();
3234
3235 let status_col = result_df.column("status").unwrap();
3237 let values: Vec<Option<&str>> = status_col.str().unwrap().into_iter().collect();
3238
3239 assert_eq!(values[0], Some("minor")); assert_eq!(values[1], Some("adult")); assert_eq!(values[2], Some("adult")); }
3243
3244 #[test]
3246 fn test_when_then_otherwise_numeric_then_else() {
3247 let df = df!("a" => &[1i64, 3, 5]).unwrap();
3248 let condition = col("a").gt(polars::prelude::lit(2));
3249 let result = when(&condition).then(&lit_i64(100)).otherwise(&lit_i64(0));
3250 let result_df = df
3251 .lazy()
3252 .with_column(result.into_expr().alias("out"))
3253 .collect()
3254 .unwrap();
3255 let out = result_df.column("out").unwrap();
3256 let out_i64 = out.cast(&DataType::Int64).unwrap();
3258 let values: Vec<Option<i64>> = out_i64.i64().unwrap().into_iter().collect();
3259 assert_eq!(values, vec![Some(0), Some(100), Some(100)]);
3260 }
3261
3262 #[test]
3263 fn test_coalesce_returns_first_non_null() {
3264 let df = df!(
3266 "a" => &[Some(1), None, None],
3267 "b" => &[None, Some(2), None],
3268 "c" => &[None, None, Some(3)]
3269 )
3270 .unwrap();
3271
3272 let col_a = col("a");
3273 let col_b = col("b");
3274 let col_c = col("c");
3275 let result = coalesce(&[&col_a, &col_b, &col_c]);
3276
3277 let result_df = df
3279 .lazy()
3280 .with_column(result.into_expr().alias("coalesced"))
3281 .collect()
3282 .unwrap();
3283
3284 let coalesced_col = result_df.column("coalesced").unwrap();
3286 let values: Vec<Option<i32>> = coalesced_col.i32().unwrap().into_iter().collect();
3287
3288 assert_eq!(values[0], Some(1)); assert_eq!(values[1], Some(2)); assert_eq!(values[2], Some(3)); }
3292
3293 #[test]
3294 fn test_coalesce_with_literal_fallback() {
3295 let df = df!(
3297 "a" => &[Some(1), None],
3298 "b" => &[None::<i32>, None::<i32>]
3299 )
3300 .unwrap();
3301
3302 let col_a = col("a");
3303 let col_b = col("b");
3304 let fallback = lit_i32(0);
3305 let result = coalesce(&[&col_a, &col_b, &fallback]);
3306
3307 let result_df = df
3309 .lazy()
3310 .with_column(result.into_expr().alias("coalesced"))
3311 .collect()
3312 .unwrap();
3313
3314 let coalesced_col = result_df.column("coalesced").unwrap();
3316 let values: Vec<Option<i32>> = coalesced_col.i32().unwrap().into_iter().collect();
3317
3318 assert_eq!(values[0], Some(1)); assert_eq!(values[1], Some(0)); }
3321
3322 #[test]
3323 #[should_panic(expected = "coalesce requires at least one column")]
3324 fn test_coalesce_empty_panics() {
3325 let columns: [&Column; 0] = [];
3326 let _ = coalesce(&columns);
3327 }
3328
3329 #[test]
3330 fn test_cast_double_string_column_strict_ok() {
3331 let df = df!(
3333 "s" => &["123", " 45.5 ", "0"]
3334 )
3335 .unwrap();
3336
3337 let s_col = col("s");
3338 let cast_col = cast(&s_col, "double").unwrap();
3339
3340 let out = df
3341 .lazy()
3342 .with_column(cast_col.into_expr().alias("v"))
3343 .collect()
3344 .unwrap();
3345
3346 let v = out.column("v").unwrap();
3347 let vals: Vec<Option<f64>> = v.f64().unwrap().into_iter().collect();
3348 assert_eq!(vals, vec![Some(123.0), Some(45.5), Some(0.0)]);
3349 }
3350
3351 #[test]
3352 fn test_try_cast_double_string_column_invalid_to_null() {
3353 let df = df!(
3355 "s" => &["123", " 45.5 ", "abc", ""]
3356 )
3357 .unwrap();
3358
3359 let s_col = col("s");
3360 let try_cast_col = try_cast(&s_col, "double").unwrap();
3361
3362 let out = df
3363 .lazy()
3364 .with_column(try_cast_col.into_expr().alias("v"))
3365 .collect()
3366 .unwrap();
3367
3368 let v = out.column("v").unwrap();
3369 let vals: Vec<Option<f64>> = v.f64().unwrap().into_iter().collect();
3370 assert_eq!(vals, vec![Some(123.0), Some(45.5), None, None]);
3371 }
3372
3373 #[test]
3375 fn test_cast_double_to_int_nan_to_null() {
3376 use polars::prelude::df;
3377
3378 let df = df!(
3379 "x" => &[1.7, 2.3, f64::NAN, f64::INFINITY]
3380 )
3381 .unwrap();
3382
3383 let col_x = col("x");
3384 let try_cast_col = try_cast(&col_x, "int").unwrap();
3385
3386 let out = df
3387 .lazy()
3388 .with_column(try_cast_col.into_expr().alias("v"))
3389 .collect()
3390 .unwrap();
3391
3392 let v = out.column("v").unwrap();
3393 let vals: Vec<Option<i32>> = v.i32().unwrap().into_iter().collect();
3394 assert_eq!(vals, vec![Some(1), Some(2), None, None]);
3395 }
3396
3397 #[test]
3398 fn test_to_number_and_try_to_number_numerics_and_strings() {
3399 let df = df!(
3401 "i" => &[1i32, 2, 3],
3402 "f" => &[1.5f64, 2.5, 3.5],
3403 "s" => &["10", "20.5", "xyz"]
3404 )
3405 .unwrap();
3406
3407 let i_col = col("i");
3408 let f_col = col("f");
3409 let s_col = col("s");
3410
3411 let to_number_i = to_number(&i_col, None).unwrap();
3412 let to_number_f = to_number(&f_col, None).unwrap();
3413 let try_to_number_s = try_to_number(&s_col, None).unwrap();
3414
3415 let out = df
3416 .lazy()
3417 .with_columns([
3418 to_number_i.into_expr().alias("i_num"),
3419 to_number_f.into_expr().alias("f_num"),
3420 try_to_number_s.into_expr().alias("s_num"),
3421 ])
3422 .collect()
3423 .unwrap();
3424
3425 let i_num = out.column("i_num").unwrap();
3426 let f_num = out.column("f_num").unwrap();
3427 let s_num = out.column("s_num").unwrap();
3428
3429 let i_vals: Vec<Option<f64>> = i_num.f64().unwrap().into_iter().collect();
3430 let f_vals: Vec<Option<f64>> = f_num.f64().unwrap().into_iter().collect();
3431 let s_vals: Vec<Option<f64>> = s_num.f64().unwrap().into_iter().collect();
3432
3433 assert_eq!(i_vals, vec![Some(1.0), Some(2.0), Some(3.0)]);
3434 assert_eq!(f_vals, vec![Some(1.5), Some(2.5), Some(3.5)]);
3435 assert_eq!(s_vals, vec![Some(10.0), Some(20.5), None]);
3436 }
3437}