feat(udf): isin (#3939)

4paradigm · Jun 26, 2024 · 818d292 · 818d292
1 parent e307fd9
commit 818d292
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 2 deletions.
diff --git 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
@@ -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)