feat(api): add array to string join operation

ibis-project · Feb 8, 2023 · 74de349 · 74de349
1 parent fd6ea5b
commit 74de349
Show file tree

Hide file tree

Showing 10 changed files with 132 additions and 2 deletions.
diff --git a/ibis/backends/bigquery/registry.py b/ibis/backends/bigquery/registry.py
@@ -723,6 +723,7 @@ def _interval_multiply(t, op):
     ops.RandomScalar: fixed_arity("RAND", 0),
     ops.NthValue: _nth_value,
     ops.JSONGetItem: lambda t, op: f"{t.translate(op.arg)}[{t.translate(op.index)}]",
+    ops.ArrayStringJoin: lambda t, op: f"ARRAY_TO_STRING({t.translate(op.arg)}, {t.translate(op.sep)})",
 }
 
 _invalid_operations = {

diff --git a/ibis/backends/clickhouse/compiler/values.py b/ibis/backends/clickhouse/compiler/values.py
@@ -1323,3 +1323,10 @@ def _extract_query(op, **kw):
 def _extract_fragment(op, **kw):
     arg = translate_val(op.arg, **kw)
     return f"nullIf(fragment({arg}), '')"
+
+
+@translate_val.register(ops.ArrayStringJoin)
+def _array_string_join(op, **kw):
+    arg = translate_val(op.arg, **kw)
+    sep = translate_val(op.sep, **kw)
+    return f"arrayStringConcat({arg}, {sep})"
diff --git a/ibis/backends/duckdb/registry.py b/ibis/backends/duckdb/registry.py
@@ -275,6 +275,9 @@ def _struct_column(t, op):
         ops.IntervalAdd: fixed_arity(operator.add, 2),
         ops.IntervalSubtract: fixed_arity(operator.sub, 2),
         ops.Capitalize: alchemy.sqlalchemy_operation_registry[ops.Capitalize],
+        ops.ArrayStringJoin: fixed_arity(
+            lambda sep, arr: sa.func.array_aggr(arr, sa.text("'string_agg'"), sep), 2
+        ),
     }
 )
 

diff --git a/ibis/backends/postgres/registry.py b/ibis/backends/postgres/registry.py
@@ -617,5 +617,8 @@ def translate(t, op: ops.ArgMin | ops.ArgMax) -> str:
                 else_=sa.null(),
             )
         ),
+        ops.ArrayStringJoin: fixed_arity(
+            lambda sep, arr: sa.func.array_to_string(arr, sep), 2
+        ),
     }
 )
diff --git a/ibis/backends/pyspark/compiler.py b/ibis/backends/pyspark/compiler.py
@@ -1917,3 +1917,10 @@ def compile_argmin(t, op, **kwargs):
 @compiles(ops.ArgMax)
 def compile_argmax(t, op, **kwargs):
     return compile_aggregator(t, op, fn=F.max_by, **kwargs)
+
+
+@compiles(ops.ArrayStringJoin)
+def compile_array_string_join(t, op, **kwargs):
+    arg = t.translate(op.arg, **kwargs)
+    sep = t.translate(op.sep, raw=True, **kwargs)
+    return F.concat_ws(sep, arg)
diff --git a/ibis/backends/tests/test_string.py b/ibis/backends/tests/test_string.py
@@ -542,3 +542,15 @@ def test_capitalize(con):
     expected = "Abc"
     expr = s.capitalize()
     assert con.execute(expr) == expected
+
+
+@pytest.mark.notimpl(["dask", "datafusion", "pandas", "polars"])
+@pytest.mark.notyet(["impala", "mssql", "mysql", "sqlite"], reason="no arrays")
+def test_array_string_join(con):
+    s = ibis.array(["a", "b", "c"])
+    expected = "a,b,c"
+    expr = ibis.literal(",").join(s)
+    assert con.execute(expr) == expected
+
+    expr = s.join(",")
+    assert con.execute(expr) == expected
diff --git a/ibis/backends/trino/registry.py b/ibis/backends/trino/registry.py
@@ -317,6 +317,9 @@ def _cot(t, op):
         ),
         ops.TypeOf: unary(sa.func.typeof),
         ops.Unnest: _unnest,
+        ops.ArrayStringJoin: fixed_arity(
+            lambda sep, arr: sa.func.array_join(arr, sep), 2
+        ),
     }
 )
 

diff --git a/ibis/expr/operations/strings.py b/ibis/expr/operations/strings.py
@@ -137,6 +137,15 @@ def output_shape(self):
         return rlz.highest_precedence_shape(self.arg)
 
 
+@public
+class ArrayStringJoin(Value):
+    sep = rlz.string
+    arg = rlz.value(dt.Array(dt.string))
+
+    output_dtype = dt.string
+    output_shape = rlz.shape_like("args")
+
+
 @public
 class StartsWith(Value):
     arg = rlz.string

diff --git a/ibis/expr/types/arrays.py b/ibis/expr/types/arrays.py
@@ -341,6 +341,53 @@ def unnest(self) -> ir.Value:
         except com.ExpressionError:
             return expr
 
+    def join(self, sep: str | ir.StringValue) -> ir.StringValue:
+        """Join the elements of this array expression with `sep`.
+
+        Parameters
+        ----------
+        sep
+            Separator to use for joining array elements
+
+        Returns
+        -------
+        StringValue
+            Elements of `self` joined with `sep`
+
+        Examples
+        --------
+        >>> import ibis
+        >>> ibis.options.interactive = True
+        >>> t = ibis.memtable({"arr": [["a", "b", "c"], None, [], ["b", None]]})
+        >>> t
+        ┏━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ arr                  ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━┩
+        │ array<string>        │
+        ├──────────────────────┤
+        │ ['a', 'b', ... +1]   │
+        │ ∅                    │
+        │ []                   │
+        │ ['b', None]          │
+        └──────────────────────┘
+        >>> t.arr.join("|")
+        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ ArrayStringJoin('|', arr) ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+        │ string                    │
+        ├───────────────────────────┤
+        │ a|b|c                     │
+        │ ∅                         │
+        │ ∅                         │
+        │ b                         │
+        └───────────────────────────┘
+
+        See Also
+        --------
+        [`StringValue.join`][ibis.expr.types.strings.StringValue.join]
+        """
+        return ops.ArrayStringJoin(sep, self).to_expr()
+
 
 @public
 class ArrayScalar(Scalar, ArrayValue):

diff --git a/ibis/expr/types/strings.py b/ibis/expr/types/strings.py
@@ -383,7 +383,7 @@ def find_in_set(self, str_list: Sequence[str]) -> ir.IntegerValue:
         """
         return ops.FindInSet(self, str_list).to_expr()
 
-    def join(self, strings: Sequence[str | StringValue]) -> StringValue:
+    def join(self, strings: Sequence[str | StringValue] | ir.ArrayValue) -> StringValue:
         """Join a list of strings using `self` as the separator.
 
         Parameters
@@ -401,8 +401,46 @@ def join(self, strings: Sequence[str | StringValue]) -> StringValue:
         -------
         StringValue
             Joined string
+
+        Examples
+        --------
+        >>> import ibis
+        >>> ibis.options.interactive = True
+        >>> t = ibis.memtable({"arr": [["a", "b", "c"], None, [], ["b", None]]})
+        >>> t
+        ┏━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ arr                  ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━┩
+        │ array<string>        │
+        ├──────────────────────┤
+        │ ['a', 'b', ... +1]   │
+        │ ∅                    │
+        │ []                   │
+        │ ['b', None]          │
+        └──────────────────────┘
+        >>> ibis.literal("|").join(t.arr)
+        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+        ┃ ArrayStringJoin('|', arr) ┃
+        ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+        │ string                    │
+        ├───────────────────────────┤
+        │ a|b|c                     │
+        │ ∅                         │
+        │ ∅                         │
+        │ b                         │
+        └───────────────────────────┘
+
+        See Also
+        --------
+        [`ArrayValue.join`][ibis.expr.types.arrays.ArrayValue.join]
         """
-        return ops.StringJoin(self, strings).to_expr()
+        import ibis.expr.types as ir
+
+        if isinstance(strings, ir.ArrayValue):
+            cls = ops.ArrayStringJoin
+        else:
+            cls = ops.StringJoin
+        return cls(self, strings).to_expr()
 
     def startswith(self, start: str | StringValue) -> ir.BooleanValue:
         """Determine whether `self` starts with `end`.