Repository: narwhals-dev/narwhals
Pull request #525 (2024-07-15 06:50:40)

feat: pyarrow `concat`
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #372 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index d6fe38ba..6b7ca2fe 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -6,8 +6,11 @@ from typing import Any
 from typing import Iterable
 
 from narwhals import dtypes
+from narwhals._arrow.dataframe import ArrowDataFrame
 from narwhals._arrow.expr import ArrowExpr
 from narwhals._arrow.series import ArrowSeries
+from narwhals._arrow.utils import horizontal_concat
+from narwhals._arrow.utils import vertical_concat
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals.dependencies import get_pyarrow
 from narwhals.utils import flatten
@@ -15,7 +18,6 @@ from narwhals.utils import flatten
 if TYPE_CHECKING:
     from typing import Callable
 
-    from narwhals._arrow.dataframe import ArrowDataFrame
     from narwhals._arrow.typing import IntoArrowExpr
 
 
@@ -159,3 +161,23 @@ class ArrowNamespace:
                 namespace=self,
             ),
         )
+
+    def concat(
+        self,
+        items: Iterable[ArrowDataFrame],
+        *,
+        how: str = "vertical",
+    ) -> ArrowDataFrame:
+        dfs: list[Any] = [item._native_dataframe for item in items]
+
+        if how == "horizontal":
+            return ArrowDataFrame(
+                horizontal_concat(dfs),
+                backend_version=self._backend_version,
+            )
+        if how == "vertical":
+            return ArrowDataFrame(
+                vertical_concat(dfs),
+                backend_version=self._backend_version,
+            )
+        raise NotImplementedError
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index 6bda7e8d..13fea297 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -141,3 +141,45 @@ def validate_dataframe_comparand(length: int, other: Any) -> Any:
         return other._native_series
     msg = "Please report a bug"  # pragma: no cover
     raise AssertionError(msg)
+
+
+def horizontal_concat(dfs: list[Any]) -> Any:
+    """
+    Concatenate (native) DataFrames horizontally.
+
+    Should be in namespace.
+    """
+    pa = get_pyarrow()
+    if not dfs:
+        msg = "No dataframes to concatenate"  # pragma: no cover
+        raise AssertionError(msg)
+
+    names = [name for df in dfs for name in df.column_names]
+
+    if len(set(names)) < len(names):  # pragma: no cover
+        msg = "Expected unique column names"
+        raise ValueError(msg)
+
+    arrays = [a for df in dfs for a in df]
+    return pa.Table.from_arrays(arrays, names=names)
+
+
+def vertical_concat(dfs: list[Any]) -> Any:
+    """
+    Concatenate (native) DataFrames vertically.
+
+    Should be in namespace.
+    """
+    if not dfs:
+        msg = "No dataframes to concatenate"  # pragma: no cover
+        raise AssertionError(msg)
+
+    cols = set(dfs[0].column_names)
+    for df in dfs:
+        cols_current = set(df.column_names)
+        if cols_current != cols:
+            msg = "unable to vstack, column names don't match"
+            raise TypeError(msg)
+
+    pa = get_pyarrow()
+    return pa.concat_tables(dfs).combine_chunks()


Test patch:
diff --git a/tests/frame/test_common.py b/tests/frame/test_common.py
index 8e9c83ac..01eb3495 100644
--- a/tests/frame/test_common.py
+++ b/tests/frame/test_common.py
@@ -190,10 +190,7 @@ def test_drop(constructor: Any, drop: list[str], left: list[str]) -> None:
     assert df.drop(*drop).columns == left
 
 
-def test_concat_horizontal(request: Any, constructor_with_lazy: Any) -> None:
-    if "pyarrow_table" in str(constructor_with_lazy):
-        request.applymarker(pytest.mark.xfail)
-
+def test_concat_horizontal(constructor_with_lazy: Any) -> None:
     df_left = nw.from_native(constructor_with_lazy(data))
     df_right = nw.from_native(constructor_with_lazy(data_right))
     result = nw.concat([df_left, df_right], how="horizontal")
@@ -210,10 +207,7 @@ def test_concat_horizontal(request: Any, constructor_with_lazy: Any) -> None:
         nw.concat([])
 
 
-def test_concat_vertical(request: Any, constructor_with_lazy: Any) -> None:
-    if "pyarrow_table" in str(constructor_with_lazy):
-        request.applymarker(pytest.mark.xfail)
-
+def test_concat_vertical(constructor_with_lazy: Any) -> None:
     df_left = (
         nw.from_native(constructor_with_lazy(data))
         .rename({"a": "c", "b": "d"})
@@ -228,7 +222,8 @@ def test_concat_vertical(request: Any, constructor_with_lazy: Any) -> None:
 
     with pytest.raises(ValueError, match="No items"):
         nw.concat([], how="vertical")
-    with pytest.raises(Exception, match="unable to vstack"):
+
+    with pytest.raises((Exception, TypeError), match="unable to vstack"):
         nw.concat([df_left, df_right.rename({"d": "i"})], how="vertical").collect()