narwhals_dataframe_agnostic_function_pandas_polars_interop.py

1import pandas as pd
2import polars as pl
3import narwhals as nw
4from typing import Any
5
6# Define a dataframe-agnostic function
7def my_agnostic_function(df_any: Any) -> Any:
8    # Use narwhals.from_native to wrap the dataframe
9    df = nw.from_native(df_any)
10    
11    # Use the Narwhals API to perform operations
12    result = df.filter(nw.col("a") > 1).group_by("b").agg(nw.col("c").mean())
13    
14    # Return the result back to its original native format
15    return nw.to_native(result)
16
17# Example data
18data = {"a": [1, 2, 3], "b": [1, 1, 2], "c": [4.0, 5.0, 6.0]}
19
20# Works with pandas
21df_pd = pd.DataFrame(data)
22print("Pandas result:")
23print(my_agnostic_function(df_pd))
24
25# Works with Polars
26df_pl = pl.DataFrame(data)
27print("\nPolars result:")
28print(my_agnostic_function(df_pl))