diff --git a/stratum/optimizer/logical/_column_expr.py b/stratum/optimizer/logical/_column_expr.py index cfc75c88..57fa670b 100644 --- a/stratum/optimizer/logical/_column_expr.py +++ b/stratum/optimizer/logical/_column_expr.py @@ -28,7 +28,8 @@ ColumnMethodOp, get_column_method_spec, polars_dtype) from stratum.optimizer.logical._map_ops import ( AccessorMapOp, DatetimeConversionOp, StringMethodOp, STR_POLARS_METHODS, - polars_datetime_kwargs, polars_str_method_unsupported, polars_string_method) + polars_datetime_kwargs, polars_dt_attribute, polars_str_method_unsupported, + polars_string_method) from stratum.optimizer.logical._projection_ops import ColumnProjectionOp # operator callable -> symbol. A binary/unary op whose callable is not in the @@ -404,8 +405,8 @@ def has_aggregate(self): class DtExpr(ColumnExpr): """Datetime accessor attribute (``.dt.``). - pandas reads the attribute off ``.dt``; polars calls a method, remapping a - few names via ``AccessorMapOp.POLARS_ATTR_NAME_MAP``. + pandas reads the attribute off ``.dt``; polars calls a method through + :func:`polars_dt_attribute`, which remaps names and values to pandas'. """ __slots__ = ("operand", "attr") @@ -424,11 +425,7 @@ def to_pandas(self, ctx): return getattr(obj.dt, self.attr) def to_polars(self, ctx): - obj = self.operand.to_polars(ctx) - if self.attr == "is_month_end": - return obj.dt.month_end() == obj - name = AccessorMapOp.POLARS_ATTR_NAME_MAP.get(self.attr, self.attr) - return getattr(obj.dt, name)() + return polars_dt_attribute(self.operand.to_polars(ctx), self.attr) def _polars_unsupported_here(self): name = AccessorMapOp.POLARS_ATTR_NAME_MAP.get(self.attr, self.attr) diff --git a/stratum/optimizer/logical/_map_ops.py b/stratum/optimizer/logical/_map_ops.py index 3ad77073..d5c6dd66 100644 --- a/stratum/optimizer/logical/_map_ops.py +++ b/stratum/optimizer/logical/_map_ops.py @@ -312,10 +312,14 @@ class AccessorMapOp(MapOp): """ fields = ["attr_name"] - # NOTE: Polars and Pandas differ in semantics for some datetime attributes: - # - dayofweek: Pandas uses Monday=0, Polars weekday() uses Monday=1 (ISO 8601) - # - dayofyear: Pandas is 1-indexed, Polars ordinal_day() is also 1-indexed (same) - POLARS_ATTR_NAME_MAP = {"dayofweek": "weekday","dayofyear": "ordinal_day"} + # pandas `.dt.` -> polars `.dt.`, for the attributes whose names + # differ between backends. Some also differ in *values*; those are corrected + # in :func:`polars_dt_attribute`. + POLARS_ATTR_NAME_MAP = { + "dayofweek": "weekday", "day_of_week": "weekday", + "dayofyear": "ordinal_day", "day_of_year": "ordinal_day", + "daysinmonth": "days_in_month", + } def __init__(self, attr_name: list[str] | str = None, inputs: list[Op] = None, outputs: list[Op] = None): if attr_name is None: @@ -365,6 +369,26 @@ def propagate_output_schema(self): self.output_schema = None +def polars_dt_attribute(obj, attr: str): + """Read pandas' ``.dt.`` off a polars Series or expression. + + polars exposes the attributes as methods, some under other names (see + :data:`AccessorMapOp.POLARS_ATTR_NAME_MAP`), and counts two of them + differently: weekdays from Monday=1 (ISO 8601, pandas: Monday=0) and + nanoseconds past the second (pandas: past the microsecond, 0..999). The + result carries pandas' values in polars' integer dtype. + """ + if attr == "is_month_end": + return obj.dt.month_end() == obj + name = AccessorMapOp.POLARS_ATTR_NAME_MAP.get(attr, attr) + result = getattr(obj.dt, name)() + if name == "weekday": + return result - 1 + if name == "nanosecond": + return result % 1000 + return result + + class AssignMapOp(MapOp): """``df.assign(...)`` with each assigned column folded to a ``ColumnExpr``. diff --git a/stratum/optimizer/physical/_map_execs.py b/stratum/optimizer/physical/_map_execs.py index 2b9630b0..529265bd 100644 --- a/stratum/optimizer/physical/_map_execs.py +++ b/stratum/optimizer/physical/_map_execs.py @@ -24,7 +24,7 @@ from stratum.optimizer.logical._map_ops import ( AccessorMapOp, ApplyUDFOp, AssignMapOp, AssignOp, CallMapOp, DatetimeConversionOp, StringMethodOp, polars_datetime_kwargs, - polars_string_method, polars_str_method_unsupported) + polars_dt_attribute, polars_string_method, polars_str_method_unsupported) from stratum.optimizer.physical._physical_ops import PhysicalOp from stratum.optimizer.physical._representation import refuse_unsupported_exprs from stratum.optimizer.physical._registry import physical_impl @@ -184,18 +184,20 @@ def _to_datetime_via_pandas(obj, args, kwargs): @physical_impl(of=DatetimeConversionOp, backend="pandas") class PandasDatetimeConversionOp(DatetimeConversionOp, PhysicalOp): def process(self, mode: str, inputs: list): - return pd.to_datetime(inputs[0], *self.args, **self.kwargs) + _obj, _args, _kwargs = self._extract_args_and_kwargs(inputs) + return pd.to_datetime(_obj, *_args, **_kwargs) @physical_impl(of=DatetimeConversionOp, backend="polars") class PolarsDatetimeConversionOp(DatetimeConversionOp, PhysicalOp): def process(self, mode: str, inputs: list): - translated = polars_datetime_kwargs(self.args, self.kwargs) + _obj, _args, _kwargs = self._extract_args_and_kwargs(inputs) + translated = polars_datetime_kwargs(_args, _kwargs) if translated is not None: # TODO: Support already-datetime and numeric operands natively; # the Polars string namespace only accepts string input. - return inputs[0].str.to_datetime(**translated) - return _to_datetime_via_pandas(inputs[0], self.args, self.kwargs) + return _obj.str.to_datetime(**translated) + return _to_datetime_via_pandas(_obj, _args, _kwargs) # --- StringMethodOp (col.str.) ---------------------------------------- @@ -261,16 +263,12 @@ def supports(cls, op: AccessorMapOp, ctx) -> bool: return op.attr_name != ["index"] def process(self, mode: str, inputs: list): - result = inputs[0] - tmp = result + if len(self.attr_name) == 2 and self.attr_name[0] == "dt": + return polars_dt_attribute(inputs[0], self.attr_name[1]) + tmp = inputs[0] for attr in self.attr_name: - attr = self.POLARS_ATTR_NAME_MAP.get(attr, attr) - # TODO find better way to handle this - if attr == "is_month_end": - return result.dt.month_end() == result - # polars implements dt.day as a method, not an attribute; getattr - # handles both attributes and methods. tmp = getattr(tmp, attr) if len(self.attr_name) == 2: + # polars implements accessor attributes as methods. return tmp() return tmp diff --git a/tests/optimizer/logical/test_map_ops.py b/tests/optimizer/logical/test_map_ops.py index 29ec6155..0c533e27 100644 --- a/tests/optimizer/logical/test_map_ops.py +++ b/tests/optimizer/logical/test_map_ops.py @@ -1176,6 +1176,29 @@ def test_polars_fallback_keeps_scalar_result(self): result = run_op(op, "01/02/2020") self.assertEqual(pd.Timestamp("2020-02-01"), result) + def _convert_with_option(self, polars, values, args, kwargs, option): + # A plan-fed option is an OperandRef into the op's inputs: here inputs[1]. + series = pl.Series("dt", values) if polars else pd.Series(values) + with force_polars(polars): + op = DatetimeConversionOp(args=args, kwargs=kwargs) + return list(run_op(op, series, option)) + + def test_resolves_graph_fed_kwarg(self): + # On polars the kwargs translation must see the resolved format. + for polars in (False, True): + with self.subTest(polars=polars): + result = self._convert_with_option( + polars, ["2025-01-06"], (), {"format": OperandRef(1)}, "%Y-%m-%d") + self.assertEqual([pd.Timestamp("2025-01-06")], result) + + def test_resolves_graph_fed_positional(self): + for polars in (False, True): + with self.subTest(polars=polars): + result = self._convert_with_option( + polars, ["2025-01-06", "bad"], (OperandRef(1),), {}, "coerce") + self.assertEqual(pd.Timestamp("2025-01-06"), result[0]) + self.assertTrue(pd.isna(result[1])) + def test_polars_datetime_kwargs_rejects_unsupported_errors(self): self.assertIsNone( polars_datetime_kwargs((), {"errors": "ignore"})) @@ -1198,10 +1221,12 @@ def test_polars_year(self): result = self._run_polars(["2025-01-15", "2025-06-20"], ["dt", "year"]) self.assertEqual([2025, 2025], result.to_list()) - def test_polars_dayofweek(self): - # polars: Monday=1 (pandas: Monday=0) - result = self._run_polars(["2025-01-06"], ["dt", "dayofweek"]) - self.assertEqual([1], result.to_list()) + def test_polars_matches_pandas_values(self): + values = _DATES["d"] + for attr in _DT_ATTRS_PANDAS_SEMANTICS: + with self.subTest(attr=attr): + result = self._run_polars(values, ["dt", attr]) + self.assertEqual(getattr(values.dt, attr).tolist(), result.to_list()) def test_polars_is_month_end(self): result = self._run_polars(["2025-01-31", "2025-01-15"], @@ -1230,6 +1255,66 @@ def test_shared_accessor_keeps_original_edge_when_fused(self): self.assertIn(new_op, source.outputs) +_DATES = pd.DataFrame({ + "d": pd.to_datetime(["2025-01-06 00:00:00.000001001", # a Monday + "2025-01-12 12:00:00.123456789"]), # a Sunday + "s": ["2025-01-06", "2025-01-12"]}) + +# The `.dt` attributes where polars disagrees with pandas: it counts weekdays +# from Monday=1 and nanoseconds past the second (pandas: Monday=0, past the +# microsecond), and it lacks some of pandas' alias spellings. +_DT_ATTRS_PANDAS_SEMANTICS = ("dayofweek", "day_of_week", "weekday", "nanosecond", + "day_of_year", "daysinmonth") + + +def _run_plan(dag, env, polars): + """Plan ``dag`` against one backend, run it and return a pandas result.""" + from stratum.runtime._buffer_pool import BufferPool + with force_polars(polars): + ops = optimize(dag, env=env) + pool = BufferPool() + for op in ops: + pool.put(op, op.process("fit_transform", [pool.pin(inp) for inp in op.inputs])) + result = pool.pin(ops[-1]) + # Every op here has a native polars impl, so a pandas result would mean a + # fallback hid the polars kernel. + assert isinstance(result, (pl.Series, pl.DataFrame)) == polars + return result.to_pandas() if polars else result + + +@pytest.mark.parametrize("polars", [False, True], ids=["pandas", "polars"]) +@pytest.mark.parametrize("path", ["accessor", "assign", "selection"]) +@pytest.mark.parametrize("attr", _DT_ATTRS_PANDAS_SEMANTICS) +def test_dt_attribute_matches_pandas(attr, path, polars): + # The standalone accessor runs PolarsAccessorMapOp; assign and mask + # selection fold the read into a DtExpr. + values = getattr(_DATES["d"].dt, attr) + expected = {"accessor": values, "assign": _DATES.assign(v=values), + "selection": _DATES[values == values.iloc[-1]]}[path] + with skrub.config_context(eager_data_ops=False): + src = skrub.var("x", _DATES) + read = getattr(src["d"].dt, attr) + dag = {"accessor": read, "assign": src.assign(v=read), + "selection": src[read == values.iloc[-1]]}[path] + result = _run_plan(dag, {"x": _DATES}, polars) + if path == "accessor": + assert expected.tolist() == result.tolist() + else: + pd.testing.assert_frame_equal(expected.reset_index(drop=True), + result.reset_index(drop=True), + check_dtype=False) + + +@pytest.mark.parametrize("polars", [False, True], ids=["pandas", "polars"]) +def test_to_datetime_resolves_graph_fed_format(polars): + with skrub.config_context(eager_data_ops=False): + src = skrub.var("x", _DATES) + fmt = skrub.deferred(lambda: "%Y-%m-%d")() + dag = skrub.deferred(pd.to_datetime)(src["s"], format=fmt) + result = _run_plan(dag, {"x": _DATES}, polars) + assert pd.to_datetime(_DATES["s"]).tolist() == result.tolist() + + class TestStringMethodOp(unittest.TestCase): """`col.str.(...)` fuses the .str accessor and the call into one op."""