add binary bitmap function (apache#51)

HappenLee · Aug 10, 2021 · 64aec63 · 64aec63
1 parent f959f24
commit 64aec63
Show file tree

Hide file tree

Showing 5 changed files with 105 additions and 3 deletions.
diff --git a/be/src/vec/data_types/data_type_bitmap.h b/be/src/vec/data_types/data_type_bitmap.h
@@ -27,6 +27,8 @@ class DataTypeBitMap : public IDataType {
     DataTypeBitMap() = default;
     ~DataTypeBitMap() = default;
 
+    using FieldType = BitmapValue;
+
     std::string doGetName() const override { return getFamilyName(); }
     const char* getFamilyName() const override { return "BitMap"; }
 
@@ -66,4 +68,14 @@ class DataTypeBitMap : public IDataType {
                         ErrorCodes::NOT_IMPLEMENTED);
     }
 };
+
+template <typename T>
+struct is_complex : std::false_type {};
+
+template <>
+struct is_complex<DataTypeBitMap::FieldType> : std::true_type {};
+
+template <class T>
+constexpr bool is_complex_v = is_complex<T>::value;
+
 } // namespace doris::vectorized
diff --git a/be/src/vec/functions/function_bitmap.cpp b/be/src/vec/functions/function_bitmap.cpp
@@ -53,12 +53,42 @@ struct ToBitmapImpl {
         return Status::OK();
     }
 };
-// using FunctionToBitmap = FunctionBitmap<ToBitmapImpl, NameToBitmap>;
+
+// B00LEAN BITMAP_CONTAINS(BITMAP bitmap, BIGINT input)
+// B00LEAN BITMAP_HAS_ANY(BITMAP lhs, BITMAP rhs)
+// BITMAP BITMAP_OR(BITMAP lhs, BITMAP rhs)
+// BITMAP BITMAP_XOR(BITMAP lhs, BITMAP rhs)
+// BITMAP BITMAP_NOT(BITMAP lhs, BITMAP rhs)
+// BITMAP BITMAP_XOR(BITMAP lhs, BITMAP rhs)
+
+struct NameBitmapAnd {
+    static constexpr auto name = "bitmap_and";
+};
+
+template <typename LeftDataType, typename RightDataType>
+struct BitmapAnd {
+    using ResultDataType = DataTypeBitMap;
+    using T0 = typename LeftDataType::FieldType;
+    using T1 = typename RightDataType::FieldType;
+    using TData = std::vector<BitmapValue>;
+
+    static Status vector_vector(const TData& lvec, const TData& rvec, TData& res) {
+        int size = lvec.size();
+        for (int i = 0; i < size; ++i) {
+            res[i] = lvec[i];
+            res[i] &= rvec[i];
+        }
+        return Status::OK();
+    }
+};
 
 using FunctionToBitmap = FunctionUnaryToType<ToBitmapImpl, NameToBitmap>;
+using FunctionBitmapAnd =
+        FunctionBinaryToType<DataTypeBitMap, DataTypeBitMap, BitmapAnd, NameBitmapAnd>;
 
 void registerFunctionBitmap(SimpleFunctionFactory& factory) {
     factory.registerFunction<FunctionToBitmap>();
+    factory.registerFunction<FunctionBitmapAnd>();
 }
 
 } // namespace doris::vectorized
diff --git a/be/src/vec/functions/function_totype.h b/be/src/vec/functions/function_totype.h
@@ -22,6 +22,8 @@
 #include "vec/columns/column_vector.h"
 #include "vec/data_types/data_type.h"
 #include "vec/data_types/data_type_bitmap.h"
+#include "vec/data_types/data_type_number.h"
+#include "vec/functions/cast_type_to_either.h"
 #include "vec/functions/function.h"
 
 namespace doris::vectorized {
@@ -65,4 +67,62 @@ class FunctionUnaryToType : public IFunction {
                                                 getName()));
     }
 };
+
+template <typename LeftDataType,typename RightDataType, template <typename, typename> typename Impl, typename Name>
+class FunctionBinaryToType : public IFunction {
+
+public:
+    static constexpr auto name = Name::name;
+    static FunctionPtr create() { return std::make_shared<FunctionBinaryToType>(); }
+    String getName() const override { return name; }
+    size_t getNumberOfArguments() const override { return 2; }
+    DataTypePtr getReturnTypeImpl(const DataTypes& arguments) const override {
+        using ResultDataType = typename Impl<LeftDataType,RightDataType>::ResultDataType;
+        return std::make_shared<ResultDataType>();
+    }
+
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    Status executeImpl(Block& block, const ColumnNumbers& arguments, size_t result,
+                       size_t /*input_rows_count*/) override {
+        DCHECK_EQ(arguments.size(), 2);
+        const auto& left = block.getByPosition(arguments[0]);
+        const auto& right = block.getByPosition(arguments[1]);
+        const auto& left_type = left.type;
+        const auto& right_type = left.type;
+
+        using ResultDataType = typename Impl<LeftDataType, RightDataType>::ResultDataType;
+
+        using T0 = typename LeftDataType::FieldType;
+        using T1 = typename RightDataType::FieldType;
+        using ResultType = typename ResultDataType::FieldType;
+
+        using ColVecLeft = std::conditional_t<is_complex_v<T0>, ColumnComplexType<T0>,
+                                                            ColumnVector<T0>>;
+        using ColVecRight = std::conditional_t<is_complex_v<T1>, ColumnComplexType<T1>,
+                                                            ColumnVector<T1>>;
+
+        using ColVecResult = std::conditional_t<is_complex_v<ResultType>,
+                                                                ColumnComplexType<ResultType>,
+                                                                ColumnVector<ResultType>>;
+
+
+        typename ColVecResult::MutablePtr col_res = nullptr;
+
+        col_res = ColVecResult::create();
+        auto& vec_res = col_res->getData();
+        vec_res.resize(block.rows());
+
+        if (auto col_left = checkAndGetColumn<ColVecLeft>(left.column.get())) {
+            if (auto col_right = checkAndGetColumn<ColVecRight>(right.column.get())) {
+                  Impl<LeftDataType, RightDataType>::vector_vector(col_left->getData(), col_right->getData(),
+                                              vec_res);
+                  block.getByPosition(result).column = std::move(col_res);
+            }
+        }
+
+        return Status::OK();
+    }
+};
+
 } // namespace doris::vectorized
diff --git a/be/src/vec/sink/result_sink.cpp b/be/src/vec/sink/result_sink.cpp
@@ -101,7 +101,7 @@ Status ResultSink::send(RuntimeState* state, Block* block) {
 }
 
 Status ResultSink::close(RuntimeState* state, Status exec_status) {
-    if (_closed) {
+    if (_closed || _writer == nullptr|| _sender == nullptr) {
         return Status::OK();
     }
 

diff --git a/gensrc/script/doris_builtins_functions.py b/gensrc/script/doris_builtins_functions.py
@@ -813,7 +813,7 @@
     [['bitmap_not'], 'BITMAP', ['BITMAP','BITMAP'],
         '_ZN5doris15BitmapFunctions10bitmap_notEPN9doris_udf15FunctionContextERKNS1_9StringValES6_'],
     [['bitmap_and'], 'BITMAP', ['BITMAP','BITMAP'],
-        '_ZN5doris15BitmapFunctions10bitmap_andEPN9doris_udf15FunctionContextERKNS1_9StringValES6_'],
+        '_ZN5doris15BitmapFunctions10bitmap_andEPN9doris_udf15FunctionContextERKNS1_9StringValES6_', 'vec'],
     [['bitmap_to_string'], 'VARCHAR', ['BITMAP'],
         '_ZN5doris15BitmapFunctions16bitmap_to_stringEPN9doris_udf15FunctionContextERKNS1_9StringValE'],
     [['bitmap_from_string'], 'BITMAP', ['VARCHAR'],