feat: Add Spark varchar_type_write_side_check function (facebookincubator#13034)

Summary:
This PR adds the `varchar_type_write_side_check` function. This function trims
trailing space characters (ASCII 32) from input string and checks if it fits
within the specified length limit, throwing exceptions if the string exceeds
the limit after trimming or if the specified limit is not greater than 0.

Examples:

```
varchar_type_write_side_check("abc   ", 3) -- "abc"
varchar_type_write_side_check("abcd", 3) -- VeloxUserError: "Exceeds allowed length limitation: '3'"
varchar_type_write_side_check("   ", 0) -- VeloxUserError: "The length limit must be greater than 0."
```

**Note:** This PR is a splitted work of PR facebookincubator#12773, as outlined in PR facebookincubator#12772.

Pull Request resolved: facebookincubator#13034

Reviewed By: Yuhta

Differential Revision: D75963922

Pulled By: kKPulla

fbshipit-source-id: 34bba67e95645ab45f2493fd30f6635a6b3cf9a0

Author: Yifeng-Wang

velox/docs/functions/spark/string.rst

@@ -438,3 +438,20 @@ String Functions
     Returns string with all characters changed to uppercase. ::
 
         SELECT upper('SparkSql'); -- SPARKSQL
+
+.. spark:function:: varchar_type_write_side_check(string, limit) -> varchar
+
+    Removes trailing space characters (ASCII 32) that exceed the length ``limit`` from the end of input ``string``. ``limit`` is the maximum length of characters that can be allowed.
+    Throws exception when ``string`` still exceeds ``limit`` after trimming trailing spaces or when ``limit`` is not greater than 0.
+    Empty strings are returned as-is since they always satisfy any length ``limit`` greater than 0.
+    Note: This function is not directly callable in Spark SQL, but internally used for length check when writing string type columns. ::
+
+        -- Function call examples (this function is not directly callable in Spark SQL).
+        varchar_type_write_side_check("abc", 3) -- "abc"
+        varchar_type_write_side_check("abc   ", 3) -- "abc"
+        varchar_type_write_side_check("abcd", 3) -- VeloxUserError: "Exceeds allowed length limitation: '3'"
+        varchar_type_write_side_check("中国", 3) -- "中国"
+        varchar_type_write_side_check("中文中国", 3) -- VeloxUserError: "Exceeds allowed length limitation: '3'"
+        varchar_type_write_side_check("   ", 0) -- VeloxUserError: "The length limit must be greater than 0."
+        varchar_type_write_side_check("", 3) -- ""
+

velox/functions/sparksql/CMakeLists.txt

@@ -17,6 +17,7 @@ velox_add_library(
   velox_functions_spark_impl
   ArrayGetFunction.cpp
   ArraySort.cpp
+  CharVarcharUtils.cpp
   Comparisons.cpp
   ConcatWs.cpp
   DecimalArithmetic.cpp

velox/functions/sparksql/CharVarcharUtils.cpp

@@ -0,0 +1,43 @@
+/*
+ * 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/CharVarcharUtils.h"
+
+namespace facebook::velox::functions::sparksql {
+
+void trimTrailingSpaces(
+    exec::StringWriter& output,
+    StringView inputStr,
+    int32_t numChars,
+    uint32_t limit) {
+  const auto numTailSpacesToTrim = numChars - limit;
+  VELOX_USER_CHECK_GT(numTailSpacesToTrim, 0);
+
+  auto curPos = inputStr.end() - 1;
+  const auto trimTo = inputStr.end() - numTailSpacesToTrim;
+
+  while (curPos >= trimTo && stringImpl::isAsciiSpace(*curPos)) {
+    curPos--;
+  }
+  // Get the length of the trimmed string in characters.
+  const auto trimmedSize = numChars - std::distance(curPos + 1, inputStr.end());
+
+  VELOX_USER_CHECK_LE(
+      trimmedSize, limit, "Exceeds allowed length limitation: {}", limit);
+  output.setNoCopy(
+      StringView(inputStr.data(), std::distance(inputStr.begin(), curPos + 1)));
+}
+
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/CharVarcharUtils.h

@@ -0,0 +1,35 @@
+/*
+ * 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.
+ */
+#pragma once
+
+#include <string>
+
+#include "velox/expression/StringWriter.h"
+#include "velox/functions/lib/string/StringImpl.h"
+
+namespace facebook::velox::functions::sparksql {
+
+/// Trims trailing ASCII space characters (0x20) from `'abc'`
+/// to ensure its length does not exceed the specified Unicode string length
+/// `limit` (must be greater than 0) in characters. Throws an exception if the
+/// string still exceeds `limit` after trimming.
+void trimTrailingSpaces(
+    exec::StringWriter& output,
+    StringView inputStr,
+    int32_t numChars,
+    uint32_t limit);
+
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/VarcharTypeWriteSideCheck.h

@@ -0,0 +1,74 @@
+/*
+ * 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.
+ */
+#pragma once
+
+#include <string>
+
+#include "velox/functions/Macros.h"
+#include "velox/functions/lib/string/StringImpl.h"
+#include "velox/functions/sparksql/CharVarcharUtils.h"
+
+namespace facebook::velox::functions::sparksql {
+
+/// Ensures that `'abc'` fits within the specified length `n` in
+/// characters. If the length of `'abc'` exceeds `n`, trailing spaces
+/// are trimmed to fit within `n`. If the length of `'abc'` is less than
+/// or equal to `n`, it is returned as-is. Throws an exception if the
+/// trimmed string still exceeds `n` or if `n` is negative. This function
+/// will trim at most (length of `'abc'` - `n`) space characters (ASCII 32)
+/// from the end of `'abc'`.
+template <typename T>
+struct VarcharTypeWriteSideCheckFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  // Results refer to strings in the first argument.
+  static constexpr int32_t reuse_strings_from_arg = 0;
+
+  // ASCII input always produces ASCII result.
+  static constexpr bool is_default_ascii_behavior = true;
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& input,
+      int32_t limit) {
+    doCall<false>(result, input, limit);
+  }
+
+  FOLLY_ALWAYS_INLINE void callAscii(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& input,
+      int32_t limit) {
+    doCall<true>(result, input, limit);
+  }
+
+ private:
+  template <bool isAscii>
+  FOLLY_ALWAYS_INLINE void doCall(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& input,
+      int32_t limit) {
+    VELOX_USER_CHECK_GT(limit, 0, "The length limit must be greater than 0.");
+
+    auto numCharacters = stringImpl::length<isAscii>(input);
+    if (numCharacters <= limit) {
+      result.setNoCopy(input);
+    } else {
+      trimTrailingSpaces(result, input, numCharacters, limit);
+    }
+  }
+};
+
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/registration/RegisterString.cpp

@@ -24,6 +24,7 @@
 #include "velox/functions/sparksql/Split.h"
 #include "velox/functions/sparksql/String.h"
 #include "velox/functions/sparksql/StringToMap.h"
+#include "velox/functions/sparksql/VarcharTypeWriteSideCheck.h"
 
 namespace facebook::velox::functions {
 void registerSparkStringFunctions(const std::string& prefix) {
@@ -161,6 +162,11 @@ void registerStringFunctions(const std::string& prefix) {
       prefix + "lower",
       SparkLowerFunction::signatures(),
       std::make_unique<SparkLowerFunction>());
+  registerFunction<
+      VarcharTypeWriteSideCheckFunction,
+      Varchar,
+      Varchar,
+      int32_t>({prefix + "varchar_type_write_side_check"});
 }
 } // namespace sparksql
 } // namespace facebook::velox::functions

velox/functions/sparksql/tests/CMakeLists.txt

@@ -68,6 +68,7 @@ add_executable(
   UnscaledValueFunctionTest.cpp
   UpperLowerTest.cpp
   UuidTest.cpp
+  VarcharTypeWriteSideCheckTest.cpp
   XxHash64Test.cpp)
 
 add_test(velox_functions_spark_test velox_functions_spark_test)

velox/functions/sparksql/tests/StringTest.cpp

@@ -1055,5 +1055,6 @@ TEST_F(StringTest, empty2Null) {
   EXPECT_EQ(empty2Null(""), std::nullopt);
   EXPECT_EQ(empty2Null("abc"), "abc");
 }
+
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test

velox/functions/sparksql/tests/VarcharTypeWriteSideCheckTest.cpp

@@ -0,0 +1,76 @@
+/*
+ * 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/common/base/tests/GTestUtils.h"
+#include "velox/functions/sparksql/tests/SparkFunctionBaseTest.h"
+
+namespace facebook::velox::functions::sparksql::test {
+namespace {
+
+class VarcharTypeWriteSideCheckTest : public SparkFunctionBaseTest {};
+
+TEST_F(VarcharTypeWriteSideCheckTest, varcharTypeWriteSideCheck) {
+  const auto varcharTypeWriteSideCheck =
+      [&](const std::optional<std::string>& input,
+          const std::optional<int32_t>& limit) {
+        return evaluateOnce<std::string>(
+            "varchar_type_write_side_check(c0, c1)", input, limit);
+      };
+
+  // Basic cases - string length <= limit.
+  EXPECT_EQ(varcharTypeWriteSideCheck("abc", 3), "abc");
+  EXPECT_EQ(varcharTypeWriteSideCheck("ab", 3), "ab");
+  EXPECT_EQ(varcharTypeWriteSideCheck("", 5), "");
+
+  // Cases with trailing spaces.
+  // Edge cases - input string is longer than limit but trims to exactly limit.
+  EXPECT_EQ(varcharTypeWriteSideCheck("abc  ", 3), "abc");
+  EXPECT_EQ(varcharTypeWriteSideCheck("abc  ", 4), "abc ");
+  EXPECT_EQ(varcharTypeWriteSideCheck("abc  ", 5), "abc  ");
+
+  // Unicode string cases with trailing spaces.
+  EXPECT_EQ(varcharTypeWriteSideCheck("世界   ", 2), "世界");
+  EXPECT_EQ(varcharTypeWriteSideCheck("世界", 2), "世界");
+
+  // Error cases - string length > limit even after trimming trailing spaces.
+  VELOX_ASSERT_USER_THROW(
+      varcharTypeWriteSideCheck("abcd", 3),
+      "Exceeds allowed length limitation: 3");
+  VELOX_ASSERT_USER_THROW(
+      varcharTypeWriteSideCheck("世界人", 2),
+      "Exceeds allowed length limitation: 2");
+  VELOX_ASSERT_USER_THROW(
+      varcharTypeWriteSideCheck("abc def", 5),
+      "Exceeds allowed length limitation: 5");
+
+  // Null input cases.
+  EXPECT_EQ(varcharTypeWriteSideCheck(std::nullopt, 5), std::nullopt);
+
+  // Edge cases - length limit must be positive
+  VELOX_ASSERT_USER_THROW(
+      varcharTypeWriteSideCheck("abc", 0),
+      "The length limit must be greater than 0.");
+  VELOX_ASSERT_USER_THROW(
+      varcharTypeWriteSideCheck("abc", -1),
+      "The length limit must be greater than 0.");
+
+  // Edge cases - input string is all spaces.
+  EXPECT_EQ(varcharTypeWriteSideCheck("   ", 2), "  ");
+  EXPECT_EQ(varcharTypeWriteSideCheck("   ", 3), "   ");
+  EXPECT_EQ(varcharTypeWriteSideCheck("   ", 1), " ");
+}
+
+} // namespace
+} // namespace facebook::velox::functions::sparksql::test