Add narwhals result types

Signed-off-by: Pascal Tomecek <pascal.tomecek@cubistsystematic.com>

Author: ptomecek

ccflow/result/narwhals.py

@@ -0,0 +1,28 @@
+import narwhals.stable.v1 as nw
+from pydantic import Field
+
+from ..base import ResultBase
+from ..exttypes.narwhals import DataFrameT, FrameT
+
+__all__ = (
+    "NarwhalsFrameResult",
+    "NarwhalsDataFrameResult",
+)
+
+
+class NarwhalsFrameResult(ResultBase):
+    df: FrameT = Field(description="Narwhals DataFrame or LazyFrame")
+
+    def collect(self) -> "NarwhalsDataFrameResult":
+        """Collects the result into a NarwhalsDataFrameResult."""
+        if isinstance(self.df, nw.LazyFrame):
+            return NarwhalsDataFrameResult(df=self.df.collect(), **self.model_dump(exclude={"df", "type_"}))
+        return NarwhalsDataFrameResult(df=self.df, **self.model_dump(exclude={"df", "type_"}))
+
+
+class NarwhalsDataFrameResult(NarwhalsFrameResult):
+    df: DataFrameT = Field(description="Narwhals eager Dataframe")
+
+    def collect(self) -> "NarwhalsDataFrameResult":
+        """Collects the result into a NarwhalsDataFrameResult."""
+        return self

ccflow/tests/result/test_narwhals.py

@@ -0,0 +1,71 @@
+from typing import Annotated
+
+import narwhals.stable.v1 as nw
+import polars as pl
+import pytest
+
+from ccflow.exttypes.narwhals import (
+    DataFrameT,
+    SchemaValidator,
+)
+from ccflow.result.narwhals import NarwhalsDataFrameResult, NarwhalsFrameResult
+
+
+@pytest.fixture
+def data():
+    return {
+        "a": [1.0, 2.0, 3.0],
+        "b": [4, 5, 6],
+        "c": ["foo", "bar", "baz"],
+        "d": [0, 0, 0],
+    }
+
+
+@pytest.fixture
+def schema():
+    return {
+        "a": nw.Float64,
+        "b": nw.Int64,
+        "c": nw.String,
+        "d": nw.Float64,
+    }
+
+
+def test_narwhals_frame_result(data):
+    df = pl.DataFrame(data)
+    result = NarwhalsFrameResult(df=df)
+    assert isinstance(result.df, nw.DataFrame)
+    assert result.df.to_native() is df
+
+    df = pl.DataFrame(data).lazy()
+    result = NarwhalsFrameResult(df=df)
+    assert isinstance(result.df, nw.LazyFrame)
+    assert result.df.to_native() is df
+
+
+def test_narwhals_dataframe_result(data):
+    df = pl.DataFrame(data)
+    result = NarwhalsDataFrameResult(df=df)
+    assert isinstance(result.df, nw.DataFrame)
+    assert result.df.to_native() is df
+
+    df = pl.DataFrame(data).lazy()
+    result = NarwhalsDataFrameResult(df=df)
+    assert isinstance(result.df, nw.DataFrame)
+
+
+def test_collect(data):
+    df = pl.DataFrame(data)
+    result = NarwhalsFrameResult(df=df)
+    result2 = result.collect()
+    assert isinstance(result2, NarwhalsDataFrameResult)
+    assert isinstance(result2.df, nw.DataFrame)
+
+
+def test_custom(data, schema):
+    class MyNarwhalsResult(NarwhalsDataFrameResult):
+        df: Annotated[DataFrameT, SchemaValidator(schema, cast=True)]
+
+    df = pl.DataFrame(data)
+    result = MyNarwhalsResult(df=df)
+    assert result.df.schema["d"] == nw.Float64()

docs/wiki/Key-Features.md

@@ -134,11 +134,13 @@ A Result is an object that holds the results from a callable model. It provides
 The following table summarizes the "result" models.
 
 | Name | Path | Description |
-| :-- | :-- | :--- |
+|:---------------------------|:-------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | `GenericResult` | `ccflow.result` | A generic result (holds anything). |
 | `DictResult` | `ccflow.result` | A generic dict (key/value) result. |
 | `ArrowResult` | `ccflow.result.pyarrow` | Holds an arrow table. |
 | `ArrowDateRangeResult` | `ccflow.result.pyarrow` | Extension of `ArrowResult` for representing a table over a date range that can be divided by date, such that generation of any sub-range of dates gives the same results as the original table filtered for those dates. |
+| `NarwhalsResult` | `ccflow.result.narwhals` | Holds a narwhals `DataFrame` or `LazyFrame`. |
+| `NarwhalsDataFrameResult` | `ccflow.result.narwhals` | Holds a narwhals eager `DataFrame`. |
 | `NumpyResult` | `ccflow.result.numpy` | Holds a numpy array. |
 | `PandasResult` | `ccflow.result.pandas` | Holds a pandas dataframe. |
 | `XArrayResult` | `ccflow.result.xarray` | Holds an xarray. |