feat(mssql): add lpad and rpad ops

ibis-project · Sep 9, 2024 · 7c128db · 7c128db
1 parent 422d361
commit 7c128db
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 14 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,17 @@ 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 self.f.left(
+            self.f.right(
+                self.f.concat(self.f.replicate(pad, length - self.f.length(arg)), arg),
+                length + self.f.length(arg),
+            ),
+            length,
+        )
+
+    def visit_RPad(self, op, *, arg, length, pad=" "):
+        return self.f.left(self.f.concat(arg, self.f.replicate(pad, length)), 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"]),
@@ -1105,10 +1103,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,
@@ -1136,10 +1130,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,

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"),