feat: Add `linear_space` #20678

mcrumiller · 2025-01-12T23:29:18Z

First pass at #5255.

Update: linear space with temporals is now implemented, along with an f32 path, and stronger testing.

import polars as pl
from datetime import date, datetime, time, timedelta

pl.linear_space(date(2025, 1, 1), date(2025, 2, 1), 3, eager=True)
# shape: (3,)
# Series: 'literal' [datetime[ms]]
# [
#         2025-01-01 00:00:00
#         2025-01-16 12:00:00
#         2025-02-01 00:00:00
# ]

pl.linear_space(pl.lit(0, pl.Float32), pl.lit(1, pl.Float32), 5, eager=True)
# shape: (5,)
# Series: 'literal' [f32]
# [
#         0.0
#         0.25
#         0.5
#         0.75
#         1.0
# ]

I kept the name linear_space since that's a more canonical polars name, although linspace is far more recognizable. The usage is fairly simple, and the closed is one of both (default), left, right, or none. In each case, the number of samples asked for is provided.

import polars as pl

pl.linear_space(0, 5, 3, closed="both", eager=True)
# shape: (3,)
# Series: 'literal' [f64]
# [
#         0.0
#         2.5
#         5.0
# ]

pl.linear_space(0, 1, 3, closed="left", eager=True)
# shape: (3,)
# Series: 'literal' [f64]
# [
#         0.0
#         0.333333
#         0.666667
# ]

>>> pl.linear_space(0, 1, 5, closed="none", eager=True)
# shape: (5,)
# Series: 'literal' [f64]
# [
#         0.166667
#         0.333333
#         0.5
#         0.666667
#         0.833333
# ]

mcrumiller · 2025-01-12T23:30:54Z

crates/polars-ops/src/series/ops/is_between.rs


-#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Default, IntoStaticStr)]


I moved ClosedInterval into linear_space since it's feature-gated here and I wanted to re-use it. I'm not sure if there's a better place for it.

codecov · 2025-01-13T08:01:45Z

Codecov Report

Attention: Patch coverage is 90.37433% with 18 lines in your changes missing coverage. Please review.

Project coverage is 79.74%. Comparing base (ce9ad17) to head (02e6603).
Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
crates/polars-ops/src/series/ops/linear_space.rs	84.37%	10 Missing ⚠️
...s-plan/src/dsl/function_expr/range/linear_space.rs	88.13%	7 Missing ⚠️
...tes/polars-plan/src/dsl/function_expr/range/mod.rs	94.44%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #20678    +/-   ##
========================================
  Coverage   79.73%   79.74%            
========================================
  Files        1561     1564     +3     
  Lines      221913   222097   +184     
  Branches     2530     2531     +1     
========================================
+ Hits       176943   177105   +162     
- Misses      44388    44410    +22     
  Partials      582      582

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

mcrumiller · 2025-01-13T17:20:21Z

I don't think closed = "none" should exist.

I wasn't sure about that one myself. Will update.

orlp · 2025-01-13T17:24:59Z

@mcrumiller I changed my mind right after and deleted my comment, but I guess I wasn't fast enough. I think it's alright to have - equidistant points over the range, not including either start or end.

I'm not sure about the name "none" though. Maybe it should be "neither"? Do we have precedence for a closed argument with "none" anywhere else?

mcrumiller · 2025-01-13T17:59:04Z

I'm not sure about the name "none" though. Maybe it should be "neither"? Do we have precedence for a closed argument with "none" anywhere else?

I was thinking about this as well. A better-suited parameter name might be something like interval_type where we would specify closed-left, closed-right, closed, or open. The parameter name "closed" is essentially a question "how is the interval closed?" and the answer none is saying "it's not". It's a bit awkward.

I grabbed the interval enumeration from is_between. So in the interest of continuity it might be best to leave the not-great "closed" option, but I don't have a strong preference either way.

orlp · 2025-01-14T08:52:05Z

Let's keep using the established closed terminology then.

mcrumiller · 2025-01-14T12:11:16Z

Sure. Do you have any comments on improvements? Should we include a return 'dtype' parameter like in numpy's implementation?

orlp · 2025-01-14T12:25:35Z

Rust side mostly looks good to me. I don't think we need a dtype parameter at this time.

However I think there's still a couple things missing around dtypes:

If the input types are all f32, I think the output should probably also be f32.
If the input types are all dates/times, I think the output should also be dates/times.
I think we should reject inputs which are not strictly numeric, strictly date or strictly time. You should not be able to specify a start date but end time for example.

mcrumiller · 2025-01-14T21:25:33Z

@orlp, latest round of updates addressing your suggestions:

If the input types are all f32, I think the output should probably also be f32.

I've added separate floating implementations (I originally went for generics but it felt messy to me, so I just made two separate implementations). Included new tests.

If the input types are all dates/times, I think the output should also be dates/times.

I've added a temporal path. start/end points with dates produces a Datetime("ms") (similar to mean/median) and all other temporal types require the same temporal type. If the start/end data types do not match exactly (except for numeric, which are fine to mix and match) (even the time precision units are required to match), we raise.

I think we should reject inputs which are not strictly numeric, strictly date or strictly time. You should not be able to specify a start date but end time for example.

Done--only numeric/temporal work. Again, the start/end must match, and anything outside of these dtypes throws an Exception.

orlp · 2025-01-15T08:43:53Z

py-polars/polars/functions/range/linear_space.py

+    closed: ClosedInterval = "both",
+    eager: bool = False,
+) -> Expr | Series:
+    """Linearly-spaced elements."""


I think the documentation could be a bit better, perhaps with an example for start = 0, end = 1, num_samples = 4 for each of the closed options?

Yep forgot docstrings.

mcrumiller · 2025-01-15T13:15:39Z

@orlp I think we may want num_samples to be expressionable, as a somewhat obvious use case is to generate something like pl.linspace(0, 1, pl.len()), e.g. use the height of the frame as the number of samples. Do you agree?

orlp · 2025-01-15T14:10:08Z

@mcrumiller Yes, I agree, but it must resolve to a single value, similar to this error:

>>> df = pl.DataFrame({"x": [1, 2, 3]})
>>> df.select(pl.int_range(0, pl.col.x))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/orlp/.localpython/lib/python3.11/site-packages/polars/dataframe/frame.py", line 9323, in select
    return self.lazy().select(*exprs, **named_exprs).collect(_eager=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/orlp/.localpython/lib/python3.11/site-packages/polars/lazyframe/frame.py", line 2057, in collect
    return wrap_df(ldf.collect(callback))
                   ^^^^^^^^^^^^^^^^^^^^^
polars.exceptions.ComputeError: `end` must contain exactly one value, got 3 values

mcrumiller · 2025-01-16T00:26:37Z

@orlp updated num_samples to be an expression that only allows one value, added some tests as well, and updated the docstrings.

orlp · 2025-01-16T08:53:56Z

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

@@ -70,6 +75,7 @@ impl RangeFunction {
        match self {
            IntRange { dtype, .. } => mapper.with_dtype(dtype.clone()),
            IntRanges => mapper.with_dtype(DataType::List(Box::new(DataType::Int64))),
+            LinearSpace { closed: _ } => mapper.with_dtype(DataType::Float64),


This is no longer correct, right? The output dtype isn't just Float64.

Thanks Orson, should have used LazyFrames in my tests. Fixed.

I now assume in the actual implementation that the schema has been validated already, which lets me remove a match arm. Is this ok to do? In other words, do I need to re-check the dtypes in the linear space implementation, or can I assume that an invalid input schema would have already been caught in the dsl?

@mcrumiller Please verify the types in the implementation as extra backup, I don't think eliding it makes much difference for performance.

Ok, I've reverted to include the type verification in the implementation.

orlp · 2025-01-16T08:54:53Z

py-polars/polars/functions/range/linear_space.py

+    start
+        Lower bound of the range.
+    end
+        Upper bound of the time range.


Why does this say "time range"?

orlp · 2025-01-16T08:56:30Z

py-polars/polars/functions/range/linear_space.py

+            0.333333
+            0.666667
+            1.0
+    ]


Missing a closed="both" example.

I've added a closed="none" example which was also missing. Do you think I should be explicit in the base example with the closed="both" parameter? This is the default so it's not strictly necessary.

>>> pl.linear_space(start=0, end=1, num_samples=3, eager=True) shape: (3,) Series: 'literal' [f64] [ 0.0 0.5 1.0 ]

(I have now explicitly provided closed="both" in the example).

I meant closed="none", sorry.

Got it. I've removed the closed="both" explicit declaration since it shows the default usage a bit better.

mcrumiller · 2025-01-16T16:05:08Z

@orlp I think I have addressed all of your concerns. Thanks for the thorough review.

github-actions bot added enhancement New feature or an improvement of an existing feature python Related to Python Polars rust Related to Rust Polars labels Jan 12, 2025

mcrumiller commented Jan 12, 2025

View reviewed changes

mcrumiller marked this pull request as ready for review January 13, 2025 03:26

mcrumiller requested review from ritchie46, c-peters, alexander-beedie, MarcoGorelli, reswqa and orlp as code owners January 13, 2025 03:26

mcrumiller marked this pull request as draft January 14, 2025 12:51

mcrumiller force-pushed the linear-space branch from d64b15c to 97bbae8 Compare January 14, 2025 20:40

mcrumiller marked this pull request as ready for review January 14, 2025 21:18

orlp requested changes Jan 15, 2025

View reviewed changes

mcrumiller marked this pull request as draft January 15, 2025 13:01

mcrumiller force-pushed the linear-space branch 2 times, most recently from 2f39ae6 to 6627366 Compare January 15, 2025 23:07

mcrumiller marked this pull request as ready for review January 16, 2025 00:26

orlp requested changes Jan 16, 2025

View reviewed changes

mcrumiller marked this pull request as draft January 16, 2025 13:27

mcrumiller added 7 commits January 16, 2025 09:27

Add linear_space

01615e5

Add missing func

e708d2c

Add more dtypes and tests

61b849f

Make 'num_samples' expr and add docstring

186f18e

Add special condition message for streaming test

77f94a5

Remove the word 'time' in upper bound description

526b45f

Check schema in plan

29b6992

mcrumiller force-pushed the linear-space branch 2 times, most recently from e513cca to 29b6992 Compare January 16, 2025 15:15

Remove 'both' in default example

02e6603

mcrumiller marked this pull request as ready for review January 16, 2025 15:49

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add `linear_space` #20678

feat: Add `linear_space` #20678

mcrumiller commented Jan 12, 2025 •

edited

Loading

mcrumiller Jan 12, 2025 •

edited

Loading

codecov bot commented Jan 13, 2025 •

edited

Loading

mcrumiller commented Jan 13, 2025

orlp commented Jan 13, 2025 •

edited

Loading

mcrumiller commented Jan 13, 2025 •

edited

Loading

orlp commented Jan 14, 2025

mcrumiller commented Jan 14, 2025

orlp commented Jan 14, 2025

mcrumiller commented Jan 14, 2025 •

edited

Loading

orlp Jan 15, 2025

mcrumiller Jan 15, 2025

mcrumiller commented Jan 15, 2025

orlp commented Jan 15, 2025

mcrumiller commented Jan 16, 2025

orlp Jan 16, 2025

mcrumiller Jan 16, 2025

orlp Jan 16, 2025

mcrumiller Jan 16, 2025

orlp Jan 16, 2025

mcrumiller Jan 16, 2025

orlp Jan 16, 2025

mcrumiller Jan 16, 2025

mcrumiller Jan 16, 2025

orlp Jan 16, 2025

mcrumiller Jan 16, 2025

mcrumiller commented Jan 16, 2025


		#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq, Default, IntoStaticStr)]

+.333333
+.666667
+.0
+                  ]

feat: Add linear_space #20678

Are you sure you want to change the base?

feat: Add linear_space #20678

Conversation

mcrumiller commented Jan 12, 2025 • edited Loading

mcrumiller Jan 12, 2025 • edited Loading

Choose a reason for hiding this comment

codecov bot commented Jan 13, 2025 • edited Loading

Codecov Report

mcrumiller commented Jan 13, 2025

orlp commented Jan 13, 2025 • edited Loading

mcrumiller commented Jan 13, 2025 • edited Loading

orlp commented Jan 14, 2025

mcrumiller commented Jan 14, 2025

orlp commented Jan 14, 2025

mcrumiller commented Jan 14, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mcrumiller commented Jan 15, 2025

orlp commented Jan 15, 2025

mcrumiller commented Jan 16, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mcrumiller commented Jan 16, 2025

feat: Add `linear_space` #20678

feat: Add `linear_space` #20678

mcrumiller commented Jan 12, 2025 •

edited

Loading

mcrumiller Jan 12, 2025 •

edited

Loading

codecov bot commented Jan 13, 2025 •

edited

Loading

orlp commented Jan 13, 2025 •

edited

Loading

mcrumiller commented Jan 13, 2025 •

edited

Loading

mcrumiller commented Jan 14, 2025 •

edited

Loading