From 818d292cc5190c65ef3ead678f6742e828d0f163 Mon Sep 17 00:00:00 2001
From: aceforeverd <teapot@aceforeverd.com>
Date: Wed, 26 Jun 2024 12:34:11 +0800
Subject: [PATCH] feat(udf): isin (#3939)

---
 cases/query/udf_query.yaml                 | 19 ++++++++++
 hybridse/src/udf/default_defs/array_def.cc | 41 ++++++++++++++++++++--
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/cases/query/udf_query.yaml b/cases/query/udf_query.yaml
index fefe1380dbb..ee9cad2d667 100644
--- a/cases/query/udf_query.yaml
+++ b/cases/query/udf_query.yaml
@@ -536,6 +536,25 @@ cases:
       data: |
         true, true, false, false, true, false, true, false, true, false, true
 
+  - id: isin
+    mode: request-unsupport
+    inputs:
+      - name: t1
+        columns: ["col1:int32", "std_ts:timestamp", "col2:string"]
+        indexs: ["index1:col1:std_ts"]
+        rows:
+          - [1, 1590115420001, "ABCabcabc"]
+    sql: |
+      select
+        isin(2, [2,2]) as c0,
+        isin(cast(3 as int64), ARRAY<INT64>[NULL, 1, 2]) as c1
+    expect:
+      columns:
+        - c0 bool
+        - c1 bool
+      data: |
+        true, false
+
   - id: array_split
     mode: request-unsupport
     inputs:
diff --git a/hybridse/src/udf/default_defs/array_def.cc b/hybridse/src/udf/default_defs/array_def.cc
index a1f35f38e35..b5c36bc3d7e 100644
--- a/hybridse/src/udf/default_defs/array_def.cc
+++ b/hybridse/src/udf/default_defs/array_def.cc
@@ -37,8 +37,30 @@ struct ArrayContains {
     // - bool/intxx/float/double -> bool/intxx/float/double
     // - Timestamp/Date/StringRef -> Timestamp*/Date*/StringRef*
     bool operator()(ArrayRef<T>* arr, ParamType v, bool is_null) {
-        // NOTE: array_contains([null], null) returns null
-        // this might not expected
+        for (uint64_t i = 0; i < arr->size; ++i) {
+            if constexpr (std::is_pointer_v<ParamType>) {
+                // null or same value returns true
+                if ((is_null && arr->nullables[i]) || (!arr->nullables[i] && *arr->raw[i] == *v)) {
+                    return true;
+                }
+            } else {
+                if ((is_null && arr->nullables[i]) || (!arr->nullables[i] && arr->raw[i] == v)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+};
+
+template <typename T>
+struct IsIn {
+    // udf registry types
+    using Args = std::tuple<Nullable<T>, ArrayRef<T>>;
+
+    using ParamType = typename DataTypeTrait<T>::CCallArgType;
+
+    bool operator()(ParamType v, bool is_null, ArrayRef<T>* arr) {
         for (uint64_t i = 0; i < arr->size; ++i) {
             if constexpr (std::is_pointer_v<ParamType>) {
                 // null or same value returns true
@@ -98,6 +120,21 @@ void DefaultUdfLibrary::InitArrayUdfs() {
              @since 0.7.0
              )");
 
+    RegisterExternalTemplate<IsIn>("isin")
+        .args_in<bool, int16_t, int32_t, int64_t, float, double, Timestamp, Date, StringRef>()
+        .doc(R"(
+             @brief isin(value, array) - Returns true if the array contains the value.
+
+             Example:
+
+             @code{.sql}
+                 select isin(2, [2,2]) as c0;
+                 -- output true
+             @endcode
+
+             @since 0.9.1
+             )");
+
     RegisterExternal("split_array")
         .returns<ArrayRef<StringRef>>()
         .return_by_arg(true)