Allow selection with Series in col() (#1602)

pola-rs · Oct 26, 2021 · 4509e76 · 4509e76
1 parent e1a60a9
commit 4509e76
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/py-polars/polars/lazy/functions.py b/py-polars/polars/lazy/functions.py
@@ -157,6 +157,8 @@ def col(
     ╰───────────┴─────╯
 
     """
+    if isinstance(name, pl.Series):
+        name = name.to_list()
 
     if isclass(name) and issubclass(name, DataType):  # type: ignore
         name = [name]  # type: ignore

diff --git a/py-polars/tests/test_lazy.py b/py-polars/tests/test_lazy.py
@@ -362,6 +362,13 @@ def test_exclude_selection():
     assert df.select([pl.exclude("a")]).columns == ["b", "c"]
 
 
+def test_col_series_selection():
+    df = pl.DataFrame({"a": [1], "b": [1], "c": [1]}).lazy()
+    srs = pl.Series(["b", "c"])
+
+    assert df.select(pl.col(srs)).columns == ["b", "c"]
+
+
 def test_literal_projection():
     df = pl.DataFrame({"a": [1, 2]})
     assert df.select([True]).dtypes == [pl.Boolean]