fixup lazy power

dataframe_api_compat/polars_standard/polars_standard.py

@@ -274,15 +274,23 @@ def __lt__(self, other: Column[DType] | Any) -> PolarsColumn[Bool]:
         return PolarsColumn(self.column < other, hash=self._hash, dtype=pl.Boolean())
 
     def __mul__(self, other: Column[DType] | Any) -> PolarsColumn[Any]:
-        if isinstance(self.column, pl.Expr):
-            raise NotImplementedError("mul not implemented for lazy columns")
-        if isinstance(other, PolarsColumn) and isinstance(other.column, pl.Expr):
-            raise NotImplementedError("mul not implemented for lazy columns")
         if isinstance(other, PolarsColumn):
             res = self.column * other.column
-            return PolarsColumn(res, dtype=res.dtype, hash=self._hash)  # type: ignore[arg-type]
+            res_dtype = (
+                pl.DataFrame(
+                    {"a": [1], "b": [1]}, schema={"a": self._dtype, "b": other._dtype}
+                )
+                .select(result=pl.col("a") * pl.col("b"))
+                .schema["result"]
+            )
+            return PolarsColumn(res, dtype=res_dtype, hash=self._hash)  # type: ignore[arg-type]
         res = self.column * other
-        return PolarsColumn(res, dtype=res.dtype, hash=self._hash)  # type: ignore[arg-type]
+        res_dtype = (
+            pl.DataFrame({"a": [1]}, schema={"a": self._dtype})
+            .select(result=pl.col("a") * other)
+            .schema["result"]
+        )
+        return PolarsColumn(res, dtype=res_dtype, hash=self._hash)  # type: ignore[arg-type]
 
     def __floordiv__(self, other: Column[DType] | Any) -> PolarsColumn[Any]:
         if isinstance(other, PolarsColumn):
@@ -304,25 +312,30 @@ def __truediv__(self, other: Column[DType] | Any) -> PolarsColumn[Any]:
         return PolarsColumn(res, dtype=res.dtype, hash=self._hash)  # type: ignore[arg-type]
 
     def __pow__(self, other: Column[DType] | Any) -> PolarsColumn[Any]:
-        if isinstance(self.column, pl.Expr) or (
-            isinstance(other, PolarsColumn) and isinstance(other.column, pl.Expr)
-        ):
-            raise NotImplementedError("pow not implemented for lazy columns")
-        original_type = self.column.dtype
+        original_type = self._dtype
         if isinstance(other, PolarsColumn):
-            assert isinstance(other.column, pl.Series)  # help mypy
             ret = self.column.pow(other.column)
-            if _is_integer_dtype(original_type) and _is_integer_dtype(other.column.dtype):
-                if (other.column < 0).any():
-                    raise ValueError("Cannot raise integer to negative power")
-                ret = ret.cast(original_type)
+            ret_type = (
+                pl.DataFrame(
+                    {"a": [1], "b": [1]}, schema={"a": original_type, "b": other._dtype}
+                )
+                .select(result=pl.col("a") ** pl.col("b"))
+                .schema["result"]
+            )
+            if _is_integer_dtype(original_type) and _is_integer_dtype(other._dtype):
+                ret_type = original_type
+                ret = ret.cast(ret_type)
         else:
             ret = self.column.pow(other)  # type: ignore[arg-type]
+            ret_type = (
+                pl.DataFrame({"a": [1]}, schema={"a": original_type})
+                .select(result=pl.col("a") ** other)
+                .schema["result"]
+            )
             if _is_integer_dtype(original_type) and isinstance(other, int):
-                if other < 0:
-                    raise ValueError("Cannot raise integer to negative power")
-                ret = ret.cast(original_type)
-        return PolarsColumn(ret, dtype=ret.dtype)  # type: ignore[arg-type]
+                ret_type = original_type
+                ret = ret.cast(ret_type)
+        return PolarsColumn(ret, dtype=ret_type, hash=self._hash)  # type: ignore[arg-type]
 
     def __mod__(self, other: Column[DType] | Any) -> PolarsColumn[Any]:
         if isinstance(self.column, pl.Expr) or (

tests/column/pow_test.py

@@ -1,42 +1,58 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import pandas as pd
-import pytest
 
-from tests.utils import convert_series_to_pandas_numpy
-from tests.utils import integer_series_1
+from tests.utils import convert_dataframe_to_pandas_numpy
+from tests.utils import integer_dataframe_1
+from tests.utils import interchange_to_pandas
+
+if TYPE_CHECKING:
+    import pytest
 
 
 def test_float_powers_column(library: str, request: pytest.FixtureRequest) -> None:
-    ser = integer_series_1(library, request)
-    other = integer_series_1(library, request) * 1.0
-    result = ser.__pow__(other)
-    namespace = ser.__column_namespace__()
-    result_pd = pd.api.interchange.from_dataframe(
-        namespace.dataframe_from_dict({"result": (result).rename("result")}).dataframe
-    )["result"]
-    expected = pd.Series([1.0, 4.0, 27.0], name="result")
-    result_pd = convert_series_to_pandas_numpy(result_pd)
-    pd.testing.assert_series_equal(result_pd, expected)
+    df = integer_dataframe_1(library)
+    ser = df.get_column_by_name("a")
+    other = df.get_column_by_name("b") * 1.0
+    result = df.insert(0, "result", ser.__pow__(other))
+    result_pd = interchange_to_pandas(result, library)
+    expected = pd.DataFrame(
+        {"result": [1.0, 32.0, 729.0], "a": [1, 2, 3], "b": [4, 5, 6]}
+    )
+    result_pd = convert_dataframe_to_pandas_numpy(result_pd)
+    pd.testing.assert_frame_equal(result_pd, expected)
 
 
 def test_float_powers_scalar_column(library: str, request: pytest.FixtureRequest) -> None:
-    ser = integer_series_1(library, request)
+    df = integer_dataframe_1(library)
+    ser = df.get_column_by_name("a")
     other = 1.0
-    result = ser.__pow__(other)
-    namespace = ser.__column_namespace__()
-    result_pd = pd.api.interchange.from_dataframe(
-        namespace.dataframe_from_dict({"result": (result).rename("result")}).dataframe
-    )["result"]
-    expected = pd.Series([1.0, 2.0, 3.0], name="result")
-    result_pd = convert_series_to_pandas_numpy(result_pd)
-    pd.testing.assert_series_equal(result_pd, expected)
-
-
-def test_negative_powers_column(library: str, request: pytest.FixtureRequest) -> None:
-    ser = integer_series_1(library, request)
-    other = integer_series_1(library, request) * -1
-    with pytest.raises(ValueError):
-        ser.__pow__(-1)
-    with pytest.raises(ValueError):
-        ser.__pow__(other)
+    result = df.insert(0, "result", ser.__pow__(other))
+    result_pd = interchange_to_pandas(result, library)
+    expected = pd.DataFrame({"result": [1.0, 2.0, 3.0], "a": [1, 2, 3], "b": [4, 5, 6]})
+    result_pd = convert_dataframe_to_pandas_numpy(result_pd)
+    pd.testing.assert_frame_equal(result_pd, expected)
+
+
+def test_int_powers_column(library: str, request: pytest.FixtureRequest) -> None:
+    df = integer_dataframe_1(library)
+    ser = df.get_column_by_name("a")
+    other = df.get_column_by_name("b") * 1
+    result = df.insert(0, "result", ser.__pow__(other))
+    result_pd = interchange_to_pandas(result, library)
+    expected = pd.DataFrame({"result": [1, 32, 729], "a": [1, 2, 3], "b": [4, 5, 6]})
+    result_pd = convert_dataframe_to_pandas_numpy(result_pd)
+    pd.testing.assert_frame_equal(result_pd, expected)
+
+
+def test_int_powers_scalar_column(library: str, request: pytest.FixtureRequest) -> None:
+    df = integer_dataframe_1(library)
+    ser = df.get_column_by_name("a")
+    other = 1
+    result = df.insert(0, "result", ser.__pow__(other))
+    result_pd = interchange_to_pandas(result, library)
+    expected = pd.DataFrame({"result": [1, 2, 3], "a": [1, 2, 3], "b": [4, 5, 6]})
+    result_pd = convert_dataframe_to_pandas_numpy(result_pd)
+    pd.testing.assert_frame_equal(result_pd, expected)