feat: Allow step parameter in int_ranges to take an expression

crates/polars-plan/src/dsl/function_expr/range/date_range.rs

@@ -5,7 +5,7 @@ use polars_time::{datetime_range_impl, ClosedWindow, Duration};
 
 use super::datetime_range::{datetime_range, datetime_ranges};
 use super::utils::{
-    ensure_range_bounds_contain_exactly_one_value, ranges_impl_broadcast,
+    ensure_range_bounds_contain_exactly_one_value, temporal_ranges_impl_broadcast,
     temporal_series_to_i64_scalar,
 };
 use crate::dsl::function_expr::FieldsMapper;
@@ -106,7 +106,7 @@ fn date_ranges(s: &[Series], interval: Duration, closed: ClosedWindow) -> Polars
         Ok(())
     };
 
-    let out = ranges_impl_broadcast(&start, &end, range_impl, &mut builder)?;
+    let out = temporal_ranges_impl_broadcast(&start, &end, range_impl, &mut builder)?;
 
     let to_type = DataType::List(Box::new(DataType::Date));
     out.cast(&to_type)

crates/polars-plan/src/dsl/function_expr/range/datetime_range.rs

@@ -5,7 +5,7 @@ use polars_core::series::Series;
 use polars_time::{datetime_range_impl, ClosedWindow, Duration};
 
 use super::utils::{
-    ensure_range_bounds_contain_exactly_one_value, ranges_impl_broadcast,
+    ensure_range_bounds_contain_exactly_one_value, temporal_ranges_impl_broadcast,
     temporal_series_to_i64_scalar,
 };
 use crate::dsl::function_expr::FieldsMapper;
@@ -202,7 +202,7 @@ pub(super) fn datetime_ranges(
                 Ok(())
             };
 
-            ranges_impl_broadcast(start, end, range_impl, &mut builder)?
+            temporal_ranges_impl_broadcast(start, end, range_impl, &mut builder)?
         },
         _ => unimplemented!(),
     };

crates/polars-plan/src/dsl/function_expr/range/int_range.rs

@@ -2,7 +2,7 @@ use polars_core::prelude::*;
 use polars_core::series::{IsSorted, Series};
 use polars_core::with_match_physical_integer_polars_type;
 
-use super::utils::{ensure_range_bounds_contain_exactly_one_value, ranges_impl_broadcast};
+use super::utils::{ensure_range_bounds_contain_exactly_one_value, numeric_ranges_impl_broadcast};
 
 const CAPACITY_FACTOR: usize = 5;
 
@@ -70,15 +70,18 @@ where
     Ok(ca.into_series())
 }
 
-pub(super) fn int_ranges(s: &[Series], step: i64) -> PolarsResult<Series> {
+pub(super) fn int_ranges(s: &[Series]) -> PolarsResult<Series> {
     let start = &s[0];
     let end = &s[1];
+    let step = &s[2];
 
     let start = start.cast(&DataType::Int64)?;
     let end = end.cast(&DataType::Int64)?;
+    let step = step.cast(&DataType::Int64)?;
 
     let start = start.i64()?;
     let end = end.i64()?;
+    let step = step.i64()?;
 
     let len = std::cmp::max(start.len(), end.len());
     let mut builder = ListPrimitiveChunkedBuilder::<Int64Type>::new(
@@ -88,18 +91,19 @@ pub(super) fn int_ranges(s: &[Series], step: i64) -> PolarsResult<Series> {
         DataType::Int64,
     );
 
-    let range_impl = |start, end, builder: &mut ListPrimitiveChunkedBuilder<Int64Type>| {
-        match step {
-            1 => builder.append_iter_values(start..end),
-            2.. => builder.append_iter_values((start..end).step_by(step as usize)),
-            _ => builder.append_iter_values(
-                (end..start)
-                    .step_by(step.unsigned_abs() as usize)
-                    .map(|x| start - (x - end)),
-            ),
+    let range_impl =
+        |start, end, step: i64, builder: &mut ListPrimitiveChunkedBuilder<Int64Type>| {
+            match step {
+                1 => builder.append_iter_values(start..end),
+                2.. => builder.append_iter_values((start..end).step_by(step as usize)),
+                _ => builder.append_iter_values(
+                    (end..start)
+                        .step_by(step.unsigned_abs() as usize)
+                        .map(|x| start - (x - end)),
+                ),
+            };
+            Ok(())
         };
-        Ok(())
-    };
 
-    ranges_impl_broadcast(start, end, range_impl, &mut builder)
+    numeric_ranges_impl_broadcast(start, end, step, range_impl, &mut builder)
 }

crates/polars-plan/src/dsl/function_expr/range/mod.rs

@@ -28,9 +28,7 @@ pub enum RangeFunction {
         step: i64,
         dtype: DataType,
     },
-    IntRanges {
-        step: i64,
-    },
+    IntRanges,
     #[cfg(feature = "temporal")]
     DateRange {
         interval: Duration,
@@ -76,7 +74,7 @@ impl RangeFunction {
         use RangeFunction::*;
         let field = match self {
             IntRange { dtype, .. } => Field::new("int", dtype.clone()),
-            IntRanges { .. } => Field::new("int_range", DataType::List(Box::new(DataType::Int64))),
+            IntRanges => Field::new("int_range", DataType::List(Box::new(DataType::Int64))),
             #[cfg(feature = "temporal")]
             DateRange {
                 interval,
@@ -158,7 +156,7 @@ impl Display for RangeFunction {
         use RangeFunction::*;
         let s = match self {
             IntRange { .. } => "int_range",
-            IntRanges { .. } => "int_ranges",
+            IntRanges => "int_ranges",
             #[cfg(feature = "temporal")]
             DateRange { .. } => "date_range",
             #[cfg(feature = "temporal")]
@@ -183,8 +181,8 @@ impl From<RangeFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
             IntRange { step, dtype } => {
                 map_as_slice!(int_range::int_range, step, dtype.clone())
             },
-            IntRanges { step } => {
-                map_as_slice!(int_range::int_ranges, step)
+            IntRanges => {
+                map_as_slice!(int_range::int_ranges)
             },
             #[cfg(feature = "temporal")]
             DateRange {

crates/polars-plan/src/dsl/function_expr/range/time_range.rs

@@ -3,7 +3,7 @@ use polars_core::series::Series;
 use polars_time::{time_range_impl, ClosedWindow, Duration};
 
 use super::utils::{
-    ensure_range_bounds_contain_exactly_one_value, ranges_impl_broadcast,
+    ensure_range_bounds_contain_exactly_one_value, temporal_ranges_impl_broadcast,
     temporal_series_to_i64_scalar,
 };
 
@@ -59,7 +59,7 @@ pub(super) fn time_ranges(
         Ok(())
     };
 
-    let out = ranges_impl_broadcast(start, end, range_impl, &mut builder)?;
+    let out = temporal_ranges_impl_broadcast(start, end, range_impl, &mut builder)?;
 
     let to_type = DataType::List(Box::new(DataType::Time));
     out.cast(&to_type)

crates/polars-plan/src/dsl/function_expr/range/utils.rs

@@ -1,5 +1,5 @@
 use polars_core::prelude::{
-    polars_bail, polars_ensure, ChunkedArray, IntoSeries, ListBuilderTrait,
+    polars_bail, polars_ensure, ChunkedArray, Int64Chunked, IntoSeries, ListBuilderTrait,
     ListPrimitiveChunkedBuilder, PolarsIntegerType, PolarsResult, Series,
 };
 
@@ -21,8 +21,124 @@ pub(super) fn ensure_range_bounds_contain_exactly_one_value(
     Ok(())
 }
 
+/// Create a numeric ranges column from the given start/end/step columns and a range function.
+pub(super) fn numeric_ranges_impl_broadcast<T, U, F>(
+    start: &ChunkedArray<T>,
+    end: &ChunkedArray<T>,
+    step: &Int64Chunked,
+    range_impl: F,
+    builder: &mut ListPrimitiveChunkedBuilder<U>,
+) -> PolarsResult<Series>
+where
+    T: PolarsIntegerType,
+    U: PolarsIntegerType,
+    F: Fn(T::Native, T::Native, i64, &mut ListPrimitiveChunkedBuilder<U>) -> PolarsResult<()>,
+{
+    match (start.len(), end.len(), step.len()) {
+        (len_start, len_end, len_step) if len_start == len_end && len_start == len_step => {
+            build_numeric_ranges::<_, _, _, T, U, F>(
+                start.downcast_iter().flatten(),
+                end.downcast_iter().flatten(),
+                step.downcast_iter().flatten(),
+                range_impl,
+                builder,
+            )?;
+        },
+        (1, len_end, 1) => {
+            let start_scalar = start.get(0);
+            let step_scalar = step.get(0);
+            match (start_scalar, step_scalar) {
+                (Some(start), Some(step)) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    std::iter::repeat(Some(&start)),
+                    end.downcast_iter().flatten(),
+                    std::iter::repeat(Some(&step)),
+                    range_impl,
+                    builder,
+                )?,
+                _ => build_nulls(builder, len_end),
+            }
+        },
+        (len_start, 1, 1) => {
+            let end_scalar = end.get(0);
+            let step_scalar = step.get(0);
+            match (end_scalar, step_scalar) {
+                (Some(end), Some(step)) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    start.downcast_iter().flatten(),
+                    std::iter::repeat(Some(&end)),
+                    std::iter::repeat(Some(&step)),
+                    range_impl,
+                    builder,
+                )?,
+                _ => build_nulls(builder, len_start),
+            }
+        },
+        (1, 1, len_step) => {
+            let start_scalar = start.get(0);
+            let end_scalar = end.get(0);
+            match (start_scalar, end_scalar) {
+                (Some(start), Some(end)) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    std::iter::repeat(Some(&start)),
+                    std::iter::repeat(Some(&end)),
+                    step.downcast_iter().flatten(),
+                    range_impl,
+                    builder,
+                )?,
+                _ => build_nulls(builder, len_step),
+            }
+        },
+        (len_start, len_end, 1) if len_start == len_end => {
+            let step_scalar = step.get(0);
+            match step_scalar {
+                Some(step) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    start.downcast_iter().flatten(),
+                    end.downcast_iter().flatten(),
+                    std::iter::repeat(Some(&step)),
+                    range_impl,
+                    builder,
+                )?,
+                None => build_nulls(builder, len_start),
+            }
+        },
+        (len_start, 1, len_step) if len_start == len_step => {
+            let end_scalar = end.get(0);
+            match end_scalar {
+                Some(end) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    start.downcast_iter().flatten(),
+                    std::iter::repeat(Some(&end)),
+                    step.downcast_iter().flatten(),
+                    range_impl,
+                    builder,
+                )?,
+                None => build_nulls(builder, len_start),
+            }
+        },
+        (1, len_end, len_step) if len_end == len_step => {
+            let start_scalar = start.get(0);
+            match start_scalar {
+                Some(start) => build_numeric_ranges::<_, _, _, T, U, F>(
+                    std::iter::repeat(Some(&start)),
+                    end.downcast_iter().flatten(),
+                    step.downcast_iter().flatten(),
+                    range_impl,
+                    builder,
+                )?,
+                None => build_nulls(builder, len_end),
+            }
+        },
+        (len_start, len_end, len_step) => {
+            polars_bail!(
+                ComputeError:
+                "lengths of `start` ({}), `end` ({}) and `step` ({}) do not match",
+                len_start, len_end, len_step
+            )
+        },
+    };
+    let out = builder.finish().into_series();
+    Ok(out)
+}
+
 /// Create a ranges column from the given start/end columns and a range function.
-pub(super) fn ranges_impl_broadcast<T, U, F>(
+pub(super) fn temporal_ranges_impl_broadcast<T, U, F>(
     start: &ChunkedArray<T>,
     end: &ChunkedArray<T>,
     range_impl: F,
@@ -35,7 +151,7 @@ where
 {
     match (start.len(), end.len()) {
         (len_start, len_end) if len_start == len_end => {
-            build_ranges::<_, _, T, U, F>(
+            build_temporal_ranges::<_, _, T, U, F>(
                 start.downcast_iter().flatten(),
                 end.downcast_iter().flatten(),
                 range_impl,
@@ -45,7 +161,7 @@ where
         (1, len_end) => {
             let start_scalar = start.get(0);
             match start_scalar {
-                Some(start) => build_ranges::<_, _, T, U, F>(
+                Some(start) => build_temporal_ranges::<_, _, T, U, F>(
                     std::iter::repeat(Some(&start)),
                     end.downcast_iter().flatten(),
                     range_impl,
@@ -57,7 +173,7 @@ where
         (len_start, 1) => {
             let end_scalar = end.get(0);
             match end_scalar {
-                Some(end) => build_ranges::<_, _, T, U, F>(
+                Some(end) => build_temporal_ranges::<_, _, T, U, F>(
                     start.downcast_iter().flatten(),
                     std::iter::repeat(Some(&end)),
                     range_impl,
@@ -78,8 +194,33 @@ where
     Ok(out)
 }
 
+/// Iterate over a start and end column and create a range with the step for each entry.
+fn build_numeric_ranges<'a, I, J, K, T, U, F>(
+    start: I,
+    end: J,
+    step: K,
+    range_impl: F,
+    builder: &mut ListPrimitiveChunkedBuilder<U>,
+) -> PolarsResult<()>
+where
+    I: Iterator<Item = Option<&'a T::Native>>,
+    J: Iterator<Item = Option<&'a T::Native>>,
+    K: Iterator<Item = Option<&'a i64>>,
+    T: PolarsIntegerType,
+    U: PolarsIntegerType,
+    F: Fn(T::Native, T::Native, i64, &mut ListPrimitiveChunkedBuilder<U>) -> PolarsResult<()>,
+{
+    for ((start, end), step) in start.zip(end).zip(step) {
+        match (start, end, step) {
+            (Some(start), Some(end), Some(step)) => range_impl(*start, *end, *step, builder)?,
+            _ => builder.append_null(),
+        }
+    }
+    Ok(())
+}
+
 /// Iterate over a start and end column and create a range for each entry.
-fn build_ranges<'a, I, J, T, U, F>(
+fn build_temporal_ranges<'a, I, J, T, U, F>(
     start: I,
     end: J,
     range_impl: F,

crates/polars-plan/src/dsl/functions/range.rs

@@ -22,12 +22,12 @@ pub fn int_range(start: Expr, end: Expr, step: i64, dtype: DataType) -> Expr {
 }
 
 /// Generate a range of integers for each row of the input columns.
-pub fn int_ranges(start: Expr, end: Expr, step: i64) -> Expr {
-    let input = vec![start, end];
+pub fn int_ranges(start: Expr, end: Expr, step: Expr) -> Expr {
+    let input = vec![start, end, step];
 
     Expr::Function {
         input,
-        function: FunctionExpr::Range(RangeFunction::IntRanges { step }),
+        function: FunctionExpr::Range(RangeFunction::IntRanges),
         options: FunctionOptions {
             allow_rename: true,
             ..Default::default()