Making the blosc2 engine respect the default pandas udf signatures

Author: datapythonista

src/blosc2/proxy.py

@@ -699,9 +699,7 @@ def map(cls, data, func, args, kwargs, decorator, skip_na):
         with the NumPy array as the function parameter, instead of calling the
         function once for each element.
         """
-        data = cls._ensure_numpy_data(data)
-        func = decorator(func)
-        return func(data, *args, **kwargs)
+        raise NotImplementedError("The Blosc2 engine does not support map. Use apply instead.")
 
     @classmethod
     def apply(cls, data, func, args, kwargs, decorator, axis):
@@ -713,7 +711,23 @@ def apply(cls, data, func, args, kwargs, decorator, axis):
         """
         data = cls._ensure_numpy_data(data)
         func = decorator(func)
-        return func(data, *args, **kwargs)
+        if data.ndim == 1 or axis is None:
+            # pandas Series.apply or pipe
+            return func(data, *args, **kwargs)
+        elif axis in (0, "index"):
+            # pandas apply(axis=0) column-wise
+            result = []
+            for row_idx in range(data.shape[1]):
+                result.append(func(data[:, row_idx], *args, **kwargs))
+            return np.vstack(result).transpose()
+        elif axis == (1, "columns"):
+            # pandas apply(axis=1) row-wise
+            result = []
+            for col_idx in range(data.shape[0]):
+                result.append(func(data[col_idx, :], *args, **kwargs))
+            return np.vstack(result)
+        else:
+            raise NotImplementedError(f"Unknown axis '{axis}'. Use one of 0, 1 or None.")
 
 
 jit.__pandas_udf__ = PandasUdfEngine

tests/test_pandas_udf_engine.py

@@ -7,72 +7,88 @@
 #######################################################################
 
 import numpy as np
+import pytest
 
 import blosc2
 
 
 class TestPandasUDF:
-    def test_map_1d(self):
+    def test_map(self):
         def add_one(x):
             return x + 1
 
         data = np.array([1, 2])
 
-        result = blosc2.jit.__pandas_udf__.map(
+        with pytest.raises(NotImplementedError):
+            blosc2.jit.__pandas_udf__.map(
+                data,
+                add_one,
+                args=(),
+                kwargs={},
+                decorator=blosc2.jit,
+                skip_na=False,
+            )
+
+    def test_apply_1d(self):
+        def add_one(x):
+            return x + 1
+
+        data = np.array([1, 2])
+
+        result = blosc2.jit.__pandas_udf__.apply(
             data,
             add_one,
             args=(),
             kwargs={},
             decorator=blosc2.jit,
-            skip_na=False,
+            axis=0,
         )
         assert result.shape == (2,)
         assert result[0] == 2
         assert result[1] == 3
 
-    def test_map_1d_with_args(self):
+    def test_apply_1d_with_args(self):
         def add_numbers(x, num1, num2):
             return x + num1 + num2
 
         data = np.array([1, 2])
 
-        result = blosc2.jit.__pandas_udf__.map(
+        result = blosc2.jit.__pandas_udf__.apply(
             data,
             add_numbers,
             args=(10,),
             kwargs={"num2": 100},
             decorator=blosc2.jit,
-            skip_na=False,
+            axis=0,
         )
         assert result.shape == (2,)
         assert result[0] == 111
         assert result[1] == 112
 
-    def test_map_2d(self):
+    def test_apply_2d(self):
         def add_one(x):
+            assert x.shape == (2, 3)
             return x + 1
 
-        data = np.array([[1, 2], [3, 4]])
+        data = np.array([[1, 2, 3], [4, 5, 6]])
 
-        result = blosc2.jit.__pandas_udf__.map(
+        result = blosc2.jit.__pandas_udf__.apply(
             data,
             add_one,
             args=(),
             kwargs={},
             decorator=blosc2.jit,
-            skip_na=False,
+            axis=None,
         )
-        assert result.shape == (2, 2)
-        assert result[0, 0] == 2
-        assert result[0, 1] == 3
-        assert result[1, 0] == 4
-        assert result[1, 1] == 5
+        expected = np.array([[2, 3, 4], [5, 6, 7]])
+        assert np.array_equal(result, expected)
 
-    def test_apply_1d(self):
+    def test_apply_2d_by_column(self):
         def add_one(x):
+            assert x.shape == (2,)
             return x + 1
 
-        data = np.array([1, 2])
+        data = np.array([[1, 2, 3], [4, 5, 6]])
 
         result = blosc2.jit.__pandas_udf__.apply(
             data,
@@ -82,44 +98,23 @@ def add_one(x):
             decorator=blosc2.jit,
             axis=0,
         )
-        assert result.shape == (2,)
-        assert result[0] == 2
-        assert result[1] == 3
+        expected = np.array([[2, 3, 4], [5, 6, 7]])
+        assert np.array_equal(result, expected)
 
-    def test_apply_1d_with_args(self):
-        def add_numbers(x, num1, num2):
-            return x + num1 + num2
-
-        data = np.array([1, 2])
-
-        result = blosc2.jit.__pandas_udf__.apply(
-            data,
-            add_numbers,
-            args=(10,),
-            kwargs={"num2": 100},
-            decorator=blosc2.jit,
-            axis=0,
-        )
-        assert result.shape == (2,)
-        assert result[0] == 111
-        assert result[1] == 112
-
-    def test_apply_2d(self):
+    def test_apply_2d_by_row(self):
         def add_one(x):
+            assert x.shape == (3,)
             return x + 1
 
-        data = np.array([[1, 2], [3, 4]])
+        data = np.array([[1, 2, 3], [4, 5, 6]])
 
         result = blosc2.jit.__pandas_udf__.apply(
             data,
             add_one,
             args=(),
             kwargs={},
             decorator=blosc2.jit,
-            axis=0,
+            axis=1,
         )
-        assert result.shape == (2, 2)
-        assert result[0, 0] == 2
-        assert result[0, 1] == 3
-        assert result[1, 0] == 4
-        assert result[1, 1] == 5
+        expected = np.array([[2, 3, 4], [5, 6, 7]])
+        assert np.array_equal(result, expected)