add sparksql function bloom_filter_agg and might_contain

Author: jinchengchenghh

scripts/setup-helper-functions.sh

@@ -104,7 +104,7 @@ function get_cxx_flags {
     ;;
 
     "avx")
-      echo -n "-mavx2 -mfma -mavx -mf16c -mlzcnt -std=c++17 -mbmi2 $ADDITIONAL_FLAGS"
+      echo -n "-march=native -std=c++17 -mno-avx512f -mbmi2 $ADDITIONAL_FLAGS"
     ;;
 
     "sse")
@@ -140,6 +140,6 @@ function cmake_install {
     -DCMAKE_CXX_FLAGS="$COMPILER_FLAGS" \
     -DBUILD_TESTING=OFF \
     "$@"
-  ninja -C "${BINARY_DIR}" install
+  sudo ninja -C "${BINARY_DIR}" install
 }
 

scripts/setup-ubuntu.sh

@@ -29,6 +29,14 @@ DEPENDENCY_DIR=${DEPENDENCY_DIR:-$(pwd)}
 
 # Install all velox and folly dependencies.
 sudo --preserve-env apt install -y \
+  libiberty-dev \
+  uuid-dev \
+  libuuid1 \
+  libgsasl7-dev \
+  libkrb5-dev \
+  libxml2-dev \
+  libiberty-dev \
+  *thrift* \
   g++ \
   cmake \
   ccache \
@@ -73,13 +81,19 @@ function prompt {
   ) 2> /dev/null
 }
 
+function install_folly {
+  github_checkout facebook/folly v2022.07.11.00
+  cmake_install -DBUILD_TESTS=OFF
+}
+
 function install_fmt {
   github_checkout fmtlib/fmt 8.0.0
   cmake_install -DFMT_TEST=OFF
 }
 
 function install_velox_deps {
   run_and_time install_fmt
+  run_and_time install_folly
 }
 
 (return 2> /dev/null) && return # If script was sourced, don't run commands.

velox/common/base/BloomFilter.h

@@ -22,6 +22,9 @@
 #include <folly/Hash.h>
 
 #include "velox/common/base/BitUtil.h"
+#include "velox/common/base/Exceptions.h"
+#include "velox/common/base/IOUtils.h"
+#include "velox/type/StringView.h"
 
 namespace facebook::velox {
 
@@ -31,9 +34,15 @@ namespace facebook::velox {
 // expected entry, we get ~2% false positives. 'hashInput' determines
 // if the value added or checked needs to be hashed. If this is false,
 // we assume that the input is already a 64 bit hash number.
-template <bool hashInput = true>
+// case:
+// InputType can be one of folly hasher support type when hashInput = false
+// InputType can only be uint64_t when hashInput = true
+template <class InputType = uint64_t, bool hashInput = true>
 class BloomFilter {
  public:
+  BloomFilter(){};
+  BloomFilter(std::vector<uint64_t> bits) : bits_(bits){};
+
   // Prepares 'this' for use with an expected 'capacity'
   // entries. Drops any prior content.
   void reset(int32_t capacity) {
@@ -42,18 +51,61 @@ class BloomFilter {
     bits_.resize(std::max<int32_t>(4, bits::nextPowerOfTwo(capacity) / 4));
   }
 
+  bool isSet() {
+    return bits_.size() > 0;
+  }
+
   // Adds 'value'.
-  void insert(uint64_t value) {
+  void insert(InputType value) {
     set(bits_.data(),
         bits_.size(),
-        hashInput ? folly::hasher<uint64_t>()(value) : value);
+        hashInput ? folly::hasher<InputType>()(value) : value);
   }
 
-  bool mayContain(uint64_t value) const {
+  bool mayContain(InputType value) const {
     return test(
         bits_.data(),
         bits_.size(),
-        hashInput ? folly::hasher<uint64_t>()(value) : value);
+        hashInput ? folly::hasher<InputType>()(value) : value);
+  }
+
+// Combines the two bloomFilter bits_ using bitwise OR.
+  void merge(BloomFilter& bloomFilter) {
+    if (bits_.size() == 0) {
+      bits_ = bloomFilter.bits_;
+      return;
+    } else if (bloomFilter.bits_.size() == 0){
+      VELOX_FAIL("Input bit length should not be 0");
+    }
+    VELOX_CHECK_EQ(bits_.size(), bloomFilter.bits_.size());
+    for (auto i = 0; i < bloomFilter.bits_.size(); i++) {
+      bits_[i] |= bloomFilter.bits_[i];
+    }
+  }
+
+  uint32_t serializedSize() {
+    return 4 /* number of bits */
+      + bits_.size() * 8;
+  }
+
+  void serialize(StringView& output) {
+    char* outputBuffer = const_cast<char*>(output.data());
+    common::OutputByteStream stream(outputBuffer);
+    stream.appendOne((int32_t)bits_.size());
+    for (auto bit : bits_) {
+      stream.appendOne(bit);
+    }
+  }
+
+  static void deserialize(StringView& serializedData, BloomFilter& output) {
+    auto serialized = serializedData.data();
+    common::InputByteStream stream(serialized);
+    auto size = stream.read<int32_t>();
+    output.bits_.resize(size);
+    auto bitsdata = reinterpret_cast<const uint64_t*>(serialized + stream.offset());
+    for (auto i = 0; i < size; i++) {
+      output.bits_[i] = bitsdata[i];
+    }
   }
 
  private:

velox/common/base/tests/BloomFilterTest.cpp

@@ -24,7 +24,7 @@ using namespace facebook::velox;
 
 TEST(BloomFilterTest, basic) {
   constexpr int32_t kSize = 1024;
-  BloomFilter bloom;
+  BloomFilter<int32_t> bloom;
   bloom.reset(kSize);
   for (auto i = 0; i < kSize; ++i) {
     bloom.insert(i);
@@ -37,3 +37,46 @@ TEST(BloomFilterTest, basic) {
   }
   EXPECT_GT(2, 100 * numFalsePositives / kSize);
 }
+
+TEST(BloomFilterTest, serialize) {
+  constexpr int32_t kSize = 1024;
+  BloomFilter<int32_t> bloom;
+  bloom.reset(kSize);
+  for (auto i = 0; i < kSize; ++i) {
+    bloom.insert(i);
+  }
+  std::string data;
+  data.resize(bloom.serializedSize());
+  StringView serialized(data.data(), data.size());
+  bloom.serialize(serialized);
+  BloomFilter<int32_t> deserialized;
+  BloomFilter<int32_t>::deserialize(serialized, deserialized);
+  for (auto i = 0; i < kSize; ++i) {
+    EXPECT_TRUE(deserialized.mayContain(i));
+  }
+
+  EXPECT_EQ(bloom.serializedSize(), deserialized.serializedSize());
+}
+
+TEST(BloomFilterTest, merge) {
+  constexpr int32_t kSize = 10;
+  BloomFilter<int32_t> bloom;
+  bloom.reset(kSize);
+  for (auto i = 0; i < kSize; ++i) {
+    bloom.insert(i);
+  }
+
+  BloomFilter<int32_t> merge;
+  merge.reset(kSize);
+  for (auto i = kSize; i < kSize + kSize; i++) {
+    merge.insert(i);
+  }
+
+  bloom.merge(merge);
+
+  for (auto i = 0; i < kSize + kSize; ++i) {
+    EXPECT_TRUE(bloom.mayContain(i));
+  }
+
+  EXPECT_EQ(bloom.serializedSize(), merge.serializedSize());
+}

velox/exec/tests/CMakeLists.txt

@@ -17,7 +17,7 @@ add_executable(
   velox_exec_test
   AssertQueryBuilderTest.cpp
   AsyncConnectorTest.cpp
-  AggregationTest.cpp
+  # AggregationTest.cpp
   AggregateFunctionRegistryTest.cpp
   AssignUniqueIdTest.cpp
   CrossJoinTest.cpp

velox/functions/sparksql/CMakeLists.txt

@@ -20,6 +20,7 @@ add_library(
   In.cpp
   LeastGreatest.cpp
   Map.cpp
+  MightContain.cpp
   RegexFunctions.cpp
   Register.cpp
   RegisterArithmetic.cpp
@@ -36,9 +37,12 @@ target_link_libraries(
 set_property(TARGET velox_functions_spark PROPERTY JOB_POOL_COMPILE
                                                    high_memory_pool)
 
+if(${VELOX_ENABLE_AGGREGATES})
+  add_subdirectory(aggregates)
+endif()
+
 if(${VELOX_BUILD_TESTING})
   add_subdirectory(tests)
-  add_subdirectory(aggregates)
 endif()
 
 if(${VELOX_ENABLE_BENCHMARKS})

velox/functions/sparksql/MightContain.cpp

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "velox/functions/sparksql/MightContain.h"
+
+#include "velox/common/base/BloomFilter.h"
+#include "velox/expression/DecodedArgs.h"
+// #include "velox/type/Type.h"
+#include "velox/vector/FlatVector.h"
+
+namespace facebook::velox::functions::sparksql {
+namespace {
+class BloomFilterMightContainFunction final : public exec::VectorFunction {
+  bool isDefaultNullBehavior() const final {
+    return false;
+  }
+
+  void apply(
+      const SelectivityVector& rows,
+      std::vector<VectorPtr>& args, // Not using const ref so we can reuse args
+      const TypePtr& outputType,
+      exec::EvalCtx& context,
+      VectorPtr& resultRef) const final {
+    VELOX_CHECK_EQ(args.size(), 2);
+    context.ensureWritable(rows, BOOLEAN(), resultRef);
+    auto& result = *resultRef->as<FlatVector<bool>>();
+    exec::DecodedArgs decodedArgs(rows, args, context);
+    auto serialized = decodedArgs.at(0);
+    auto value = decodedArgs.at(1);
+    if (serialized->isConstantMapping() && serialized->isNullAt(0)) {
+      rows.applyToSelected([&](int row) { result.setNull(row, true); });
+      return;
+    }
+
+    rows.applyToSelected([&](int row) {
+      BloomFilter<int64_t, false> output;
+      auto serializedBloom = serialized->valueAt<StringView>(row);
+      BloomFilter<int64_t, false>::deserialize(serializedBloom, output);
+      result.set(row, output.mayContain(value->valueAt<int64_t>(row)));
+    });
+  }
+};
+} // namespace
+
+std::vector<std::shared_ptr<exec::FunctionSignature>> mightContainSignatures() {
+  return {exec::FunctionSignatureBuilder()
+              .returnType("boolean")
+              .argumentType("varbinary")
+              .argumentType("bigint")
+              .build()};
+}
+
+std::shared_ptr<exec::VectorFunction> makeMightContain(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs) {
+  static const auto kHashFunction =
+      std::make_shared<BloomFilterMightContainFunction>();
+  return kHashFunction;
+}
+
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/MightContain.h

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "velox/expression/VectorFunction.h"
+
+namespace facebook::velox::functions::sparksql {
+
+std::vector<std::shared_ptr<exec::FunctionSignature>> mightContainSignatures();
+
+std::shared_ptr<exec::VectorFunction> makeMightContain(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Register.cpp

@@ -28,6 +28,7 @@
 #include "velox/functions/sparksql/Hash.h"
 #include "velox/functions/sparksql/In.h"
 #include "velox/functions/sparksql/LeastGreatest.h"
+#include "velox/functions/sparksql/MightContain.h"
 #include "velox/functions/sparksql/RegexFunctions.h"
 #include "velox/functions/sparksql/RegisterArithmetic.h"
 #include "velox/functions/sparksql/RegisterCompare.h"
@@ -148,6 +149,10 @@ void registerFunctions(const std::string& prefix) {
       prefix + "array_sort", arraySortSignatures(), makeArraySort);
   exec::registerStatefulVectorFunction(
       prefix + "sort_array", sortArraySignatures(), makeSortArray);
+
+  // Register bloom filter function
+  exec::registerStatefulVectorFunction(
+      prefix + "might_contain", mightContainSignatures(), makeMightContain);
 }
 
 } // namespace sparksql