Skip to content

Commit 5b456ea

Browse files
committed
Revert "refactor: enhance expression handling by introducing _to_raw_expr function and updating _ensure_expr"
This reverts commit af3f475987a2fe089b31a4cc90483b660dc087f7.
1 parent 8349ca3 commit 5b456ea

2 files changed

Lines changed: 30 additions & 34 deletions

File tree

python/datafusion/dataframe.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
EXPR_TYPE_ERROR,
4545
Expr,
4646
SortKey,
47-
_to_raw_expr,
4847
expr_list_to_raw_expr_list,
4948
sort_list_to_raw_sort_list,
5049
)
@@ -293,19 +292,21 @@ def __init__(
293292
self.bloom_filter_ndv = bloom_filter_ndv
294293

295294

296-
def _ensure_expr(value: Expr | str) -> expr_internal.Expr:
295+
def _ensure_expr(value: Expr) -> expr_internal.Expr:
297296
"""Return the internal expression or raise ``TypeError`` if invalid.
298297
299298
Args:
300-
value: Candidate expression or column name.
299+
value: Candidate expression.
301300
302301
Returns:
303302
The internal expression representation.
304303
305304
Raises:
306-
TypeError: If ``value`` is not an instance of :class:`Expr` or ``str``.
305+
TypeError: If ``value`` is not an instance of :class:`Expr`.
307306
"""
308-
return _to_raw_expr(value)
307+
if not isinstance(value, Expr):
308+
raise TypeError(EXPR_TYPE_ERROR)
309+
return value.expr
309310

310311

311312
class DataFrame:

python/datafusion/expr.py

Lines changed: 24 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -222,29 +222,6 @@
222222
]
223223

224224

225-
def _to_raw_expr(value: Expr | str) -> expr_internal.Expr:
226-
"""Convert a Python expression or column name to its raw variant.
227-
228-
Args:
229-
value: Candidate expression or column name.
230-
231-
Returns:
232-
The internal :class:`~datafusion._internal.expr.Expr` representation.
233-
234-
Raises:
235-
TypeError: If ``value`` is neither an :class:`Expr` nor ``str``.
236-
"""
237-
if isinstance(value, str):
238-
return Expr.column(value).expr
239-
if isinstance(value, Expr):
240-
return value.expr
241-
error = (
242-
"Expected Expr or column name, found:"
243-
f" {type(value).__name__}. {EXPR_TYPE_ERROR}."
244-
)
245-
raise TypeError(error)
246-
247-
248225
def expr_list_to_raw_expr_list(
249226
expr_list: Optional[Sequence[Expr | str] | Expr | str],
250227
) -> Optional[list[expr_internal.Expr]]:
@@ -253,11 +230,23 @@ def expr_list_to_raw_expr_list(
253230
expr_list = [expr_list]
254231
if expr_list is None:
255232
return None
256-
return [_to_raw_expr(e) for e in expr_list]
233+
raw_exprs: list[expr_internal.Expr] = []
234+
for e in expr_list:
235+
if isinstance(e, str):
236+
raw_exprs.append(Expr.column(e).expr)
237+
elif isinstance(e, Expr):
238+
raw_exprs.append(e.expr)
239+
else:
240+
error = (
241+
"Expected Expr or column name, found:"
242+
f" {type(e).__name__}. {EXPR_TYPE_ERROR}."
243+
)
244+
raise TypeError(error)
245+
return raw_exprs
257246

258247

259248
def sort_or_default(e: Expr | SortExpr) -> expr_internal.SortExpr:
260-
"""Return a :class:`SortExpr`, defaulting attributes when necessary."""
249+
"""Helper function to return a default Sort if an Expr is provided."""
261250
if isinstance(e, SortExpr):
262251
return e.raw_sort
263252
return SortExpr(e, ascending=True, nulls_first=True).raw_sort
@@ -273,11 +262,17 @@ def sort_list_to_raw_sort_list(
273262
return None
274263
raw_sort_list = []
275264
for item in sort_list:
276-
if isinstance(item, SortExpr):
277-
raw_sort_list.append(sort_or_default(item))
265+
if isinstance(item, str):
266+
expr_obj = Expr.column(item)
267+
elif isinstance(item, (Expr, SortExpr)):
268+
expr_obj = item
278269
else:
279-
raw_expr = _to_raw_expr(item) # may raise ``TypeError``
280-
raw_sort_list.append(sort_or_default(Expr(raw_expr)))
270+
error = (
271+
"Expected Expr or column name, found:"
272+
f" {type(item).__name__}. {EXPR_TYPE_ERROR}."
273+
)
274+
raise TypeError(error)
275+
raw_sort_list.append(sort_or_default(expr_obj))
281276
return raw_sort_list
282277

283278

0 commit comments

Comments
 (0)