1use polars::lazy::dsl;
2use polars::prelude::*;
3use polars_plan::plans::DynLiteralValue;
4use polars_plan::prelude::UnionArgs;
5use polars_utils::python_function::PythonObject;
6use pyo3::exceptions::{PyTypeError, PyValueError};
7use pyo3::prelude::*;
8use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyString};
9
10use crate::conversion::any_value::py_object_to_any_value;
11use crate::conversion::{Wrap, get_lf};
12use crate::error::PyPolarsErr;
13use crate::expr::ToExprs;
14use crate::expr::datatype::PyDataTypeExpr;
15use crate::lazyframe::PyOptFlags;
16use crate::utils::EnterPolarsExt;
17use crate::{PyDataFrame, PyExpr, PyLazyFrame, PySeries, map};
18
19macro_rules! set_unwrapped_or_0 {
20 ($($var:ident),+ $(,)?) => {
21 $(let $var = $var.map(|e| e.inner).unwrap_or(dsl::lit(0));)+
22 };
23}
24
25#[pyfunction]
26pub fn rolling_corr(
27 x: PyExpr,
28 y: PyExpr,
29 window_size: IdxSize,
30 min_periods: IdxSize,
31 ddof: u8,
32) -> PyExpr {
33 dsl::rolling_corr(
34 x.inner,
35 y.inner,
36 RollingCovOptions {
37 min_periods,
38 window_size,
39 ddof,
40 },
41 )
42 .into()
43}
44
45#[pyfunction]
46pub fn rolling_cov(
47 x: PyExpr,
48 y: PyExpr,
49 window_size: IdxSize,
50 min_periods: IdxSize,
51 ddof: u8,
52) -> PyExpr {
53 dsl::rolling_cov(
54 x.inner,
55 y.inner,
56 RollingCovOptions {
57 min_periods,
58 window_size,
59 ddof,
60 },
61 )
62 .into()
63}
64
65#[pyfunction]
66pub fn arg_sort_by(
67 by: Vec<PyExpr>,
68 descending: Vec<bool>,
69 nulls_last: Vec<bool>,
70 multithreaded: bool,
71 maintain_order: bool,
72) -> PyExpr {
73 let by = by.into_iter().map(|e| e.inner).collect::<Vec<Expr>>();
74 dsl::arg_sort_by(
75 by,
76 SortMultipleOptions {
77 descending,
78 nulls_last,
79 multithreaded,
80 maintain_order,
81 limit: None,
82 },
83 )
84 .into()
85}
86#[pyfunction]
87pub fn arg_where(condition: PyExpr) -> PyExpr {
88 dsl::arg_where(condition.inner).into()
89}
90
91#[pyfunction]
92pub fn as_struct(exprs: Vec<PyExpr>) -> PyResult<PyExpr> {
93 let exprs = exprs.to_exprs();
94 if exprs.is_empty() {
95 return Err(PyValueError::new_err(
96 "expected at least 1 expression in 'as_struct'",
97 ));
98 }
99 Ok(dsl::as_struct(exprs).into())
100}
101
102#[pyfunction]
103pub fn field(names: Vec<String>) -> PyExpr {
104 dsl::Expr::Field(names.into_iter().map(|x| x.into()).collect()).into()
105}
106
107#[pyfunction]
108pub fn coalesce(exprs: Vec<PyExpr>) -> PyExpr {
109 let exprs = exprs.to_exprs();
110 dsl::coalesce(&exprs).into()
111}
112
113#[pyfunction]
114pub fn col(name: &str) -> PyExpr {
115 dsl::col(name).into()
116}
117
118#[pyfunction]
119pub fn element() -> PyExpr {
120 dsl::element().into()
121}
122
123fn lfs_to_plans(lfs: Vec<PyLazyFrame>) -> Vec<DslPlan> {
124 lfs.into_iter()
125 .map(|lf| lf.ldf.into_inner().logical_plan)
126 .collect()
127}
128
129#[pyfunction]
130pub fn collect_all(
131 lfs: Vec<PyLazyFrame>,
132 engine: Wrap<Engine>,
133 optflags: PyOptFlags,
134 py: Python<'_>,
135) -> PyResult<Vec<PyDataFrame>> {
136 let plans = lfs_to_plans(lfs);
137 let dfs = py.enter_polars(|| {
138 LazyFrame::collect_all_with_engine(plans, engine.0, optflags.inner.into_inner())
139 })?;
140 Ok(dfs.into_iter().map(Into::into).collect())
141}
142
143#[pyfunction]
144pub fn collect_all_lazy(lfs: Vec<PyLazyFrame>, optflags: PyOptFlags) -> PyResult<PyLazyFrame> {
145 let plans = lfs_to_plans(lfs);
146
147 for plan in &plans {
148 if !matches!(plan, DslPlan::Sink { .. }) {
149 return Err(PyValueError::new_err(
150 "all LazyFrames must end with a sink to use 'collect_all(lazy=True)'",
151 ));
152 }
153 }
154
155 Ok(LazyFrame::from_logical_plan(
156 DslPlan::SinkMultiple { inputs: plans },
157 optflags.inner.into_inner(),
158 )
159 .into())
160}
161
162#[pyfunction]
163pub fn explain_all(lfs: Vec<PyLazyFrame>, optflags: PyOptFlags, py: Python) -> PyResult<String> {
164 let plans = lfs_to_plans(lfs);
165 let explained =
166 py.enter_polars(|| LazyFrame::explain_all(plans, optflags.inner.into_inner()))?;
167 Ok(explained)
168}
169
170#[pyfunction]
171pub fn collect_all_with_callback(
172 lfs: Vec<PyLazyFrame>,
173 engine: Wrap<Engine>,
174 optflags: PyOptFlags,
175 lambda: Py<PyAny>,
176 py: Python<'_>,
177) {
178 let plans = lfs
179 .into_iter()
180 .map(|lf| lf.ldf.into_inner().logical_plan)
181 .collect();
182 let result = py
183 .enter_polars(|| {
184 LazyFrame::collect_all_with_engine(plans, engine.0, optflags.inner.into_inner())
185 })
186 .map(|dfs| {
187 dfs.into_iter()
188 .map(Into::into)
189 .collect::<Vec<PyDataFrame>>()
190 });
191
192 Python::attach(|py| match result {
193 Ok(dfs) => {
194 lambda.call1(py, (dfs,)).map_err(|err| err.restore(py)).ok();
195 },
196 Err(err) => {
197 lambda
198 .call1(py, (PyErr::from(err),))
199 .map_err(|err| err.restore(py))
200 .ok();
201 },
202 })
203}
204
205#[pyfunction]
206pub fn concat_lf(
207 seq: &Bound<'_, PyAny>,
208 rechunk: bool,
209 parallel: bool,
210 to_supertypes: bool,
211 maintain_order: bool,
212) -> PyResult<PyLazyFrame> {
213 let len = seq.len()?;
214 let mut lfs = Vec::with_capacity(len);
215
216 for res in seq.try_iter()? {
217 let item = res?;
218 let lf = get_lf(&item)?;
219 lfs.push(lf);
220 }
221
222 let lf = dsl::concat(
223 lfs,
224 UnionArgs {
225 rechunk,
226 parallel,
227 to_supertypes,
228 maintain_order,
229 ..Default::default()
230 },
231 )
232 .map_err(PyPolarsErr::from)?;
233 Ok(lf.into())
234}
235
236#[pyfunction]
237pub fn concat_list(s: Vec<PyExpr>) -> PyResult<PyExpr> {
238 let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
239 let expr = dsl::concat_list(s).map_err(PyPolarsErr::from)?;
240 Ok(expr.into())
241}
242
243#[pyfunction(name = "list")]
244pub fn as_list(s: Vec<PyExpr>) -> PyResult<PyExpr> {
245 let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
246 let expr = dsl::as_list(s).map_err(PyPolarsErr::from)?;
247 Ok(expr.into())
248}
249
250#[pyfunction]
251pub fn concat_arr(s: Vec<PyExpr>) -> PyResult<PyExpr> {
252 let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
253 let expr = dsl::concat_arr(s).map_err(PyPolarsErr::from)?;
254 Ok(expr.into())
255}
256
257#[pyfunction]
258pub fn concat_str(s: Vec<PyExpr>, separator: &str, ignore_nulls: bool) -> PyExpr {
259 let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
260 dsl::concat_str(s, separator, ignore_nulls).into()
261}
262
263#[pyfunction]
264pub fn len() -> PyExpr {
265 dsl::len().into()
266}
267
268#[pyfunction]
269pub fn cov(a: PyExpr, b: PyExpr, ddof: u8) -> PyExpr {
270 dsl::cov(a.inner, b.inner, ddof).into()
271}
272
273#[pyfunction]
274#[cfg(feature = "trigonometry")]
275pub fn arctan2(y: PyExpr, x: PyExpr) -> PyExpr {
276 y.inner.arctan2(x.inner).into()
277}
278
279#[pyfunction]
280pub fn cum_fold(
281 acc: PyExpr,
282 lambda: Py<PyAny>,
283 exprs: Vec<PyExpr>,
284 returns_scalar: bool,
285 return_dtype: Option<PyDataTypeExpr>,
286 include_init: bool,
287) -> PyExpr {
288 let exprs = exprs.to_exprs();
289 let func = PlanCallback::new_python(PythonObject(lambda));
290 dsl::cum_fold_exprs(
291 acc.inner,
292 func,
293 exprs,
294 returns_scalar,
295 return_dtype.map(|v| v.inner),
296 include_init,
297 )
298 .into()
299}
300
301#[pyfunction]
302pub fn cum_reduce(
303 lambda: Py<PyAny>,
304 exprs: Vec<PyExpr>,
305 returns_scalar: bool,
306 return_dtype: Option<PyDataTypeExpr>,
307) -> PyExpr {
308 let exprs = exprs.to_exprs();
309
310 let func = PlanCallback::new_python(PythonObject(lambda));
311 dsl::cum_reduce_exprs(func, exprs, returns_scalar, return_dtype.map(|v| v.inner)).into()
312}
313
314#[pyfunction]
315#[pyo3(signature = (year, month, day, hour=None, minute=None, second=None, microsecond=None, time_unit=Wrap(TimeUnit::Microseconds), time_zone=Wrap(None), ambiguous=PyExpr::from(dsl::lit(String::from("raise")))))]
316pub fn datetime(
317 year: PyExpr,
318 month: PyExpr,
319 day: PyExpr,
320 hour: Option<PyExpr>,
321 minute: Option<PyExpr>,
322 second: Option<PyExpr>,
323 microsecond: Option<PyExpr>,
324 time_unit: Wrap<TimeUnit>,
325 time_zone: Wrap<Option<TimeZone>>,
326 ambiguous: PyExpr,
327) -> PyExpr {
328 let year = year.inner;
329 let month = month.inner;
330 let day = day.inner;
331 set_unwrapped_or_0!(hour, minute, second, microsecond);
332 let ambiguous = ambiguous.inner;
333 let time_unit = time_unit.0;
334 let time_zone = time_zone.0;
335 let args = DatetimeArgs {
336 year,
337 month,
338 day,
339 hour,
340 minute,
341 second,
342 microsecond,
343 time_unit,
344 time_zone,
345 ambiguous,
346 };
347 dsl::datetime(args).into()
348}
349
350#[pyfunction]
351pub fn concat_lf_diagonal(
352 lfs: &Bound<'_, PyAny>,
353 rechunk: bool,
354 parallel: bool,
355 to_supertypes: bool,
356 maintain_order: bool,
357) -> PyResult<PyLazyFrame> {
358 let iter = lfs.try_iter()?;
359
360 let lfs = iter
361 .map(|item| {
362 let item = item?;
363 get_lf(&item)
364 })
365 .collect::<PyResult<Vec<_>>>()?;
366
367 let lf = dsl::functions::concat_lf_diagonal(
368 lfs,
369 UnionArgs {
370 rechunk,
371 parallel,
372 to_supertypes,
373 maintain_order,
374 ..Default::default()
375 },
376 )
377 .map_err(PyPolarsErr::from)?;
378 Ok(lf.into())
379}
380
381#[pyfunction]
382pub fn concat_lf_horizontal(
383 lfs: &Bound<'_, PyAny>,
384 parallel: bool,
385 strict: bool,
386) -> PyResult<PyLazyFrame> {
387 let iter = lfs.try_iter()?;
388
389 let lfs = iter
390 .map(|item| {
391 let item = item?;
392 get_lf(&item)
393 })
394 .collect::<PyResult<Vec<_>>>()?;
395
396 let lf = dsl::functions::concat_lf_horizontal(
397 lfs,
398 HConcatOptions {
399 parallel,
400 strict,
401 broadcast_unit_length: Default::default(),
402 },
403 )
404 .map_err(PyPolarsErr::from)?;
405 Ok(lf.into())
406}
407
408#[pyfunction]
409pub fn concat_expr(e: Vec<PyExpr>, rechunk: bool) -> PyResult<PyExpr> {
410 let e = e.to_exprs();
411 let e = dsl::functions::concat_expr(e, rechunk).map_err(PyPolarsErr::from)?;
412 Ok(e.into())
413}
414
415#[pyfunction]
416#[pyo3(signature = (weeks, days, hours, minutes, seconds, milliseconds, microseconds, nanoseconds, time_unit))]
417pub fn duration(
418 weeks: Option<PyExpr>,
419 days: Option<PyExpr>,
420 hours: Option<PyExpr>,
421 minutes: Option<PyExpr>,
422 seconds: Option<PyExpr>,
423 milliseconds: Option<PyExpr>,
424 microseconds: Option<PyExpr>,
425 nanoseconds: Option<PyExpr>,
426 time_unit: Wrap<TimeUnit>,
427) -> PyExpr {
428 set_unwrapped_or_0!(
429 weeks,
430 days,
431 hours,
432 minutes,
433 seconds,
434 milliseconds,
435 microseconds,
436 nanoseconds,
437 );
438 let args = DurationArgs {
439 weeks,
440 days,
441 hours,
442 minutes,
443 seconds,
444 milliseconds,
445 microseconds,
446 nanoseconds,
447 time_unit: time_unit.0,
448 };
449 dsl::duration(args).into()
450}
451
452#[pyfunction]
453pub fn fold(
454 acc: PyExpr,
455 lambda: Py<PyAny>,
456 exprs: Vec<PyExpr>,
457 returns_scalar: bool,
458 return_dtype: Option<PyDataTypeExpr>,
459) -> PyExpr {
460 let exprs = exprs.to_exprs();
461 let func = PlanCallback::new_python(PythonObject(lambda));
462 dsl::fold_exprs(
463 acc.inner,
464 func,
465 exprs,
466 returns_scalar,
467 return_dtype.map(|w| w.inner),
468 )
469 .into()
470}
471
472#[pyfunction]
473pub fn lit(value: &Bound<'_, PyAny>, allow_object: bool, is_scalar: bool) -> PyResult<PyExpr> {
474 let py = value.py();
475 if value.is_instance_of::<PyBool>() {
476 let val = value.extract::<bool>()?;
477 Ok(dsl::lit(val).into())
478 } else if let Ok(int) = value.cast::<PyInt>() {
479 let v = int
480 .extract::<i128>()
481 .map_err(|e| polars_err!(InvalidOperation: "integer too large for Polars: {e}"))
482 .map_err(PyPolarsErr::from)?;
483 Ok(Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Int(v))).into())
484 } else if let Ok(float) = value.cast::<PyFloat>() {
485 let val = float.extract::<f64>()?;
486 Ok(Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Float(val))).into())
487 } else if let Ok(pystr) = value.cast::<PyString>() {
488 Ok(dsl::lit(pystr.to_string()).into())
489 } else if let Ok(series) = value.extract::<PySeries>() {
490 let s = series.series.into_inner();
491 if is_scalar {
492 let av = s
493 .get(0)
494 .map_err(|_| PyValueError::new_err("expected at least 1 value"))?;
495 let av = av.into_static();
496 Ok(dsl::lit(Scalar::new(s.dtype().clone(), av)).into())
497 } else {
498 Ok(dsl::lit(s).into())
499 }
500 } else if value.is_none() {
501 Ok(dsl::lit(Null {}).into())
502 } else if let Ok(value) = value.cast::<PyBytes>() {
503 Ok(dsl::lit(value.as_bytes()).into())
504 } else {
505 let raise = || {
506 PyTypeError::new_err(format!(
507 "cannot create expression literal for value of type {}.\
508 \n\nHint: Pass `allow_object=True` to accept any value and create a literal of type Object.",
509 value
510 .get_type()
511 .qualname()
512 .map(|s| s.to_string())
513 .unwrap_or("unknown".to_owned()),
514 ))
515 };
516
517 let av = py_object_to_any_value(value, true, allow_object).map_err(|_| raise())?;
518 match av {
519 #[cfg(feature = "object")]
520 AnyValue::ObjectOwned(_) => {
521 if allow_object {
523 let s = PySeries::new_object(py, "", vec![value.extract()?], false)
524 .series
525 .into_inner();
526 Ok(dsl::lit(s).into())
527 } else {
528 Err(raise())
529 }
530 },
531 _ => Ok(Expr::Literal(LiteralValue::from(av)).into()),
532 }
533 }
534}
535
536#[pyfunction]
537#[pyo3(signature = (pyexpr, lambda, output_type, is_elementwise, returns_scalar))]
538pub fn map_expr(
539 pyexpr: Vec<PyExpr>,
540 lambda: Py<PyAny>,
541 output_type: Option<PyDataTypeExpr>,
542 is_elementwise: bool,
543 returns_scalar: bool,
544) -> PyExpr {
545 map::lazy::map_expr(&pyexpr, lambda, output_type, is_elementwise, returns_scalar)
546}
547
548#[pyfunction]
549pub fn pearson_corr(a: PyExpr, b: PyExpr) -> PyExpr {
550 dsl::pearson_corr(a.inner, b.inner).into()
551}
552
553#[pyfunction]
554pub fn reduce(
555 lambda: Py<PyAny>,
556 exprs: Vec<PyExpr>,
557 returns_scalar: bool,
558 return_dtype: Option<PyDataTypeExpr>,
559) -> PyExpr {
560 let exprs = exprs.to_exprs();
561 let func = PlanCallback::new_python(PythonObject(lambda));
562 dsl::reduce_exprs(func, exprs, returns_scalar, return_dtype.map(|v| v.inner)).into()
563}
564
565#[pyfunction]
566#[pyo3(signature = (value, n, dtype=None))]
567pub fn repeat(value: PyExpr, n: PyExpr, dtype: Option<Wrap<DataType>>) -> PyExpr {
568 let mut value = value.inner;
569 let n = n.inner;
570
571 if let Some(dtype) = dtype {
572 value = value.cast(dtype.0);
573 }
574
575 dsl::repeat(value, n).into()
576}
577
578#[pyfunction]
579pub fn spearman_rank_corr(a: PyExpr, b: PyExpr, propagate_nans: bool) -> PyExpr {
580 #[cfg(feature = "propagate_nans")]
581 {
582 dsl::spearman_rank_corr(a.inner, b.inner, propagate_nans).into()
583 }
584 #[cfg(not(feature = "propagate_nans"))]
585 {
586 panic!("activate 'propagate_nans'")
587 }
588}
589
590#[pyfunction]
591#[cfg(feature = "sql")]
592pub fn sql_expr(sql: &str) -> PyResult<PyExpr> {
593 let expr = polars::sql::sql_expr(sql).map_err(PyPolarsErr::from)?;
594 Ok(expr.into())
595}