|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | import operator as op |
| 4 | +from functools import partial |
4 | 5 | from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast |
5 | 6 |
|
6 | 7 | import daft.functions as F |
@@ -371,6 +372,33 @@ def _with_alias_output_names(self, func: AliasNames | None, /) -> DaftExpr: |
371 | 372 | version=self._version, |
372 | 373 | ) |
373 | 374 |
|
| 375 | + def _reuse_series( |
| 376 | + self, method_name: str, *, returns_scalar: bool = False, **kwargs: Any |
| 377 | + ) -> DaftExpr: |
| 378 | + """Reuse Series implementation for expression. |
| 379 | +
|
| 380 | + If Series.foo is already defined, and we'd like Expr.foo to be the same, we can |
| 381 | + leverage this method to do that for us. |
| 382 | +
|
| 383 | + Arguments: |
| 384 | + method_name: name of method. |
| 385 | + returns_scalar: whether the Series version returns a scalar. In this case, |
| 386 | + the expression version should return a 1-row Series. |
| 387 | + kwargs: keyword arguments to pass to function. |
| 388 | + """ |
| 389 | + func = partial( |
| 390 | + self._reuse_series_inner, |
| 391 | + method_name=method_name, |
| 392 | + returns_scalar=returns_scalar, |
| 393 | + **kwargs, |
| 394 | + ) |
| 395 | + return self._from_callable( |
| 396 | + func, |
| 397 | + evaluate_output_names=self._evaluate_output_names, |
| 398 | + alias_output_names=self._alias_output_names, |
| 399 | + context=self, |
| 400 | + ) |
| 401 | + |
374 | 402 | def __and__(self, other: DaftExpr) -> DaftExpr: |
375 | 403 | return self._with_binary(lambda expr, other: (expr & other), other=other) |
376 | 404 |
|
|
0 commit comments