feat(flink): implement array operators

ibis/backends/flink/registry.py

@@ -255,6 +255,21 @@ def _floor_divide(translator: ExprTranslator, op: ops.Node) -> str:
     return f"FLOOR(({left}) / ({right}))"
 
 
+def _array_column(translator: ExprTranslator, op: ops.arrays.ArrayColumn) -> str:
+    return "ARRAY[{}]".format(", ".join(map(translator.translate, op.cols)))
+
+
+def _array_contains(translator: ExprTranslator, op: ops.arrays.ArrayContains) -> str:
+    arg = translator.translate(op.arg)
+    other = translator.translate(op.other)
+    return f"ARRAY_CONTAINS({arg}, {other})"
+
+
+def _array_distinct(translator: ExprTranslator, op: ops.arrays.ArrayDistinct) -> str:
+    arg = translator.translate(op.arg)
+    return f"ARRAY_DISTINCT({arg})"
+
+
 def _array_index(translator: ExprTranslator, op: ops.arrays.ArrayIndex):
     table_column = op.arg
     index = op.index
@@ -269,6 +284,43 @@ def _array_length(translator: ExprTranslator, op: ops.arrays.ArrayLength) -> str
     return f"CARDINALITY({translator.translate(op.arg)})"
 
 
+def _array_position(translator: ExprTranslator, op: ops.arrays.ArrayPosition) -> str:
+    arg = translator.translate(op.arg)
+    other = translator.translate(op.other)
+    return f"ARRAY_POSITION({arg}, {other}) - 1"
+
+
+def _array_remove(translator: ExprTranslator, op: ops.arrays.ArrayRemove) -> str:
+    arg = translator.translate(op.arg)
+    other = translator.translate(op.other)
+    return f"ARRAY_REMOVE({arg}, {other})"
+
+
+def _array_slice(translator: ExprTranslator, op: ops.arrays.ArraySlice) -> str:
+    array = translator.translate(op.arg)
+    start = op.start.value
+    # Note (mehmet): The offsets are 1-based for ARRAY_SLICE.
+    # Ref: https://nightlies.apache.org/flink/flink-docs-master/docs/dev/table/functions/systemfunctions
+    if start >= 0:
+        start += 1
+
+    if op.stop is None:
+        return f"ARRAY_SLICE({array}, {start})"
+
+    stop = op.stop.value
+    if stop >= 0:
+        return f"ARRAY_SLICE({array}, {start}, {stop})"
+    else:
+        # Note (mehmet): To imitate the behavior of pandas array slicing.
+        return f"ARRAY_SLICE({array}, {start}, CARDINALITY({array}) - {abs(stop)})"
+
+
+def _array_union(translator: ExprTranslator, op: ops.arrays.ArrayUnion) -> str:
+    left = translator.translate(op.left)
+    right = translator.translate(op.right)
+    return f"ARRAY_UNION({left}, {right})"
+
+
 def _json_get_item(translator: ExprTranslator, op: ops.json.JSONGetItem) -> str:
     arg_translated = translator.translate(op.arg)
     if op.index.dtype.is_integer():
@@ -442,9 +494,16 @@ def _struct_field(translator, op):
         # Binary operations
         ops.Power: fixed_arity("power", 2),
         ops.FloorDivide: _floor_divide,
-        # Collection functions
+        # Collection operations
+        ops.ArrayColumn: _array_column,
+        ops.ArrayContains: _array_contains,
+        ops.ArrayDistinct: _array_distinct,
         ops.ArrayIndex: _array_index,
         ops.ArrayLength: _array_length,
+        ops.ArrayPosition: _array_position,
+        ops.ArrayRemove: _array_remove,
+        ops.ArraySlice: _array_slice,
+        ops.ArrayUnion: _array_union,
         ops.JSONGetItem: _json_get_item,
         ops.Map: _map,
         ops.MapGet: _map_get,

ibis/backends/flink/tests/conftest.py

@@ -49,14 +49,15 @@ def connect(*, tmpdir, worker_id, **kw: Any):
     def _load_data(self, **_: Any) -> None:
         import pandas as pd
 
-        from ibis.backends.tests.data import json_types, struct_types
+        from ibis.backends.tests.data import array_types, json_types, struct_types
 
         for table_name in TEST_TABLES:
             path = self.data_dir / "parquet" / f"{table_name}.parquet"
             self.connection.create_table(table_name, pd.read_parquet(path), temp=True)
 
         self.connection.create_table("json_t", json_types, temp=True)
         self.connection.create_table("struct", struct_types, temp=True)
+        self.connection.create_table("array_types", array_types, temp=True)
 
 
 class TestConfForStreaming(TestConf):