feat(mssql): add lpad and rpad ops (#10060)

Co-authored-by: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
ibis-project · Sep 10, 2024 · 77af14b · 77af14b
1 parent a32e8e5
commit 77af14b
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 16 deletions.
diff --git a/ibis/backends/sql/compilers/mssql.py b/ibis/backends/sql/compilers/mssql.py
@@ -95,7 +95,6 @@ class MSSQLCompiler(SQLGlotCompiler):
         ops.IntervalFloorDivide,
         ops.IsInf,
         ops.IsNan,
-        ops.LPad,
         ops.Levenshtein,
         ops.Map,
         ops.Median,
@@ -106,7 +105,6 @@ class MSSQLCompiler(SQLGlotCompiler):
         ops.RegexSearch,
         ops.RegexSplit,
         ops.RowID,
-        ops.RPad,
         ops.StringSplit,
         ops.StringToDate,
         ops.StringToTimestamp,
@@ -526,5 +524,23 @@ def visit_StartsWith(self, op, *, arg, start):
     def visit_EndsWith(self, op, *, arg, end):
         return arg.like(self.f.concat("%", end))
 
+    def visit_LPad(self, op, *, arg, length, pad):
+        return sge.Case(
+            ifs=[self.if_(length <= self.f.length(arg), arg)],
+            default=self.f.left(
+                self.f.concat(self.f.replicate(pad, length - self.f.length(arg)), arg),
+                length,
+            ),
+        )
+
+    def visit_RPad(self, op, *, arg, length, pad):
+        return sge.Case(
+            ifs=[self.if_(length <= self.f.length(arg), arg)],
+            default=self.f.left(
+                self.f.concat(arg, self.f.replicate(pad, length - self.f.length(arg))),
+                length,
+            ),
+        )
+
 
 compiler = MSSQLCompiler()
diff --git a/ibis/backends/tests/test_string.py b/ibis/backends/tests/test_string.py
@@ -417,13 +417,11 @@ def uses_java_re(t):
             lambda t: t.string_col.lpad(10, "a"),
             lambda t: t.string_col.str.pad(10, fillchar="a", side="left"),
             id="lpad",
-            marks=pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError),
         ),
         param(
             lambda t: t.string_col.rpad(10, "a"),
             lambda t: t.string_col.str.pad(10, fillchar="a", side="right"),
             id="rpad",
-            marks=pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError),
         ),
         param(
             lambda t: t.string_col.find_in_set(["1"]),
@@ -1112,10 +1110,6 @@ def string_temp_table(backend, con):
             lambda t: t.str[:4].str.pad(4, side="right", fillchar="-"),
             id="rpad",
             marks=[
-                pytest.mark.notimpl(
-                    ["mssql"],
-                    raises=com.OperationNotDefinedError,
-                ),
                 pytest.mark.notyet(
                     ["flink", "oracle"],
                     raises=AssertionError,
@@ -1127,7 +1121,7 @@ def string_temp_table(backend, con):
                     reason="Treats len(🐍) == 4, len(Éé) == 4",
                 ),
                 pytest.mark.notyet(
-                    ["dask", "pandas", "polars"],
+                    ["dask", "mssql", "pandas", "polars"],
                     raises=AssertionError,
                     reason="Python style padding, e.g. doesn't trim strings to pad-length",
                 ),
@@ -1143,10 +1137,6 @@ def string_temp_table(backend, con):
             lambda t: t.str[:4].str.pad(4, side="left", fillchar="-"),
             id="lpad",
             marks=[
-                pytest.mark.notimpl(
-                    ["mssql"],
-                    raises=com.OperationNotDefinedError,
-                ),
                 pytest.mark.notyet(
                     ["flink", "oracle"],
                     raises=AssertionError,
@@ -1158,7 +1148,7 @@ def string_temp_table(backend, con):
                     reason="Treats len(🐍) == 4, len(Éé) == 4",
                 ),
                 pytest.mark.notyet(
-                    ["dask", "pandas", "polars"],
+                    ["dask", "mssql", "pandas", "polars"],
                     raises=AssertionError,
                     reason="Python style padding, e.g. doesn't trim strings to pad-length",
                 ),

diff --git a/ibis/backends/tests/test_temporal.py b/ibis/backends/tests/test_temporal.py
@@ -1734,7 +1734,6 @@ def test_date_scalar_from_iso(con):
     assert result.strftime("%Y-%m-%d") == "2022-02-24"
 
 
-@pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["exasol"], raises=AssertionError, strict=False)
 def test_date_column_from_iso(backend, con, alltypes, df):
     expr = (
@@ -1821,7 +1820,6 @@ def build_date_col(t):
     ).cast("date")
 
 
-@pytest.mark.notimpl(["mssql"], raises=com.OperationNotDefinedError)
 @pytest.mark.notimpl(["druid"], raises=PyDruidProgrammingError)
 @pytest.mark.parametrize(
     ("left_fn", "right_fn"),