add mysql compatible function bit_length

facebookincubator · Feb 18, 2025 · 209c23c · 209c23c
1 parent 730842b
commit 209c23c
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 0 deletions.
diff --git a/velox/functions/prestosql/StringFunctions.h b/velox/functions/prestosql/StringFunctions.h
@@ -24,6 +24,16 @@
 
 namespace facebook::velox::functions {
 
+/// Returns the count of bits for the given string.
+template <typename T>
+struct BitLengthFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(int64_t& result, const StringView& input) {
+    result = input.size() * 8;
+  }
+};
+
 /// chr(n) → varchar
 /// Returns the Unicode code point n as a single character string.
 template <typename T>

diff --git a/velox/functions/prestosql/registration/StringFunctionsRegistration.cpp b/velox/functions/prestosql/registration/StringFunctionsRegistration.cpp
@@ -35,6 +35,7 @@ void registerSimpleFunctions(const std::string& prefix) {
   using namespace stringImpl;
 
   // Register string functions.
+  registerFunction<BitLengthFunction, Varchar, int64_t>({prefix + "bit_length"});
   registerFunction<ChrFunction, Varchar, int64_t>({prefix + "chr"});
   registerFunction<CodePointFunction, int32_t, Varchar>({prefix + "codepoint"});
   registerFunction<HammingDistanceFunction, int64_t, Varchar, Varchar>(

diff --git a/velox/functions/prestosql/tests/StringFunctionsTest.cpp b/velox/functions/prestosql/tests/StringFunctionsTest.cpp
@@ -2229,6 +2229,14 @@ TEST_F(StringFunctionsTest, normalize) {
       "Normalization form must be one of [NFD, NFC, NFKD, NFKC]");
 }
 
+TEST_F(StringFunctionsTest, bitLength) {
+  const auto runBitLength = [&](std::optional<std::string> value) {
+    return evaluateOnce<int64_t>("bit_length(c0)", value);
+  };
+  EXPECT_EQ(8, runBitLength("3"));
+  EXPECT_EQ(40, runBitLength("ab中"));
+}
+
 TEST_F(StringFunctionsTest, trail) {
   auto trail = [&](std::optional<std::string> string,
                    std::optional<int32_t> N) {