-
Notifications
You must be signed in to change notification settings - Fork 146
feat(DRAFT): Adds (Expr|Series).first()
#2528
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. Weβll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from 23 commits
ff661ae
1b77bd7
e84cba3
25ef241
78822aa
4075c50
45f24b9
4041dd1
bb9912d
6a53aa1
4efc939
fc149c1
7489e61
d2719a4
0af11db
afe20f0
6c0bd6f
e0fdf78
aa7c510
4fdc0aa
bd4ab89
9f7f5a9
ddb50d2
7146f60
8c24e6e
54a4cb4
63e0459
9493aad
88535a4
77ae9c0
c1a6173
3c4ff9b
696e35d
b2866d2
cd002f3
1458530
962ebcd
5d310bc
a417341
354da1a
0cea41b
5229096
8d3aaec
a62e3ef
9c36285
ad8e3f7
628f71e
deacc71
5c52ee4
eec2a4f
e003bab
fb2dc1c
211673b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -23,6 +23,7 @@ | |
- ewm_mean | ||
- fill_null | ||
- filter | ||
- first | ||
- gather_every | ||
- head | ||
- clip | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -28,6 +28,7 @@ | |
- ewm_mean | ||
- fill_null | ||
- filter | ||
- first | ||
- gather_every | ||
- head | ||
- hist | ||
|
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
@@ -408,6 +408,12 @@ def _clip_both( | |||||||||||||||||||||||||||||||
_clip_both, lower_bound=lower_bound, upper_bound=upper_bound | ||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||
def first(self) -> Self: | ||||||||||||||||||||||||||||||||
def fn(_input: duckdb.Expression) -> duckdb.Expression: | ||||||||||||||||||||||||||||||||
return FunctionExpression("first", _input) | ||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||
return self._with_callable(fn) | ||||||||||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. initial feedback: There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks @MarcoGorelli, so first step will be Lines 71 to 75 in dd41833
Then see what to do in each backend There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. One thing I thought worth mentioning was that I don't think There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. duckdb seems to have the same behavior as There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @MarcoGorelli these are the two other cases we have for Lines 785 to 786 in b7001e4
Lines 808 to 809 in b7001e4
We currently don't support them in narwhals/narwhals/_compliant/expr.py Lines 879 to 884 in b7001e4
I'm just pushing what I think is how to enforce the |
||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||
def sum(self) -> Self: | ||||||||||||||||||||||||||||||||
return self._with_callable(lambda _input: FunctionExpression("sum", _input)) | ||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -560,6 +560,14 @@ def _clip_both( | |
_clip_both, lower_bound=lower_bound, upper_bound=upper_bound | ||
) | ||
|
||
def first(self) -> Self: | ||
def fn(inputs: WindowInputs) -> Column: | ||
return self._F.first(inputs.expr, ignorenulls=False).over( | ||
self.partition_by(inputs).orderBy(*self._sort(inputs)) | ||
) | ||
|
||
return self._with_window_function(fn) | ||
Comment on lines
+541
to
+545
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. So it seems I missed this in the review of (#2600) π€¦ββοΈ In this diff (https://github.com/narwhals-dev/narwhals/pull/2600/files#diff-ee3b03ae02617c27c275264a02582a80e186283f3989b9e27ea7143a9161fe76) there's a reversal of an abstraction I did in (#2505) I'm not following why the below is preferred: return self._F.first(inputs.expr, ignorenulls=False).over(
self.partition_by(*inputs.partition_by).orderBy(
*self._sort(*inputs.order_by)
)
)
return self._with_window_function(fn) |
||
|
||
def is_finite(self) -> Self: | ||
def _is_finite(_input: Column) -> Column: | ||
# A value is finite if it's not NaN, and not infinite, while NULLs should be | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -801,6 +801,25 @@ def clip( | |
) | ||
) | ||
|
||
def first(self) -> Any: | ||
"""Get the first element of the Series. | ||
|
||
Returns: | ||
A scalar value or `None` if the Series is empty. | ||
|
||
Examples: | ||
>>> import polars as pl | ||
>>> import narwhals as nw | ||
>>> | ||
>>> s_native = pl.Series([1, 2, 3]) | ||
>>> s_nw = nw.from_native(s_native, series_only=True) | ||
>>> s_nw.first() | ||
1 | ||
>>> s_nw.filter(s_nw > 5).first() is None | ||
True | ||
Comment on lines
+807
to
+816
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't like the I think it's important to have an example for that case though - since |
||
""" | ||
return self._compliant_series.first() | ||
|
||
def is_in(self, other: Any) -> Self: | ||
"""Check if the elements of this Series are in the other sequence. | ||
|
||
|
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,68 @@ | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from __future__ import annotations | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from typing import TYPE_CHECKING | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from typing import Mapping | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from typing import Sequence | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
import pytest | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
import narwhals as nw | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from tests.utils import assert_equal_data | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
if TYPE_CHECKING: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from narwhals.typing import PythonLiteral | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
from tests.utils import ConstructorEager | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
data = { | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"a": [8, 2, 1, None], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"b": [58, 5, 6, 12], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"c": [2.5, 1.0, 3.0, 0.9], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"d": [2, 1, 4, 3], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
} | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
@pytest.mark.parametrize(("col", "expected"), [("a", 8), ("b", 58), ("c", 2.5)]) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def test_first_series( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
constructor_eager: ConstructorEager, col: str, expected: PythonLiteral | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
series = nw.from_native(constructor_eager(data), eager_only=True)[col] | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
result = series.first() | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assert_equal_data({col: [result]}, {col: [expected]}) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def test_first_series_empty(constructor_eager: ConstructorEager) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
series = nw.from_native(constructor_eager(data), eager_only=True)["a"] | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
series = series.filter(series > 50) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
result = series.first() | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assert result is None | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
@pytest.mark.parametrize(("col", "expected"), [("a", 8), ("b", 58), ("c", 2.5)]) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def test_first_expr_eager( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
constructor_eager: ConstructorEager, col: str, expected: PythonLiteral | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
df = nw.from_native(constructor_eager(data)) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
expr = nw.col(col).first() | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
result = df.select(expr) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assert_equal_data(result, {col: [expected]}) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
Comment on lines
+39
to
+46
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Feel like I got a bit unlucky with this being the first test I wrote π So there's a wrinkle with how the This is all good: import polars as pl
data = {
"a": [8, 2, 1, None],
"b": [58, 5, 6, 12],
"c": [2.5, 1.0, 3.0, 0.9],
"d": [2, 1, 4, 3],
"idx": [0, 1, 2, 3],
}
df = pl.DataFrame(data)
>>> df.select(pl.col("a").first())
shape: (1, 1)
βββββββ
β a β
β --- β
β i64 β
βββββββ‘
β 8 β
βββββββ
>>> df.lazy().select(pl.col("a").first()).collect()
shape: (1, 1)
βββββββ
β a β
β --- β
β i64 β
βββββββ‘
β 8 β
βββββββ We can also do use a >>> df.lazy().select(pl.col("a").sort_by("idx").first()).collect()
shape: (1, 1)
βββββββ
β a β
β --- β
β i64 β
βββββββ‘
β 8 β
βββββββ But if we do that after, the sort column has the pre-agg shape: >>> df.lazy().select(pl.col("a").first().sort_by("idx")).collect()
ShapeError: `sort_by` produced different length (4) than the Series that has to be sorted (1) If we do >>> df.lazy().select(pl.col("a").first().over(pl.lit(1), order_by="idx")).collect()
shape: (4, 1)
βββββββ
β a β
β --- β
β i64 β
βββββββ‘
β 8 β
β 8 β
β 8 β
β 8 β
βββββββ @MarcoGorelli would we want to land (#2534) first so that we have a way to specify this as an aggregation? I do hope there's another way we can do this with the existing There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The example for Lines 724 to 742 in 6c110ca
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
narwhals/tests/expr_and_series/arg_min_test.py Lines 44 to 56 in 6c110ca
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The thing is that arg_min is not supported in over context for any other backend than polars.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @FBruzzesi I know There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I guess the point I'm trying to make is that adding the constraint of an This is what we'd need to suggest, since that's the way to maintain the aggregation in
I'm just a little lost since the rules we've been working on are for after the aggregation - whereas this is flipped π€ There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Ah yeah I'm getting that locally as well - I'll push the tests as-is for now |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
@pytest.mark.parametrize( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
"expected", | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
[{"a": [8], "c": [2.5]}, {"d": [2], "b": [58]}, {"c": [2.5], "a": [8], "d": [2]}], | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def test_first_expr_eager_expand( | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
constructor_eager: ConstructorEager, expected: Mapping[str, Sequence[PythonLiteral]] | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
df = nw.from_native(constructor_eager(data)) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
expr = nw.col(expected).first() | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
result = df.select(expr) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assert_equal_data(result, expected) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
|
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
def test_first_expr_eager_expand_sort(constructor_eager: ConstructorEager) -> None: | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
df = nw.from_native(constructor_eager(data)) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
expr = nw.col("d", "a", "b", "c").first() | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
result = df.sort("d").select(expr) | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
expected = {"d": [1], "a": [2], "b": [5], "c": [1.0]} | ||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
assert_equal_data(result, expected) |
Uh oh!
There was an error while loading. Please reload this page.