cudf-polars string/numeric casting

python/cudf_polars/cudf_polars/dsl/expressions/unary.py

@@ -9,6 +9,7 @@
 
 import pyarrow as pa
 import pylibcudf as plc
+from pylibcudf.traits import is_floating_point, is_integral_not_bool
 
 from cudf_polars.containers import Column
 from cudf_polars.dsl.expressions.base import AggInfo, ExecutionContext, Expr
@@ -23,6 +24,10 @@
 __all__ = ["Cast", "UnaryFunction", "Len"]
 
 
+def _is_int_or_float(dtype: plc.DataType) -> bool:
+    return is_integral_not_bool(dtype) or is_floating_point(dtype)
+
+
 class Cast(Expr):
     """Class representing a cast of an expression."""
 
@@ -33,7 +38,19 @@ class Cast(Expr):
     def __init__(self, dtype: plc.DataType, value: Expr) -> None:
         super().__init__(dtype)
         self.children = (value,)
-        if not dtypes.can_cast(value.dtype, self.dtype):
+        if (
+            self.dtype.id() == plc.TypeId.STRING
+            or value.dtype.id() == plc.TypeId.STRING
+        ):
+            if not (
+                (self.dtype.id() == plc.TypeId.STRING and _is_int_or_float(value.dtype))
+                or (
+                    _is_int_or_float(self.dtype)
+                    and value.dtype.id() == plc.TypeId.STRING
+                )
+            ):
+                raise NotImplementedError("Only string to float cast is supported")
+        elif not dtypes.can_cast(value.dtype, self.dtype):
             raise NotImplementedError(
                 f"Can't cast {self.dtype.id().name} to {value.dtype.id().name}"
             )
@@ -48,7 +65,29 @@ def do_evaluate(
         """Evaluate this expression given a dataframe for context."""
         (child,) = self.children
         column = child.evaluate(df, context=context, mapping=mapping)
-        return Column(plc.unary.cast(column.obj, self.dtype)).sorted_like(column)
+        if (
+            self.dtype.id() == plc.TypeId.STRING
+            or column.obj.type().id() == plc.TypeId.STRING
+        ):
+            if self.dtype.id() == plc.TypeId.STRING:
+                if is_floating_point(column.obj.type()):
+                    result = plc.strings.convert.convert_floats.from_floats(column.obj)
+                else:
+                    result = plc.strings.convert.convert_integers.from_integers(
+                        column.obj
+                    )
+            else:
+                if is_floating_point(self.dtype):
+                    result = plc.strings.convert.convert_floats.to_floats(
+                        column.obj, self.dtype
+                    )
+                else:
+                    result = plc.strings.convert.convert_integers.to_integers(
+                        column.obj, self.dtype
+                    )
+        else:
+            result = plc.unary.cast(column.obj, self.dtype)
+        return Column(result).sorted_like(column)
 
     def collect_agg(self, *, depth: int) -> AggInfo:
         """Collect information about aggregations in groupbys."""

python/cudf_polars/tests/expressions/test_stringfunction.py

@@ -40,6 +40,34 @@ def ldf(with_nulls):
     )
 
 
+@pytest.fixture(params=[pl.Float32, pl.Float64, pl.Int8, pl.Int16, pl.Int32, pl.Int64])
+def numeric_type(request):
+    return request.param
+
+
+@pytest.fixture
+def str_to_numeric_data(with_nulls):
+    a = ["1", "2", "3", "4", "5", "6"]
+    if with_nulls:
+        a[4] = None
+    return pl.LazyFrame({"a": a})
+
+
+@pytest.fixture
+def str_from_numeric_data(with_nulls, numeric_type):
+    a = [
+        1,
+        2,
+        3,
+        4,
+        5,
+        6,
+    ]
+    if with_nulls:
+        a[4] = None
+    return pl.LazyFrame({"a": pl.Series(a, dtype=numeric_type)})
+
+
 slice_cases = [
     (1, 3),
     (0, 3),
@@ -337,3 +365,13 @@ def test_unsupported_regex_raises(pattern):
 
     q = df.select(pl.col("a").str.contains(pattern, strict=True))
     assert_ir_translation_raises(q, NotImplementedError)
+
+
+def test_string_to_numeric(str_to_numeric_data, numeric_type):
+    query = str_to_numeric_data.select(pl.col("a").cast(numeric_type))
+    assert_gpu_result_equal(query)
+
+
+def test_string_from_numeric(str_from_numeric_data):
+    query = str_from_numeric_data.select(pl.col("a").cast(pl.String))
+    assert_gpu_result_equal(query)