ENH: Introduce pandas.col

doc/source/reference/general_functions.rst

@@ -71,6 +71,7 @@ Top-level evaluation
 .. autosummary::
    :toctree: api/
 
+   col
    eval
 
 Datetime formats

pandas/core/col.py

@@ -43,7 +43,7 @@ def _parse_args(df: DataFrame, *args: Any) -> tuple[Series]:
     return tuple([x(df) if isinstance(x, Expression) else x for x in args])
 
 
-def _parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[Hashable, Series]:
+def _parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[str, Any]:
     # Parse `kwargs`, evaluating any expressions we encounter.
     return {
         key: val(df) if isinstance(val, Expression) else val
@@ -260,9 +260,15 @@ def col(col_name: Hashable) -> Expression:
 
     def func(df: DataFrame) -> Series:
         if col_name not in df.columns:
+            columns_list = df.columns.tolist()
+            max_cols = 10
+            if len(columns_list) > max_cols:
+                columns_hint = columns_list[:max_cols] + ["..."]
+            else:
+                columns_hint = columns_list
             msg = (
                 f"Column '{col_name}' not found in given DataFrame.\n\n"
-                f"Hint: did you mean one of {df.columns.tolist()} instead?"
+                f"Hint: did you mean one of {columns_hint} instead?"
             )
             raise ValueError(msg)
         return df[col_name]

pandas/tests/test_col.py

@@ -32,7 +32,7 @@
         (pd.col("a") < 1, [False, False], "(col('a') < 1)"),
         (pd.col("a") <= 1, [True, False], "(col('a') <= 1)"),
         (pd.col("a") == 1, [True, False], "(col('a') == 1)"),
-        (np.log(pd.col("a")), [0.0, 0.6931471805599453], "log(col('a'))"),
+        (np.power(pd.col("a"), 2), [1, 4], "power(col('a'), 2)"),
         (np.divide(pd.col("a"), pd.col("a")), [1.0, 1.0], "divide(col('a'), col('a'))"),
     ],
 )
@@ -67,9 +67,19 @@ def test_namespaces(
 
 
 def test_invalid() -> None:
-    df = pd.DataFrame({"a": [1, 2]})
-    with pytest.raises(ValueError, match="did you mean"):
-        df.assign(c=pd.col("b").mean())
+    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+    with pytest.raises(ValueError, match=r"did you mean one of \['a', 'b'\] instead"):
+        df.assign(c=pd.col("c").mean())
+    df = pd.DataFrame({f"col_{i}": [0] for i in range(11)})
+    msg = (
+        "did you mean one of "
+        r"\['col_0', 'col_1', 'col_2', 'col_3', "
+        "'col_4', 'col_5', 'col_6', 'col_7', "
+        r"'col_8', 'col_9', '\.\.\.'\] instead"
+    )
+    ""
+    with pytest.raises(ValueError, match=msg):
+        df.assign(c=pd.col("c").mean())
 
 
 def test_custom_accessor() -> None: