Repository: narwhals-dev/narwhals

Issue:
[Bug]: Pandas Dataframe as native does not work with .shift().over()
### Describe the bug

The below code works if the native dataframe is polars. However, getting an error message using pandas.

### Steps or code to reproduce the bug

```
import polars as pl
import narwhals as nw
from narwhals.typing import FrameT

@nw.narwhalify
def create_lag(df: FrameT) -> FrameT:
    return df.with_columns(
        nw.col("a")
        .shift(1)
        .over("grp")
        .alias("lag_a")
    )

df = pl.DataFrame(
    {
        "a": [1, 2, 3, 4, 5],
        "grp": [1, 1, 2, 2, 2]
    }
)

l = create_lag(df)

df = df.to_pandas()
l2 = create_lag(df)

```

### Expected results

Both the polars and pandas dataframe should add a new column lag_a.

### Actual results

Getting the following error when pandas is the native dataframe

```
narwhals.exceptions.ColumnNotFoundError: The following columns were not found: ['grp']

Hint: Did you mean one of these columns: ['index', 'a']?

```

### Please run narwhals.show_version() and enter the output below.

```shell
Python dependencies:
     narwhals: 1.17.0
       pandas: 2.2.2
       polars: 1.1.0
         cudf: 
        modin: 
      pyarrow: 16.0.0
        numpy: 1.26.4
```


### Relevant log output

_No response_

Discussion:

Patch:
diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml
index fea1d6b0..e19d9a7f 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.yml
+++ b/.github/ISSUE_TEMPLATE/bug_report.yml
@@ -58,4 +58,4 @@ body:
   attributes:
     label: Relevant log output
     description: Please copy and paste any relevant log output. This will be automatically formatted into code, so no need for backticks.
-    render: shell
\ No newline at end of file
+    render: shell
diff --git a/.github/ISSUE_TEMPLATE/doc_issue.yml b/.github/ISSUE_TEMPLATE/doc_issue.yml
index 954320c9..bc28e287 100644
--- a/.github/ISSUE_TEMPLATE/doc_issue.yml
+++ b/.github/ISSUE_TEMPLATE/doc_issue.yml
@@ -32,4 +32,4 @@ body:
   attributes:
     value: > 
         ### If you'd be interested in opening a pull request to fix this, please let us know!
-      
\ No newline at end of file
+      
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
index b7a84a59..2d0556ed 100644
--- a/.github/dependabot.yml
+++ b/.github/dependabot.yml
@@ -9,4 +9,4 @@ updates:
       interval: "monthly"
     commit-message:
       prefix: "skip changelog" # So this PR will not be added to release-drafter
-      include: "scope" # List of the updated dependencies in the commit will be added
\ No newline at end of file
+      include: "scope" # List of the updated dependencies in the commit will be added
diff --git a/.github/workflows/release-drafter.yml b/.github/workflows/release-drafter.yml
index 9b37e420..0aee4d21 100644
--- a/.github/workflows/release-drafter.yml
+++ b/.github/workflows/release-drafter.yml
@@ -35,4 +35,4 @@ jobs:
            config-name: release-drafter.yml
         #   disable-autolabeler: true
         env:
-          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
\ No newline at end of file
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 9600d697..fb149f30 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -95,4 +95,7 @@ repos:
   rev: v5.0.0
   hooks:
   -   id: name-tests-test
-      exclude: ^tests/utils\.py
\ No newline at end of file
+      exclude: ^tests/utils\.py
+  -   id: no-commit-to-branch
+  -   id: end-of-file-fixer
+      exclude: .svg$
diff --git a/docs/css/extra.css b/docs/css/extra.css
index 2d2769e0..6d19bea3 100644
--- a/docs/css/extra.css
+++ b/docs/css/extra.css
@@ -1,4 +1,4 @@
 .md-typeset ol li,
 .md-typeset ul li {
   margin-bottom: 0.1em !important;
-}
\ No newline at end of file
+}
diff --git a/docs/javascripts/extra.js b/docs/javascripts/extra.js
index 4a50c954..718f6033 100644
--- a/docs/javascripts/extra.js
+++ b/docs/javascripts/extra.js
@@ -64,4 +64,4 @@ function cleanupClipboardText(targetSelector) {
   document$.subscribe(function () {
     setCopyText();
   });
-  
\ No newline at end of file
+  
diff --git a/docs/javascripts/katex.js b/docs/javascripts/katex.js
index 3828300a..eb5813a0 100644
--- a/docs/javascripts/katex.js
+++ b/docs/javascripts/katex.js
@@ -7,4 +7,4 @@ document$.subscribe(({ body }) => {
         { left: "\\[", right: "\\]", display: true }
       ],
     })
-  })
\ No newline at end of file
+  })
diff --git a/mkdocs.yml b/mkdocs.yml
index bd24010d..5bbbd12d 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -131,4 +131,4 @@ extra_javascript:
 
 extra_css:
   - https://unpkg.com/katex@0/dist/katex.min.css
-  - css/extra.css
\ No newline at end of file
+  - css/extra.css
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
index f8f5c992..2d683a15 100644
--- a/narwhals/_arrow/group_by.py
+++ b/narwhals/_arrow/group_by.py
@@ -21,11 +21,16 @@ if TYPE_CHECKING:
     from narwhals._arrow.typing import IntoArrowExpr
 
 POLARS_TO_ARROW_AGGREGATIONS = {
-    "len": "count",
+    "sum": "sum",
+    "mean": "mean",
     "median": "approximate_median",
-    "n_unique": "count_distinct",
+    "max": "max",
+    "min": "min",
     "std": "stddev",
     "var": "variance",  # currently unused, we don't have `var` yet
+    "len": "count",
+    "n_unique": "count_distinct",
+    "count": "count",
 }
 
 
@@ -36,7 +41,7 @@ def get_function_name_option(
     import pyarrow.compute as pc
 
     function_name_to_options = {
-        "count": pc.CountOptions(mode="all"),
+        "count": pc.CountOptions(mode="only_valid"),
         "count_distinct": pc.CountOptions(mode="all"),
         "stddev": pc.VarianceOptions(ddof=1),
         "variance": pc.VarianceOptions(ddof=1),
@@ -131,7 +136,11 @@ def agg_arrow(
 
     all_simple_aggs = True
     for expr in exprs:
-        if not is_simple_aggregation(expr):
+        if not (
+            is_simple_aggregation(expr)
+            and remove_prefix(expr._function_name, "col->")
+            in POLARS_TO_ARROW_AGGREGATIONS
+        ):
             all_simple_aggs = False
             break
 
@@ -185,7 +194,7 @@ def agg_arrow(
         return from_dataframe(result_simple)
 
     msg = (
-        "Non-trivial complex found.\n\n"
+        "Non-trivial complex aggregation found.\n\n"
         "Hint: you were probably trying to apply a non-elementary aggregation with a "
         "pyarrow table.\n"
         "Please rewrite your query such that group-by aggregations "
diff --git a/narwhals/_dask/group_by.py b/narwhals/_dask/group_by.py
index 3bc21efb..7165b8c7 100644
--- a/narwhals/_dask/group_by.py
+++ b/narwhals/_dask/group_by.py
@@ -35,8 +35,16 @@ def n_unique() -> dd.Aggregation:
 
 
 POLARS_TO_DASK_AGGREGATIONS = {
+    "sum": "sum",
+    "mean": "mean",
+    "median": "median",
+    "max": "max",
+    "min": "min",
+    "std": "std",
+    "var": "var",
     "len": "size",
     "n_unique": n_unique,
+    "count": "count",
 }
 
 
@@ -108,7 +116,10 @@ def agg_dask(
 
     all_simple_aggs = True
     for expr in exprs:
-        if not is_simple_aggregation(expr):
+        if not (
+            is_simple_aggregation(expr)
+            and remove_prefix(expr._function_name, "col->") in POLARS_TO_DASK_AGGREGATIONS
+        ):
             all_simple_aggs = False
             break
 
@@ -143,15 +154,11 @@ def agg_dask(
 
             for root_name, output_name in zip(expr._root_names, expr._output_names):
                 simple_aggregations[output_name] = (root_name, function_name)
-        try:
-            result_simple = grouped.agg(**simple_aggregations)
-        except ValueError as exc:
-            msg = "Failed to aggregated - does your aggregation function return a scalar?"
-            raise RuntimeError(msg) from exc
+        result_simple = grouped.agg(**simple_aggregations)
         return from_dataframe(result_simple.reset_index())
 
     msg = (
-        "Non-trivial complex found.\n\n"
+        "Non-trivial complex aggregation found.\n\n"
         "Hint: you were probably trying to apply a non-elementary aggregation with a "
         "dask dataframe.\n"
         "Please rewrite your query such that group-by aggregations "
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 0ddad326..ce90db0c 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -40,7 +40,6 @@ class PandasLikeExpr:
         self._depth = depth
         self._function_name = function_name
         self._root_names = root_names
-        self._depth = depth
         self._output_names = output_names
         self._implementation = implementation
         self._backend_version = backend_version
@@ -53,6 +52,7 @@ class PandasLikeExpr:
             f"function_name={self._function_name}, "
             f"root_names={self._root_names}, "
             f"output_names={self._output_names}"
+            ")"
         )
 
     def __narwhals_namespace__(self) -> PandasLikeNamespace:
diff --git a/narwhals/_pandas_like/group_by.py b/narwhals/_pandas_like/group_by.py
index 66e90c83..5735953b 100644
--- a/narwhals/_pandas_like/group_by.py
+++ b/narwhals/_pandas_like/group_by.py
@@ -14,6 +14,7 @@ from narwhals._pandas_like.utils import native_series_from_iterable
 from narwhals._pandas_like.utils import rename
 from narwhals._pandas_like.utils import select_columns_by_name
 from narwhals.utils import Implementation
+from narwhals.utils import find_stacklevel
 from narwhals.utils import remove_prefix
 from narwhals.utils import tupleify
 
@@ -23,8 +24,16 @@ if TYPE_CHECKING:
     from narwhals._pandas_like.typing import IntoPandasLikeExpr
 
 POLARS_TO_PANDAS_AGGREGATIONS = {
+    "sum": "sum",
+    "mean": "mean",
+    "median": "median",
+    "max": "max",
+    "min": "min",
+    "std": "std",
+    "var": "var",
     "len": "size",
     "n_unique": "nunique",
+    "count": "count",
 }
 
 
@@ -144,7 +153,11 @@ def agg_pandas(  # noqa: PLR0915
     """
     all_aggs_are_simple = True
     for expr in exprs:
-        if not is_simple_aggregation(expr):
+        if not (
+            is_simple_aggregation(expr)
+            and remove_prefix(expr._function_name, "col->")
+            in POLARS_TO_PANDAS_AGGREGATIONS
+        ):
             all_aggs_are_simple = False
             break
 
@@ -193,11 +206,7 @@ def agg_pandas(  # noqa: PLR0915
             simple_aggs[named_agg[0]].append(named_agg[1])
             name_mapping[f"{named_agg[0]}_{named_agg[1]}"] = output_name
         if simple_aggs:
-            try:
-                result_simple_aggs = grouped.agg(simple_aggs)
-            except AttributeError as exc:
-                msg = "Failed to aggregated - does your aggregation function return a scalar?"
-                raise RuntimeError(msg) from exc
+            result_simple_aggs = grouped.agg(simple_aggs)
             result_simple_aggs.columns = [
                 f"{a}_{b}" for a, b in result_simple_aggs.columns
             ]
@@ -264,7 +273,7 @@ def agg_pandas(  # noqa: PLR0915
         "pandas API. If you can, please rewrite your query such that group-by aggregations "
         "are simple (e.g. mean, std, min, max, ...).",
         UserWarning,
-        stacklevel=2,
+        stacklevel=find_stacklevel(),
     )
 
     def func(df: Any) -> Any:
@@ -272,6 +281,9 @@ def agg_pandas(  # noqa: PLR0915
         out_names = []
         for expr in exprs:
             results_keys = expr._call(from_dataframe(df))
+            if not all(len(x) == 1 for x in results_keys):
+                msg = f"Aggregation '{expr._function_name}' failed to aggregate - does your aggregation function return a scalar?"
+                raise ValueError(msg)
             for result_keys in results_keys:
                 out_group.append(result_keys._native_series.iloc[0])
                 out_names.append(result_keys.name)
diff --git a/narwhals/_spark_like/group_by.py b/narwhals/_spark_like/group_by.py
index 3b7ad78f..a3d557c0 100644
--- a/narwhals/_spark_like/group_by.py
+++ b/narwhals/_spark_like/group_by.py
@@ -91,7 +91,7 @@ def agg_pyspark(
     for expr in exprs:
         if not is_simple_aggregation(expr):  # pragma: no cover
             msg = (
-                "Non-trivial complex found.\n\n"
+                "Non-trivial complex aggregation found.\n\n"
                 "Hint: you were probably trying to apply a non-elementary aggregation with a "
                 "dask dataframe.\n"
                 "Please rewrite your query such that group-by aggregations "


Test patch:
diff --git a/tests/group_by_test.py b/tests/group_by_test.py
index 5fb14d3c..e30db8b1 100644
--- a/tests/group_by_test.py
+++ b/tests/group_by_test.py
@@ -44,10 +44,10 @@ def test_invalid_group_by_dask() -> None:
 
     df_dask = dd.from_pandas(df_pandas)
 
-    with pytest.raises(ValueError, match=r"Non-trivial complex found"):
+    with pytest.raises(ValueError, match=r"Non-trivial complex aggregation found"):
         nw.from_native(df_dask).group_by("a").agg(nw.col("b").mean().min())
 
-    with pytest.raises(RuntimeError, match="does your"):
+    with pytest.raises(ValueError, match="Non-trivial complex aggregation"):
         nw.from_native(df_dask).group_by("a").agg(nw.col("b"))
 
     with pytest.raises(
@@ -56,9 +56,10 @@ def test_invalid_group_by_dask() -> None:
         nw.from_native(df_dask).group_by("a").agg(nw.all().mean())
 
 
+@pytest.mark.filterwarnings("ignore:Found complex group-by expression:UserWarning")
 def test_invalid_group_by() -> None:
     df = nw.from_native(df_pandas)
-    with pytest.raises(RuntimeError, match="does your"):
+    with pytest.raises(ValueError, match="does your"):
         df.group_by("a").agg(nw.col("b"))
     with pytest.raises(
         ValueError, match=r"Anonymous expressions are not supported in group_by\.agg"
@@ -68,7 +69,7 @@ def test_invalid_group_by() -> None:
         ValueError, match=r"Anonymous expressions are not supported in group_by\.agg"
     ):
         nw.from_native(pa.table({"a": [1, 2, 3]})).group_by("a").agg(nw.all().mean())
-    with pytest.raises(ValueError, match=r"Non-trivial complex found"):
+    with pytest.raises(ValueError, match=r"Non-trivial complex aggregation found"):
         nw.from_native(pa.table({"a": [1, 2, 3]})).group_by("a").agg(
             nw.col("b").mean().min()
         )
@@ -340,3 +341,27 @@ def test_group_by_categorical(
         .sort("x")
     )
     assert_equal_data(result, data)
+
+
+@pytest.mark.filterwarnings("ignore:Found complex group-by expression:UserWarning")
+def test_group_by_shift_raises(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "polars" in str(constructor):
+        # Polars supports all kinds of crazy group-by aggregations, so
+        # we don't check that it errors here.
+        request.applymarker(pytest.mark.xfail)
+    df_native = {"a": [1, 2, 3], "b": [1, 1, 2]}
+    df = nw.from_native(constructor(df_native))
+    with pytest.raises(
+        ValueError, match=".*(failed to aggregate|Non-trivial complex aggregation found)"
+    ):
+        df.group_by("b").agg(nw.col("a").shift(1))
+
+
+def test_group_by_count(constructor: Constructor) -> None:
+    data = {"a": [1, 1, 1, 2], "b": [1, None, 2, 3]}
+    df = nw.from_native(constructor(data))
+    result = df.group_by("a").agg(nw.col("b").count()).sort("a")
+    expected = {"a": [1, 2], "b": [2, 1]}
+    assert_equal_data(result, expected)