diff --git a/bolt/functions/sparksql/Bitmap.h b/bolt/functions/sparksql/Bitmap.h new file mode 100644 index 000000000..ac575ce2c --- /dev/null +++ b/bolt/functions/sparksql/Bitmap.h @@ -0,0 +1,42 @@ +/* + * Copyright (c) ByteDance Ltd. and/or 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. + */ + +#pragma once + +#include + +#include + +namespace bytedance::bolt::functions::sparksql { + +inline constexpr int64_t kBitmapNumBits = 32768; + +template +struct BitmapBitPositionFunction { + FOLLY_ALWAYS_INLINE void call(int64_t& result, int64_t value) { + if (value > 0) { + result = (value - 1) % kBitmapNumBits; + return; + } + + // Calculate the magnitude as unsigned to avoid signed overflow for + // INT64_MIN while preserving Java's two's-complement behavior. + const auto magnitude = uint64_t{0} - static_cast(value); + result = static_cast(magnitude % kBitmapNumBits); + } +}; + +} // namespace bytedance::bolt::functions::sparksql diff --git a/bolt/functions/sparksql/registration/CMakeLists.txt b/bolt/functions/sparksql/registration/CMakeLists.txt index af9a09e1b..41199105d 100644 --- a/bolt/functions/sparksql/registration/CMakeLists.txt +++ b/bolt/functions/sparksql/registration/CMakeLists.txt @@ -31,6 +31,7 @@ bolt_add_library( bolt_functions_spark Register.cpp RegisterArray.cpp + RegisterBitmap.cpp RegisterBinary.cpp RegisterBitwise.cpp RegisterCharVarchar.cpp diff --git a/bolt/functions/sparksql/registration/Register.cpp b/bolt/functions/sparksql/registration/Register.cpp index d681f8200..30297c612 100644 --- a/bolt/functions/sparksql/registration/Register.cpp +++ b/bolt/functions/sparksql/registration/Register.cpp @@ -35,6 +35,7 @@ namespace bytedance::bolt::functions::sparksql { extern void registerArrayFunctions(const std::string& prefix); +extern void registerBitmapFunctions(const std::string& prefix); extern void registerBinaryFunctions(const std::string& prefix); extern void registerBitwiseFunctions(const std::string& prefix); extern void registerCharVarcharFunctions(const std::string& prefix); @@ -53,6 +54,7 @@ extern void registerVariantFunctions(const std::string& prefix); void registerFunctions(const std::string& prefix) { registerTimestampWithTimeZoneType(); registerArrayFunctions(prefix); + registerBitmapFunctions(prefix); registerBinaryFunctions(prefix); registerBitwiseFunctions(prefix); registerCharVarcharFunctions(prefix); diff --git a/bolt/functions/sparksql/registration/RegisterBitmap.cpp b/bolt/functions/sparksql/registration/RegisterBitmap.cpp new file mode 100644 index 000000000..606688fe2 --- /dev/null +++ b/bolt/functions/sparksql/registration/RegisterBitmap.cpp @@ -0,0 +1,27 @@ +/* + * Copyright (c) ByteDance Ltd. and/or 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 "bolt/functions/lib/RegistrationHelpers.h" +#include "bolt/functions/sparksql/Bitmap.h" + +namespace bytedance::bolt::functions::sparksql { + +void registerBitmapFunctions(const std::string& prefix) { + registerFunction( + {prefix + "bitmap_bit_position"}); +} + +} // namespace bytedance::bolt::functions::sparksql diff --git a/bolt/functions/sparksql/tests/BitmapTest.cpp b/bolt/functions/sparksql/tests/BitmapTest.cpp new file mode 100644 index 000000000..055edbfba --- /dev/null +++ b/bolt/functions/sparksql/tests/BitmapTest.cpp @@ -0,0 +1,65 @@ +/* + * Copyright (c) ByteDance Ltd. and/or 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 +#include +#include + +#include "bolt/functions/sparksql/tests/SparkFunctionBaseTest.h" + +namespace bytedance::bolt::functions::sparksql::test { +namespace { + +class BitmapTest : public SparkFunctionBaseTest { + protected: + std::optional bitPosition(std::optional value) { + return evaluateOnce("bitmap_bit_position(c0)", value); + } + + std::optional bitPositionFromInt(std::optional value) { + return evaluateOnce( + "bitmap_bit_position(cast(c0 as bigint))", value); + } +}; + +TEST_F(BitmapTest, bitPositionPositive) { + EXPECT_EQ(bitPosition(1), 0); + EXPECT_EQ(bitPosition(2), 1); + EXPECT_EQ(bitPosition(32768), 32767); + EXPECT_EQ(bitPosition(32769), 0); + EXPECT_EQ(bitPosition(65536), 32767); + EXPECT_EQ(bitPosition(65537), 0); + EXPECT_EQ(bitPosition(std::numeric_limits::max()), 32766); +} + +TEST_F(BitmapTest, bitPositionNonPositive) { + EXPECT_EQ(bitPosition(0), 0); + EXPECT_EQ(bitPosition(-1), 1); + EXPECT_EQ(bitPosition(-32767), 32767); + EXPECT_EQ(bitPosition(-32768), 0); + EXPECT_EQ(bitPosition(-32769), 1); + EXPECT_EQ(bitPosition(std::numeric_limits::min()), 0); + EXPECT_EQ(bitPosition(std::numeric_limits::min() + 1), 32767); +} + +TEST_F(BitmapTest, bitPositionNullAndCast) { + EXPECT_EQ(bitPosition(std::nullopt), std::nullopt); + EXPECT_EQ(bitPositionFromInt(32769), 0); + EXPECT_EQ(bitPositionFromInt(std::nullopt), std::nullopt); +} + +} // namespace +} // namespace bytedance::bolt::functions::sparksql::test diff --git a/bolt/functions/sparksql/tests/CMakeLists.txt b/bolt/functions/sparksql/tests/CMakeLists.txt index 4e4593c35..7415a5fac 100644 --- a/bolt/functions/sparksql/tests/CMakeLists.txt +++ b/bolt/functions/sparksql/tests/CMakeLists.txt @@ -43,6 +43,7 @@ add_executable( ArraysZipTest.cpp AtLeastNNonNullsTest.cpp Base64Test.cpp + BitmapTest.cpp BitwiseTest.cpp CharVarcharTest.cpp ComparisonsTest.cpp