feat(datafusion): support substr without end

ibis-project · Feb 16, 2023 · a19fd87 · a19fd87
1 parent 5a74cb4
commit a19fd87
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/ibis/backends/datafusion/compiler.py b/ibis/backends/datafusion/compiler.py
@@ -260,8 +260,11 @@ def capitalize(op):
 def substring(op):
     arg = translate(op.arg)
     start = translate(ops.Add(left=op.start, right=1))
-    length = translate(op.length)
-    return df.functions.substr(arg, start, length)
+    if op_length := op.length:
+        length = translate(op_length)
+        return df.functions.substr(arg, start, length)
+    else:
+        return df.functions.substr(arg, start)
 
 
 @translate.register(ops.Repeat)

diff --git a/ibis/backends/tests/test_string.py b/ibis/backends/tests/test_string.py
@@ -490,10 +490,6 @@ def test_string_col_is_unicode(alltypes, df):
             lambda t: t.date_string_col.str[2:],
             id='substr-start-only',
             marks=[
-                pytest.mark.notimpl(
-                    ["datafusion"],
-                    raises=NotImplementedError,
-                ),
                 pytest.mark.notimpl(
                     ["pyspark"],
                     raises=com.OperationNotDefinedError,