perf(python): More efficient handling of *args/**kwargs (pola-rs#7026)

josemasar · Feb 21, 2023 · 10a05c8 · 10a05c8
1 parent 919f605
commit 10a05c8
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 22 deletions.
diff --git a/py-polars/polars/internals/expr/expr.py b/py-polars/polars/internals/expr/expr.py
@@ -2062,7 +2062,8 @@ def sort_by(
         if isinstance(reverse, bool):
             reverse = [reverse]
         by = selection_to_pyexpr_list(by)
-        by.extend(pli.selection_to_pyexpr_list(more_by))
+        if more_by:
+            by.extend(pli.selection_to_pyexpr_list(more_by))
         return self._from_pyexpr(self._pyexpr.sort_by(by, reverse))
 
     def take(
@@ -2866,7 +2867,8 @@ def over(self, expr: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr) -> Se
 
         """
         exprs = selection_to_pyexpr_list(expr)
-        exprs.extend(selection_to_pyexpr_list(more_exprs))
+        if more_exprs:
+            exprs.extend(selection_to_pyexpr_list(more_exprs))
         return self._from_pyexpr(self._pyexpr.over(exprs))
 
     def is_unique(self) -> Self:

diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -221,9 +221,10 @@ def col(
             dtypes = [name]
             dtypes.extend(more_names)
             return pli.wrap_expr(_dtype_cols(dtypes))
-        raise TypeError(
-            f"Invalid input for `col`. Expected `str` or `DataType`, got {type(name)!r}"
-        )
+        else:
+            raise TypeError(
+                f"Invalid input for `col`. Expected `str` or `DataType`, got {type(name)!r}"
+            )
 
     if isinstance(name, str):
         return pli.wrap_expr(pycol(name))
@@ -2715,7 +2716,8 @@ def coalesce(exprs: IntoExpr | Iterable[IntoExpr], *more_exprs: IntoExpr) -> pli
 
     """
     exprs = pli.selection_to_pyexpr_list(exprs)
-    exprs.extend(pli.selection_to_pyexpr_list(more_exprs))
+    if more_exprs:
+        exprs.extend(pli.selection_to_pyexpr_list(more_exprs))
     return pli.wrap_expr(_coalesce_exprs(exprs))
 
 

diff --git a/py-polars/polars/internals/lazyframe/frame.py b/py-polars/polars/internals/lazyframe/frame.py
@@ -906,7 +906,8 @@ def sort(
             return self._from_pyldf(self._ldf.sort(by, reverse, nulls_last))
 
         by = pli.selection_to_pyexpr_list(by)
-        by.extend(pli.selection_to_pyexpr_list(more_by))
+        if more_by:
+            by.extend(pli.selection_to_pyexpr_list(more_by))
 
         # TODO: Do this check on the Rust side
         if nulls_last and len(by) > 1:
@@ -1673,13 +1674,16 @@ def select(
         structify = bool(int(os.environ.get("POLARS_AUTO_STRUCTIFY", 0)))
 
         exprs = pli.selection_to_pyexpr_list(exprs, structify=structify)
-        exprs.extend(pli.selection_to_pyexpr_list(more_exprs, structify=structify))
-        exprs.extend(
-            pli.expr_to_lit_or_expr(
-                expr, structify=structify, name=name, str_to_lit=False
-            )._pyexpr
-            for name, expr in named_exprs.items()
-        )
+        if more_exprs:
+            exprs.extend(pli.selection_to_pyexpr_list(more_exprs, structify=structify))
+        if named_exprs:
+            exprs.extend(
+                pli.expr_to_lit_or_expr(
+                    expr, structify=structify, name=name, str_to_lit=False
+                )._pyexpr
+                for name, expr in named_exprs.items()
+            )
+
         return self._from_pyldf(self._ldf.select(exprs))
 
     def groupby(
@@ -1775,7 +1779,8 @@ def groupby(
 
         """
         exprs = selection_to_pyexpr_list(by)
-        exprs.extend(selection_to_pyexpr_list(more_by))
+        if more_by:
+            exprs.extend(selection_to_pyexpr_list(more_by))
         lgb = self._ldf.groupby(exprs, maintain_order)
         return LazyGroupBy(lgb, lazyframe_class=self.__class__)
 
@@ -2681,10 +2686,8 @@ def with_columns(
         structify = bool(int(os.environ.get("POLARS_AUTO_STRUCTIFY", 0)))
 
         exprs = pli.selection_to_pyexpr_list(exprs, structify=structify)
-
         if more_exprs:
             exprs.extend(pli.selection_to_pyexpr_list(more_exprs, structify=structify))
-
         if named_exprs:
             exprs.extend(
                 pli.expr_to_lit_or_expr(

diff --git a/py-polars/polars/internals/lazyframe/groupby.py b/py-polars/polars/internals/lazyframe/groupby.py
@@ -122,11 +122,14 @@ def agg(
             raise ValueError("Expected at least one of 'aggs' or '**named_aggs'")
 
         exprs = selection_to_pyexpr_list(aggs)
-        exprs.extend(selection_to_pyexpr_list(more_aggs))
-        exprs.extend(
-            expr_to_lit_or_expr(expr, name=name, str_to_lit=False)._pyexpr
-            for name, expr in named_aggs.items()
-        )
+        if more_aggs:
+            exprs.extend(selection_to_pyexpr_list(more_aggs))
+        if named_aggs:
+            exprs.extend(
+                expr_to_lit_or_expr(expr, name=name, str_to_lit=False)._pyexpr
+                for name, expr in named_aggs.items()
+            )
+
         return self._lazyframe_class._from_pyldf(self.lgb.agg(exprs))
 
     def apply(