add apply to cookbooks (#3504)

polars/polars-lazy/src/dsl/functions.rs

@@ -205,7 +205,9 @@ pub fn concat_str(s: Vec<Expr>, sep: &str) -> Expr {
 /// Concat lists entries.
 #[cfg(feature = "list")]
 #[cfg_attr(docsrs, doc(cfg(feature = "list")))]
-pub fn concat_lst(s: Vec<Expr>) -> Expr {
+pub fn concat_lst<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(s: E) -> Expr {
+    let s = s.as_ref().iter().map(|e| e.clone().into()).collect();
+
     let function = NoEq::new(Arc::new(move |s: &mut [Series]| {
         let mut first = std::mem::take(&mut s[0]);
         let other = &s[1..];

polars/src/docs/eager.rs

@@ -255,6 +255,37 @@
 //! # }
 //! ```
 //!
+//!
+//! ### Multiple columns
+//!
+//! ```
+//! use polars::prelude::*;
+//! fn my_black_box_function(a: f32, b: f32) -> f32 {
+//!     // do something
+//!     a
+//! }
+//!
+//! fn apply_multiples(col_a: &Series, col_b: &Series) -> Float32Chunked {
+//!     match (col_a.dtype(), col_b.dtype()) {
+//!         (DataType::Float32, DataType::Float32) => {
+//!             // downcast to `ChunkedArray`
+//!             let a = col_a.f32().unwrap();
+//!             let b = col_b.f32().unwrap();
+//!
+//!             a.into_iter()
+//!                 .zip(b.into_iter())
+//!                 .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
+//!                     (Some(a), Some(b)) => Some(my_black_box_function(a, b)),
+//!                     // if any of the two value is `None` we propagate that null
+//!                     _ => None,
+//!                 })
+//!                 .collect()
+//!         }
+//!         _ => panic!("unpexptected dtypes"),
+//!     }
+//! }
+//! ```
+//!
 //! ### DataFrame
 //!
 //! ```

polars/src/docs/lazy.rs

@@ -12,6 +12,7 @@
 //! * [GroupBy](#groupby)
 //! * [Joins](#joins)
 //! * [Conditionally apply](#conditionally-apply)
+//! * [Black box function](#black-box-function)
 //!
 //! ## Start a lazy computation
 //!
@@ -249,3 +250,51 @@
 //! # Ok(())
 //! # }
 //! ```
+//!
+//! # Black box function
+//!
+//! The expression API should be expressive enough for most of what you want to achieve, but it can happen
+//! that you need to pass the values to an external function you do not control. The snippet below
+//! shows how we use the `Struct` datatype to be able to apply a function over multiple inputs.
+//!
+//! ```
+//! use polars::prelude::*;
+//! fn my_black_box_function(a: f32, b: f32) -> f32 {
+//!     // do something
+//!     a
+//! }
+//!
+//! fn apply_multiples(lf: LazyFrame) -> Result<DataFrame> {
+//!     df![
+//!         "a" => [1.0, 2.0, 3.0],
+//!         "b" => [3.0, 5.1, 0.3]
+//!     ]?
+//!     .lazy()
+//!     .select([concat_lst(["col_a", "col_b"]).map(
+//!         |s| {
+//!             let ca = s.struct_()?;
+//!
+//!             let b = ca.field_by_name("col_a")?;
+//!             let a = ca.field_by_name("col_b")?;
+//!             let a = a.f32()?;
+//!             let b = b.f32()?;
+//!
+//!             let out: Float32Chunked = a
+//!                 .into_iter()
+//!                 .zip(b.into_iter())
+//!                 .map(|(opt_a, opt_b)| match (opt_a, opt_b) {
+//!                     (Some(a), Some(b)) => Some(my_black_box_function(a, b)),
+//!                     _ => None,
+//!                 })
+//!                 .collect();
+//!
+//!             Ok(out.into_series())
+//!         },
+//!         GetOutput::from_type(DataType::Float32),
+//!     )])
+//!     .collect()
+//! }
+//!
+//! ```
+//!
+//!

py-polars/src/lib.rs

@@ -215,7 +215,7 @@ fn concat_str(s: Vec<dsl::PyExpr>, sep: &str) -> dsl::PyExpr {
 
 #[pyfunction]
 fn concat_lst(s: Vec<dsl::PyExpr>) -> dsl::PyExpr {
-    let s = s.into_iter().map(|e| e.inner).collect();
+    let s = s.into_iter().map(|e| e.inner).collect::<Vec<_>>();
     polars::lazy::dsl::concat_lst(s).into()
 }