Repository: narwhals-dev/narwhals

Issue:
[Bug]: narwhals.DataFrame.with_columns changes types for single row dataframes
### Describe the bug

Example below of single row pandas df with categorical column 'a' being put through narwhals and returned with 'a' converted to float type, ideal behaviour would be for this to preserve type

### Steps or code to reproduce the bug

```
import polars as pl
import numpy as np
import pandas as pd
import narwhals as nw

# setup single row with a null value (this will be inferred as float later)
pandas_df=pd.DataFrame({'a': [np.NaN]},)
# convert a to category type
pandas_df['a']=pd.Categorical(pandas_df['a'], categories=['dummy_value'])

narwhals_df=nw.from_native(pandas_df)
# reassign a to itself (appreciate this step looks odd, but has been stripped back to be minimal,
# in practice we would be doing something more like `narwhals_df.with_columns(nw.col('a').fill_null('dummy_value')`
# and hitting same issue
narwhals_df=narwhals_df.with_columns(nw.col('a'))
print(narwhals_df.to_native().dtypes)

# Outputs:
# a    float64
# dtype: object

```

### Expected results

return column with preserved category type

### Actual results

returned column is converted to float type

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

```shell
System:
    python: 3.9.18 (main, Sep 11 2023, 14:09:26) [MSC v.1916 64 bit (AMD64)]
executable: c:\Miniconda\python.exe
   machine: Windows-10-10.0.17763-SP0

Python dependencies:
     narwhals: 1.9.2
       pandas: 2.0.3
       polars: 1.9.0
         cudf: 
        modin: 
      pyarrow: 17.0.0
        numpy: 1.24.4
```


### Relevant log output

_No response_

Discussion:

Patch:
diff --git a/docs/api-reference/narwhals.md b/docs/api-reference/narwhals.md
index 044b20e0..b8ec2d79 100644
--- a/docs/api-reference/narwhals.md
+++ b/docs/api-reference/narwhals.md
@@ -14,6 +14,7 @@ Here are the top-level functions available in Narwhals.
         - concat_str
         - from_dict
         - from_native
+        - from_arrow
         - get_level
         - get_native_namespace
         - is_ordered_categorical
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index 3a327aad..db8bc842 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -45,6 +45,7 @@ from narwhals.expr import sum
 from narwhals.expr import sum_horizontal
 from narwhals.expr import when
 from narwhals.functions import concat
+from narwhals.functions import from_arrow
 from narwhals.functions import from_dict
 from narwhals.functions import get_level
 from narwhals.functions import new_series
@@ -69,6 +70,7 @@ __all__ = [
     "selectors",
     "concat",
     "from_dict",
+    "from_arrow",
     "get_level",
     "new_series",
     "to_native",
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 2fe53b22..9dc9f20f 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -619,9 +619,7 @@ class PandasLikeSeries:
 
     def zip_with(self: Self, mask: Any, other: Any) -> PandasLikeSeries:
         ser = self._native_series
-        mask = validate_column_comparand(
-            ser.index, mask, treat_length_one_as_scalar=False
-        )
+        mask = validate_column_comparand(ser.index, mask)
         other = validate_column_comparand(ser.index, other)
         res = ser.where(mask, other)
         return self._from_native_series(res)
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 0773764d..5267dd07 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -32,9 +32,7 @@ PANDAS_LIKE_IMPLEMENTATION = {
 }
 
 
-def validate_column_comparand(
-    index: Any, other: Any, *, treat_length_one_as_scalar: bool = True
-) -> Any:
+def validate_column_comparand(index: Any, other: Any) -> Any:
     """Validate RHS of binary operation.
 
     If the comparison isn't supported, return `NotImplemented` so that the
@@ -55,9 +53,10 @@ def validate_column_comparand(
     if isinstance(other, PandasLikeDataFrame):
         return NotImplemented
     if isinstance(other, PandasLikeSeries):
-        if other.len() == 1 and treat_length_one_as_scalar:
+        if other.len() == 1:
             # broadcast
-            return other.item()
+            s = other._native_series
+            return s.__class__(s.iloc[0], index=index, dtype=s.dtype)
         if other._native_series.index is not index:
             return set_axis(
                 other._native_series,
@@ -83,7 +82,8 @@ def validate_dataframe_comparand(index: Any, other: Any) -> Any:
     if isinstance(other, PandasLikeSeries):
         if other.len() == 1:
             # broadcast
-            return other._native_series.iloc[0]
+            s = other._native_series
+            return s.__class__(s.iloc[0], index=index, dtype=s.dtype)
         if other._native_series.index is not index:
             return set_axis(
                 other._native_series,
diff --git a/narwhals/functions.py b/narwhals/functions.py
index b84dcb17..395da97c 100644
--- a/narwhals/functions.py
+++ b/narwhals/functions.py
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
 from typing import Literal
+from typing import Protocol
 from typing import TypeVar
 from typing import Union
 
@@ -21,6 +22,7 @@ from narwhals.utils import validate_laziness
 # The rest of the annotations seem to work fine with this anyway
 FrameT = TypeVar("FrameT", bound=Union[DataFrame, LazyFrame])  # type: ignore[type-arg]
 
+
 if TYPE_CHECKING:
     from types import ModuleType
 
@@ -29,6 +31,11 @@ if TYPE_CHECKING:
     from narwhals.series import Series
     from narwhals.typing import DTypes
 
+    class ArrowStreamExportable(Protocol):
+        def __arrow_c_stream__(
+            self, requested_schema: object | None = None
+        ) -> object: ...
+
 
 def concat(
     items: Iterable[FrameT],
@@ -406,6 +413,100 @@ def _from_dict_impl(
     return from_native(native_frame, eager_only=True)
 
 
+def from_arrow(
+    native_frame: ArrowStreamExportable, *, native_namespace: ModuleType
+) -> DataFrame[Any]:
+    """
+    Construct a DataFrame from an object which supports the PyCapsule Interface.
+
+    Arguments:
+        native_frame: Object which implements `__arrow_c_stream__`.
+        native_namespace: The native library to use for DataFrame creation.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import narwhals as nw
+        >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+
+        Let's define a dataframe-agnostic function which creates a PyArrow
+        Table.
+
+        >>> @nw.narwhalify
+        ... def func(df):
+        ...     return nw.from_arrow(df, native_namespace=pa)
+
+        Let's see what happens when passing pandas / Polars input:
+
+        >>> func(pd.DataFrame(data))  # doctest: +SKIP
+        pyarrow.Table
+        a: int64
+        b: int64
+        ----
+        a: [[1,2,3]]
+        b: [[4,5,6]]
+        >>> func(pl.DataFrame(data))  # doctest: +SKIP
+        pyarrow.Table
+        a: int64
+        b: int64
+        ----
+        a: [[1,2,3]]
+        b: [[4,5,6]]
+    """
+    if not hasattr(native_frame, "__arrow_c_stream__"):
+        msg = f"Given object of type {type(native_frame)} does not support PyCapsule interface"
+        raise TypeError(msg)
+    implementation = Implementation.from_native_namespace(native_namespace)
+
+    if implementation is Implementation.POLARS and parse_version(
+        native_namespace.__version__
+    ) >= (1, 3):
+        native_frame = native_namespace.DataFrame(native_frame)
+    elif implementation in {
+        Implementation.PANDAS,
+        Implementation.MODIN,
+        Implementation.CUDF,
+        Implementation.POLARS,
+    }:
+        # These don't (yet?) support the PyCapsule Interface for import
+        # so we go via PyArrow
+        try:
+            import pyarrow as pa  # ignore-banned-import
+        except ModuleNotFoundError as exc:  # pragma: no cover
+            msg = f"PyArrow>=14.0.0 is required for `from_arrow` for object of type {native_namespace}"
+            raise ModuleNotFoundError(msg) from exc
+        if parse_version(pa.__version__) < (14, 0):  # pragma: no cover
+            msg = f"PyArrow>=14.0.0 is required for `from_arrow` for object of type {native_namespace}"
+            raise ModuleNotFoundError(msg) from None
+
+        tbl = pa.table(native_frame)
+        if implementation is Implementation.PANDAS:
+            native_frame = tbl.to_pandas()
+        elif implementation is Implementation.MODIN:  # pragma: no cover
+            from modin.pandas.utils import from_arrow
+
+            native_frame = from_arrow(tbl)
+        elif implementation is Implementation.CUDF:  # pragma: no cover
+            native_frame = native_namespace.DataFrame.from_arrow(tbl)
+        elif implementation is Implementation.POLARS:  # pragma: no cover
+            native_frame = native_namespace.from_arrow(tbl)
+        else:  # pragma: no cover
+            msg = "congratulations, you entered unrecheable code - please report a bug"
+            raise AssertionError(msg)
+    elif implementation is Implementation.PYARROW:
+        native_frame = native_namespace.table(native_frame)
+    else:  # pragma: no cover
+        try:
+            # implementation is UNKNOWN, Narwhals extension using this feature should
+            # implement PyCapsule support
+            native_frame = native_namespace.DataFrame(native_frame)
+        except AttributeError as e:
+            msg = "Unknown namespace is expected to implement `DataFrame` class which accepts object which supports PyCapsule Interface."
+            raise AttributeError(msg) from e
+    return from_native(native_frame, eager_only=True)
+
+
 def _get_sys_info() -> dict[str, str]:
     """System information
 
diff --git a/narwhals/stable/v1/__init__.py b/narwhals/stable/v1/__init__.py
index c500cdab..75da2a42 100644
--- a/narwhals/stable/v1/__init__.py
+++ b/narwhals/stable/v1/__init__.py
@@ -21,6 +21,7 @@ from narwhals.expr import When as NwWhen
 from narwhals.expr import when as nw_when
 from narwhals.functions import _from_dict_impl
 from narwhals.functions import _new_series_impl
+from narwhals.functions import from_arrow as nw_from_arrow
 from narwhals.functions import show_versions
 from narwhals.schema import Schema as NwSchema
 from narwhals.series import Series as NwSeries
@@ -66,6 +67,7 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
     from narwhals.dtypes import DType
+    from narwhals.functions import ArrowStreamExportable
     from narwhals.typing import IntoExpr
 
 T = TypeVar("T")
@@ -2183,6 +2185,52 @@ def new_series(
     )
 
 
+def from_arrow(
+    native_frame: ArrowStreamExportable, *, native_namespace: ModuleType
+) -> DataFrame[Any]:
+    """
+    Construct a DataFrame from an object which supports the PyCapsule Interface.
+
+    Arguments:
+        native_frame: Object which implements `__arrow_c_stream__`.
+        native_namespace: The native library to use for DataFrame creation.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import narwhals.stable.v1 as nw
+        >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+
+        Let's define a dataframe-agnostic function which creates a PyArrow
+        Table.
+
+        >>> @nw.narwhalify
+        ... def func(df):
+        ...     return nw.from_arrow(df, native_namespace=pa)
+
+        Let's see what happens when passing pandas / Polars input:
+
+        >>> func(pd.DataFrame(data))  # doctest: +SKIP
+        pyarrow.Table
+        a: int64
+        b: int64
+        ----
+        a: [[1,2,3]]
+        b: [[4,5,6]]
+        >>> func(pl.DataFrame(data))  # doctest: +SKIP
+        pyarrow.Table
+        a: int64
+        b: int64
+        ----
+        a: [[1,2,3]]
+        b: [[4,5,6]]
+    """
+    return _stableify(  # type: ignore[no-any-return]
+        nw_from_arrow(native_frame, native_namespace=native_namespace)
+    )
+
+
 def from_dict(
     data: dict[str, Any],
     schema: dict[str, DType] | Schema | None = None,
@@ -2307,5 +2355,6 @@ __all__ = [
     "show_versions",
     "Schema",
     "from_dict",
+    "from_arrow",
     "new_series",
 ]


Test patch:
diff --git a/tests/frame/with_columns_test.py b/tests/frame/with_columns_test.py
index 44bcd39a..8c949cc5 100644
--- a/tests/frame/with_columns_test.py
+++ b/tests/frame/with_columns_test.py
@@ -1,7 +1,10 @@
 import numpy as np
 import pandas as pd
+import pyarrow as pa
+import pytest
 
 import narwhals.stable.v1 as nw
+from narwhals.utils import parse_version
 from tests.utils import Constructor
 from tests.utils import compare_dicts
 
@@ -40,3 +43,14 @@ def test_with_columns_order_single_row(constructor: Constructor) -> None:
     assert result.collect_schema().names() == ["a", "b", "z", "d"]
     expected = {"a": [2], "b": [4], "z": [7.0], "d": [0]}
     compare_dicts(result, expected)
+
+
+def test_with_columns_dtypes_single_row(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if "pyarrow_table" in str(constructor) and parse_version(pa.__version__) < (15,):
+        request.applymarker(pytest.mark.xfail)
+    data = {"a": ["foo"]}
+    df = nw.from_native(constructor(data)).with_columns(nw.col("a").cast(nw.Categorical))
+    result = df.with_columns(nw.col("a"))
+    assert result.collect_schema() == {"a": nw.Categorical}
diff --git a/tests/from_pycapsule_test.py b/tests/from_pycapsule_test.py
new file mode 100644
index 00000000..7ab8f1fe
--- /dev/null
+++ b/tests/from_pycapsule_test.py
@@ -0,0 +1,45 @@
+import sys
+
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+import pytest
+
+import narwhals.stable.v1 as nw
+from narwhals.utils import parse_version
+from tests.utils import compare_dicts
+
+
+@pytest.mark.xfail(parse_version(pa.__version__) < (14,), reason="too old")
+def test_from_arrow_to_arrow() -> None:
+    df = nw.from_native(pl.DataFrame({"ab": [1, 2, 3], "ba": [4, 5, 6]}), eager_only=True)
+    result = nw.from_arrow(df, native_namespace=pa)
+    assert isinstance(result.to_native(), pa.Table)
+    expected = {"ab": [1, 2, 3], "ba": [4, 5, 6]}
+    compare_dicts(result, expected)
+
+
+@pytest.mark.xfail(parse_version(pa.__version__) < (14,), reason="too old")
+def test_from_arrow_to_polars(monkeypatch: pytest.MonkeyPatch) -> None:
+    tbl = pa.table({"ab": [1, 2, 3], "ba": [4, 5, 6]})
+    monkeypatch.delitem(sys.modules, "pandas")
+    df = nw.from_native(tbl, eager_only=True)
+    result = nw.from_arrow(df, native_namespace=pl)
+    assert isinstance(result.to_native(), pl.DataFrame)
+    expected = {"ab": [1, 2, 3], "ba": [4, 5, 6]}
+    compare_dicts(result, expected)
+    assert "pandas" not in sys.modules
+
+
+@pytest.mark.xfail(parse_version(pa.__version__) < (14,), reason="too old")
+def test_from_arrow_to_pandas() -> None:
+    df = nw.from_native(pa.table({"ab": [1, 2, 3], "ba": [4, 5, 6]}), eager_only=True)
+    result = nw.from_arrow(df, native_namespace=pd)
+    assert isinstance(result.to_native(), pd.DataFrame)
+    expected = {"ab": [1, 2, 3], "ba": [4, 5, 6]}
+    compare_dicts(result, expected)
+
+
+def test_from_arrow_invalid() -> None:
+    with pytest.raises(TypeError, match="PyCapsule"):
+        nw.from_arrow({"a": [1]}, native_namespace=pa)  # type: ignore[arg-type]