More array functions (#70)

More array functions

pykusto/functions.py

@@ -1,4 +1,5 @@
 import json
+from itertools import chain
 from typing import Union
 
 from pykusto.expressions import AnyTypeColumn, NumberType, NumberExpression, TimespanType, \
@@ -32,11 +33,6 @@ class Functions:
         """
         return TimespanExpression.ago(expr)
 
-    # def array_concat(): return
-    #
-    #
-    # def array_iif(): return
-
     @staticmethod
     def array_length(expr: ArrayType) -> NumberExpression:
         """
@@ -710,14 +706,102 @@ class Functions:
     # def series_subtract(self): return
     #
     #
-    # def set_difference(self): return
-    #
-    #
     # def set_intersect(self): return
     #
     #
     # def set_union(self): return
 
+    @staticmethod
+    def set_has_element(array: ArrayType, value: ExpressionType) -> BooleanExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/sethaselementfunction
+        """
+        return BooleanExpression(KQL(f'set_has_element({to_kql(array)}, {to_kql(value)})'))
+
+    @staticmethod
+    def set_difference(array1: ArrayType, array2: ArrayType, *more_arrays: ArrayType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/setdifferencefunction
+        """
+        return ArrayExpression(KQL(f'set_difference({to_kql(array1)}, {", ".join(to_kql(a) for a in chain([array2], more_arrays))})'))
+
+    @staticmethod
+    def set_intersect(array1: ArrayType, array2: ArrayType, *more_arrays: ArrayType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/setintersectfunction
+        """
+        return ArrayExpression(KQL(f'set_intersect({to_kql(array1)}, {", ".join(to_kql(a) for a in chain([array2], more_arrays))})'))
+
+    @staticmethod
+    def set_union(array1: ArrayType, array2: ArrayType, *more_arrays: ArrayType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/setunionfunction
+        """
+        return ArrayExpression(KQL(f'set_union({to_kql(array1)}, {", ".join(to_kql(a) for a in chain([array2], more_arrays))})'))
+
+    @staticmethod
+    def array_concat(array1: ArrayType, array2: ArrayType, *more_arrays: ArrayType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arrayconcatfunction
+        """
+        return ArrayExpression(KQL(f'array_concat({to_kql(array1)}, {", ".join(to_kql(a) for a in chain([array2], more_arrays))})'))
+
+    @staticmethod
+    def array_iif(condition_array: ArrayType, if_true: ArrayType, if_false: ArrayType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arrayifffunction
+        """
+        return ArrayExpression(KQL(f'array_iif({to_kql(condition_array)}, {to_kql(if_true)}, {to_kql(if_false)})'))
+
+    @staticmethod
+    def array_index_of(array: ArrayType, value: ExpressionType) -> NumberExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arrayindexoffunction
+        """
+        return NumberExpression(KQL(f'array_index_of({to_kql(array)}, {to_kql(value)})'))
+
+    @staticmethod
+    def array_rotate_left(array: ArrayType, rotate_count: NumberType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/array_rotate_leftfunction
+        """
+        return ArrayExpression(KQL(f'array_rotate_left({to_kql(array)}, {to_kql(rotate_count)})'))
+
+    @staticmethod
+    def array_rotate_right(array: ArrayType, rotate_count: NumberType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/array_rotate_rightfunction
+        """
+        return ArrayExpression(KQL(f'array_rotate_right({to_kql(array)}, {to_kql(rotate_count)})'))
+
+    @staticmethod
+    def array_shift_left(array: ArrayType, shift_count: NumberType, fill_value: ExpressionType = None) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/array_shift_leftfunction
+        """
+        return ArrayExpression(KQL(f'array_shift_left({to_kql(array)}, {to_kql(shift_count)}{"" if fill_value is None else ", " + to_kql(fill_value)})'))
+
+    @staticmethod
+    def array_shift_right(array: ArrayType, shift_count: NumberType, fill_value: ExpressionType = None) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/array_shift_rightfunction
+        """
+        return ArrayExpression(KQL(f'array_shift_right({to_kql(array)}, {to_kql(shift_count)}{"" if fill_value is None else ", " + to_kql(fill_value)})'))
+
+    @staticmethod
+    def array_slice(array: ArrayType, start: NumberType, end: NumberType) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arrayslicefunction
+        """
+        return ArrayExpression(KQL(f'array_slice({to_kql(array)}, {to_kql(start)}, {to_kql(end)})'))
+
+    @staticmethod
+    def array_split(array: ArrayType, indices: Union[NumberType, ArrayType]) -> ArrayExpression:
+        """
+        https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arraysplitfunction
+        """
+        return ArrayExpression(KQL(f'array_split({to_kql(array)}, {to_kql(indices)})'))
+
     @staticmethod
     def sign(expr: NumberType) -> NumberExpression:
         """

test/test_functions.py

@@ -827,3 +827,81 @@ class TestFunction(TestBase):
             ' | extend foo = pack_array(numField, stringField)',
             Query().extend(foo=f.pack_array(t.numField, t.stringField)).render()
         )
+
+    def test_set_has_element(self):
+        self.assertEqual(
+            ' | extend foo = set_has_element(arrayField, numField)',
+            Query().extend(foo=f.set_has_element(t.arrayField, t.numField)).render()
+        )
+
+    def test_set_difference(self):
+        self.assertEqual(
+            ' | extend foo = set_difference(arrayField, arrayField2)',
+            Query().extend(foo=f.set_difference(t.arrayField, t.arrayField2)).render()
+        )
+
+    def test_set_intersect(self):
+        self.assertEqual(
+            ' | extend foo = set_intersect(arrayField, arrayField2)',
+            Query().extend(foo=f.set_intersect(t.arrayField, t.arrayField2)).render()
+        )
+
+    def test_set_union(self):
+        self.assertEqual(
+            ' | extend foo = set_union(arrayField, arrayField2)',
+            Query().extend(foo=f.set_union(t.arrayField, t.arrayField2)).render()
+        )
+
+    def test_array_concat(self):
+        self.assertEqual(
+            ' | extend foo = array_concat(arrayField, arrayField2)',
+            Query().extend(foo=f.array_concat(t.arrayField, t.arrayField2)).render()
+        )
+
+    def test_array_iif(self):
+        self.assertEqual(
+            ' | extend foo = array_iif(arrayField, arrayField2, arrayField3)',
+            Query().extend(foo=f.array_iif(t.arrayField, t.arrayField2, t.arrayField3)).render()
+        )
+
+    def test_array_index_of(self):
+        self.assertEqual(
+            ' | extend foo = array_index_of(arrayField, numField)',
+            Query().extend(foo=f.array_index_of(t.arrayField, t.numField)).render()
+        )
+
+    def test_array_rotate_left(self):
+        self.assertEqual(
+            ' | extend foo = array_rotate_left(arrayField, numField)',
+            Query().extend(foo=f.array_rotate_left(t.arrayField, t.numField)).render()
+        )
+
+    def test_array_rotate_right(self):
+        self.assertEqual(
+            ' | extend foo = array_rotate_right(arrayField, numField)',
+            Query().extend(foo=f.array_rotate_right(t.arrayField, t.numField)).render()
+        )
+
+    def test_array_shift_left(self):
+        self.assertEqual(
+            ' | extend foo = array_shift_left(arrayField, numField)',
+            Query().extend(foo=f.array_shift_left(t.arrayField, t.numField)).render()
+        )
+
+    def test_array_shift_right(self):
+        self.assertEqual(
+            ' | extend foo = array_shift_right(arrayField, numField, numField2)',
+            Query().extend(foo=f.array_shift_right(t.arrayField, t.numField, t.numField2)).render()
+        )
+
+    def test_array_slice(self):
+        self.assertEqual(
+            ' | extend foo = array_slice(arrayField, numField, numField2)',
+            Query().extend(foo=f.array_slice(t.arrayField, t.numField, t.numField2)).render()
+        )
+
+    def test_array_split(self):
+        self.assertEqual(
+            ' | extend foo = array_split(arrayField, numField)',
+            Query().extend(foo=f.array_split(t.arrayField, t.numField)).render()
+        )