fix: Add scalar checks for n and fill_value parameters in shift

[mcrumiller]

Closes #21280.

New behavior:

import polars as pl

df = pl.DataFrame({
    "x": [1, 2, 3],
    "y": [4, 5, 6]
})

# This used to succeed but silently use y.get(0).
df.shift(2, fill_value=pl.col("y"))
# polars.exceptions.ComputeError: 'fill_value' must be scalar value

df.shift(2, fill_value=pl.col("y").first())  # Ok
# shape: (3, 2)
# ┌─────┬─────┐
# │ x   ┆ y   │
# │ --- ┆ --- │
# │ i64 ┆ i64 │
# ╞═════╪═════╡
# │ 4   ┆ 4   │
# │ 4   ┆ 4   │
# │ 1   ┆ 4   │
# └─────┴─────┘

# Expressions are not in the function signature, but they work. We can still verify.
df.shift(pl.col("y"), fill_value=1)
# polars.exceptions.ComputeError: 'n' must be scalar value

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 79.92%. Comparing base (4e286d8) to head (b408e6a).
Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main   #21292   +/-   ##
=======================================
  Coverage   79.92%   79.92%           
=======================================
  Files        1596     1596           
  Lines      228745   228745           
  Branches     2616     2616           
=======================================
+ Hits       182825   182832    +7     
+ Misses      45321    45314    -7     
  Partials      599      599           

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

[ritchie46]

I agree about this, but I think we should catch this at the IR. I believe we have more is_scalar checks in place there.

[mcrumiller]

Thanks @ritchie46, was just looking into this actually. Will follow up shortly.

[mcrumiller]

@Ritchie not sure this is possible actually. For reference, I put the check in crates/polars-plan/src/plans/conversion/functions.rs, and used all_returns_scalar(expr) on the argument in question--is this the way to do it?

The problem is that when fill_value is a datetime with a time zone, it gets replaced by a replace_time_zone() call, which doesn't return scalar. I can't see any way to ensure before evaluation that we have a scalar when this path is taken.

[mcrumiller]

I tried updating the replace_time_zone to set its returns_scalar value to be the input's return_scalar value, but I'm getting a 'returns_scalar' and 'elementwise' are mutually exclusive. I can see why this was thought to be true but perhaps they shouldn't be mutually exclusive. For example, if we can show that x is a scalar, and f doesn't change the length, then f(x) should also return scalar even if f is elementwise.

[mcrumiller]

I may open a separate PR for this to see what you think.

[nameexhaustion]

I don't think you need to remove the mutual exclusion

For this PR the function you likely need to call is ExprIR::is_scalar (or is_scalar_ae) -

polars/crates/polars-plan/src/plans/expr_ir.rs

Line 251 in de1d9d5

pub fn is_scalar(&self, expr_arena: &Arena<AExpr>) -> bool {

You should be able to call it to check the input expr at some point during the conversion.

[mcrumiller]

@nameexhaustion If you drill down into is_scalar you'll find that the scalar properties come directly from the DSL's Expr.is_scalar. In my example of using dt.round, we have:

ExprIR.is_scalar
→ is_scalar_ae
→ AExpr::Function.options.flags.contains(FunctionFlags::RETURNS_SCALAR)

The FunctionFlags::RETURNS_SCALAR in that last call there are set directly in map_many_private. In the particular case of dt.round, this flag is explicitly set to false in round:

polars/crates/polars-plan/src/dsl/dt.rs

Lines 244 to 251 in de1d9d5

	pub fn round(self, every: Expr) -> Expr {
	self.0.map_many_private(
	FunctionExpr::TemporalExpr(TemporalFunction::Round),
	&[every],
	false,
	None,
	)
	}

Instead of explicitly setting to false, it should be set to true if and only if all of its inputs are scalar, and this should be easily be checked via recursion. However, it can't, because the function is also elementwise.

In general, the issue is that certain function expressions return False for returns_scalar even if we theoretically can say for certain that they do indeed return a scalar (i.e. if a function doesn't change the length and its input is scalar).

[nameexhaustion]

In general, the issue is that certain function expressions return False for returns_scalar even if we theoretically can say for certain that they do indeed return a scalar

There is already a branch in is_scalar_ae that should correctly identify this case -

polars/crates/polars-plan/src/plans/aexpr/scalar.rs

Lines 13 to 17 in 109cc7e

	} else if options.is_elementwise()
	|| !options.flags.contains(FunctionFlags::CHANGES_LENGTH)
	{
	input.iter().all(|e| e.is_scalar(expr_arena))
	} else {

I expect it should work if you call it

[mcrumiller]

@nameexhaustion I completely missed that branch, you're right, my apologies. Let me look into why this didn't work for dates...

[mcrumiller]

@nameexhaustion I think I got it working, thanks a bunch for your help.

[mcrumiller]

I removed this check here since we now check in the IR.

[mcrumiller]

@ritchie46 I've moved the check. Let me know if you think this is a good spot for it. If we have other functions that require argument checks, this section can be switched to a match.