Implement array_iff / array_iif

ClibMouse · Aug 26, 2023 · 7adc203 · 7adc203
1 parent fc1b317
commit 7adc203
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 4 deletions.
diff --git a/src/Parsers/Kusto/KQL_ReleaseNote.md b/src/Parsers/Kusto/KQL_ReleaseNote.md
@@ -1,6 +1,14 @@
 
 ## KQL implemented features
 
+# August XX, 2022
+## Dynamic functions
+- [array_iff / array_iif](https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/arrayifffunction)
+   `print array_iif(dynamic([true, false, true]), dynamic([1, 2, 3]), dynamic([4, 5, 6])) == dynamic([1, 5, 3])`
+   `print array_iif(dynamic([true, false, true]), dynamic([1, 2, 3, 4]), dynamic([4, 5, 6])) == dynamic([1, 5, 3])`
+   `print array_iif(dynamic([true, false, true, false]), dynamic([1, 2, 3, 4]), dynamic([4, 5, 6])) == dynamic([1, 5, 3, null])`
+   `print array_iif(dynamic([1, 0, -1, 44, 0]), dynamic([1, 2, 3, 4]), dynamic([4, 5, 6])) == dynamic([1, 5, 3, 4, null])`
+
 # August 15, 2022
    **double quote support**  
    ``print res = strcat("double ","quote")``  
@@ -49,7 +57,7 @@
    `print tostring(123) == '123'`  
    `print tostring('asd') == 'asd'`  
 
-## DateType
+## Data Types
  - [dynamic](https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/scalar-data-types/dynamic)  
     *Supports only 1D array*  
     `print output = dynamic(['a', 'b', 'c'])`  

diff --git a/src/Parsers/Kusto/KustoFunctions/KQLDynamicFunctions.cpp b/src/Parsers/Kusto/KustoFunctions/KQLDynamicFunctions.cpp
@@ -15,6 +15,8 @@
 #include <Parsers/Kusto/KustoFunctions/KQLBinaryFunctions.h>
 #include <Parsers/Kusto/KustoFunctions/KQLGeneralFunctions.h>
 
+#include <format>
+
 namespace DB
 {
 
@@ -25,11 +27,23 @@ bool ArrayConcat::convertImpl(String &out,IParser::Pos &pos)
     return false;
 }
 
-bool ArrayIif::convertImpl(String &out,IParser::Pos &pos)
+bool ArrayIif::convertImpl(String & out, IParser::Pos & pos)
 {
-    String res = String(pos->begin,pos->end);
-    out = res;
+    const auto function_name = getKQLFunctionName(pos);
+    if (function_name.empty())
     return false;
+
+    const auto conditions = getArgument(function_name, pos);
+    const auto if_true = getArgument(function_name, pos);
+    const auto if_false = getArgument(function_name, pos);
+
+    out = std::format(
+        "arrayMap(x -> if(x.1 != 0, x.2, x.3), arrayZip({0}, arrayResize({1}, length({0}), null), arrayResize({2}, length({0}), null)))",
+        conditions,
+        if_true,
+        if_false);
+
+    return true;
 }
 
 bool ArrayIndexOf::convertImpl(String &out,IParser::Pos &pos)

diff --git a/src/Parsers/Kusto/KustoFunctions/KQLFunctionFactory.cpp b/src/Parsers/Kusto/KustoFunctions/KQLFunctionFactory.cpp
@@ -93,6 +93,7 @@ namespace DB
         {"url_encode", KQLFunctionValue::url_encode},
 
         {"array_concat", KQLFunctionValue::array_concat},
+        {"array_iff", KQLFunctionValue::array_iif},
         {"array_iif", KQLFunctionValue::array_iif},
         {"array_index_of", KQLFunctionValue::array_index_of},
         {"array_length", KQLFunctionValue::array_length},

diff --git a/src/Parsers/tests/KQL/gtest_KQL_Dynamic.cpp b/src/Parsers/tests/KQL/gtest_KQL_Dynamic.cpp
@@ -0,0 +1,17 @@
+#include <Parsers/tests/gtest_common.h>
+
+#include <Parsers/Kusto/ParserKQLQuery.h>
+
+INSTANTIATE_TEST_SUITE_P(ParserKQLQuery_Dynamic, ParserTest,
+    ::testing::Combine(
+        ::testing::Values(std::make_shared<DB::ParserKQLQuery>()),
+        ::testing::ValuesIn(std::initializer_list<ParserTestCase>{
+        {
+            "print array_iff(A, B, C)",
+            "SELECT arrayMap(x -> if((x.1) != 0, x.2, x.3), arrayZip(A, arrayResize(B, length(A), NULL), arrayResize(C, length(A), NULL)))"
+        },
+        {
+            "print array_iif(A, B, C)",
+            "SELECT arrayMap(x -> if((x.1) != 0, x.2, x.3), arrayZip(A, arrayResize(B, length(A), NULL), arrayResize(C, length(A), NULL)))"
+        }
+})));