[Enhancement][Refactor] optimize array_contains_all/array_contains_se…

…q function (StarRocks#51701) Signed-off-by: silverbullet233 <[email protected]>
ZiheLiu · Oct 31, 2024 · 89b3ed0 · 89b3ed0
1 parent 8c7b8f2
commit 89b3ed0
Show file tree

Hide file tree

Showing 10 changed files with 1,033 additions and 3 deletions.
diff --git a/be/src/column/column_helper.cpp b/be/src/column/column_helper.cpp
@@ -437,6 +437,17 @@ ColumnPtr ColumnHelper::convert_time_column_from_double_to_str(const ColumnPtr&
     return res;
 }
 
+std::tuple<UInt32Column::Ptr, ColumnPtr, NullColumnPtr> ColumnHelper::unpack_array_column(const ColumnPtr& column) {
+    DCHECK(!column->is_nullable() && !column->is_constant());
+    DCHECK(column->is_array());
+
+    const ArrayColumn* array_column = down_cast<ArrayColumn*>(column.get());
+    auto elements_column = down_cast<NullableColumn*>(array_column->elements_column().get())->data_column();
+    auto null_column = down_cast<NullableColumn*>(array_column->elements_column().get())->null_column();
+    auto offsets_column = array_column->offsets_column();
+    return {offsets_column, elements_column, null_column};
+}
+
 template <class Ptr>
 bool ChunkSliceTemplate<Ptr>::empty() const {
     return !chunk || offset == chunk->num_rows();

diff --git a/be/src/column/column_helper.h b/be/src/column/column_helper.h
@@ -530,6 +530,9 @@ class ColumnHelper {
     static ColumnPtr create_const_null_column(size_t chunk_size);
 
     static ColumnPtr convert_time_column_from_double_to_str(const ColumnPtr& column);
+
+    // unpack array column, return offsets_column, elements_column, elements_null_column
+    static std::tuple<UInt32Column::Ptr, ColumnPtr, NullColumnPtr> unpack_array_column(const ColumnPtr& column);
 };
 
 // Hold a slice of chunk

diff --git a/be/src/exprs/array_functions.h b/be/src/exprs/array_functions.h
@@ -151,12 +151,44 @@ class ArrayFunctions {
     DEFINE_VECTORIZED_FN(array_cum_sum_double);
 
     DEFINE_VECTORIZED_FN(array_contains_any);
+
     DEFINE_VECTORIZED_FN(array_contains_all);
+
+    template <LogicalType LT>
+    static StatusOr<ColumnPtr> array_contains_all_specific(FunctionContext* context, const Columns& columns) {
+        return ArrayContainsAll<LT, false>::process(context, columns);
+    }
+    template <LogicalType LT>
+    static Status array_contains_all_specific_prepare(FunctionContext* context,
+                                                      FunctionContext::FunctionStateScope scope) {
+        return ArrayContainsAll<LT, false>::prepare(context, scope);
+    }
+    template <LogicalType LT>
+    static Status array_contains_all_specific_close(FunctionContext* context,
+                                                    FunctionContext::FunctionStateScope scope) {
+        return ArrayContainsAll<LT, false>::close(context, scope);
+    }
+
     DEFINE_VECTORIZED_FN(array_map);
     DEFINE_VECTORIZED_FN(array_filter);
     DEFINE_VECTORIZED_FN(all_match);
     DEFINE_VECTORIZED_FN(any_match);
+
     DEFINE_VECTORIZED_FN(array_contains_seq);
+    template <LogicalType LT>
+    static StatusOr<ColumnPtr> array_contains_seq_specific(FunctionContext* context, const Columns& columns) {
+        return ArrayContainsAll<LT, true>::process(context, columns);
+    }
+    template <LogicalType LT>
+    static Status array_contains_seq_specific_prepare(FunctionContext* context,
+                                                      FunctionContext::FunctionStateScope scope) {
+        return ArrayContainsAll<LT, true>::prepare(context, scope);
+    }
+    template <LogicalType LT>
+    static Status array_contains_seq_specific_close(FunctionContext* context,
+                                                    FunctionContext::FunctionStateScope scope) {
+        return ArrayContainsAll<LT, true>::close(context, scope);
+    }
 
     // array function for nested type(Array/Map/Struct)
     DEFINE_VECTORIZED_FN(array_distinct_any_type);