[FLINK]Support function split_index

velox/experimental/stateful/udf/CMakeLists.txt

@@ -16,3 +16,7 @@ velox_add_library(
   velox_stateful_udf
   OBJECT
   Register.cpp)
+
+if(${VELOX_BUILD_TESTING})
+  add_subdirectory(tests)
+endif()

velox/experimental/stateful/udf/Register.cpp

@@ -16,6 +16,7 @@
 #include "velox/experimental/stateful/udf/Register.h"
 #include "velox/experimental/stateful/udf/BigoUDF.h"
 #include "velox/experimental/stateful/udf/ExtractDateTime.h"
+#include "velox/experimental/stateful/udf/SplitIndex.h"
 #include "velox/functions/Registerer.h"
 
 namespace facebook::velox::stateful::udf {
@@ -25,6 +26,8 @@ void registerFunctions(const std::string& prefix) {
       {prefix + "count_char"});
   registerFunction<ExtractFunction, int64_t, Varchar, Timestamp>(
       {prefix + "extract"});
+  registerFunction<SplitIndexFunction, Varchar, Varchar, Varchar, int64_t>(
+      {prefix + "split_index"});
 }
 
 } // namespace facebook::velox::stateful::udf

velox/experimental/stateful/udf/SplitIndex.h

@@ -0,0 +1,58 @@
+/*
+ * 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 "velox/functions/Macros.h"
+#include "velox/functions/lib/string/StringImpl.h"
+
+namespace facebook::velox::stateful::udf {
+
+template <typename T>
+struct SplitIndexFunction {
+  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 bool call(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& input,
+      const arg_type<Varchar>& delimiter,
+      const int64_t& index) {
+    if (index <= 0) {
+      return false;
+    }
+    return functions::stringImpl::splitPart<false>(
+        result, input, delimiter, index);
+  }
+
+  FOLLY_ALWAYS_INLINE bool callAscii(
+      out_type<Varchar>& result,
+      const arg_type<Varchar>& input,
+      const arg_type<Varchar>& delimiter,
+      const int64_t& index) {
+    if (index <= 0) {
+      return false;
+    }
+    return functions::stringImpl::splitPart<true>(
+        result, input, delimiter, index);
+  }
+};
+
+} // namespace facebook::velox::stateful::udf

velox/experimental/stateful/udf/tests/CMakeLists.txt

@@ -0,0 +1,25 @@
+# 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.
+add_executable(velox_stateful_udf_test UDFTest.cpp)
+
+add_test(velox_stateful_udf_test velox_stateful_udf_test)
+
+target_link_libraries(
+  velox_stateful_udf_test
+  velox_vector_test_lib
+  velox_exec_test_lib
+  velox_exec
+  velox_common_base
+  GTest::gtest
+  GTest::gtest_main)

velox/experimental/stateful/udf/tests/UDFTest.cpp

@@ -0,0 +1,57 @@
+/*
+ * 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/experimental/stateful/udf/Register.h"
+#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
+#include "velox/parse/TypeResolver.h"
+
+#include <folly/init/Init.h>
+#include <gtest/gtest.h>
+
+namespace facebook::velox::udf::test {
+
+class UDFTest : public functions::test::FunctionBaseTest {
+ protected:
+  static void SetUpTestCase() {
+    parse::registerTypeResolver();
+    stateful::udf::registerFunctions("");
+    memory::MemoryManager::testingSetInstance({});
+  }
+};
+
+TEST_F(UDFTest, splitIndex) {
+  const auto splitIndex = [&](const std::optional<std::string>& a,
+                              const std::optional<std::string>& d,
+                              const std::optional<int64_t>& i) {
+    return evaluateOnce<std::string>("split_index(c0, c1, c2)", a, d, i);
+  };
+  EXPECT_EQ(splitIndex("a/b/c", "/", 1), "a");
+  EXPECT_EQ(splitIndex("a/b/c", "/", 2), "b");
+  const std::optional<std::string> res0 = splitIndex("a/b/c", "/", 0);
+  const std::optional<std::string> res1 = splitIndex("a/b/c", "/", -1);
+  const std::optional<std::string> res2 = splitIndex("a/b/c", "/", 4);
+  EXPECT_EQ(res0.has_value(), false);
+  EXPECT_EQ(res1.has_value(), false);
+  EXPECT_EQ(res2.has_value(), false);
+}
+
+} // namespace facebook::velox::udf::test
+
+int main(int argc, char* argv[]) {
+  testing::InitGoogleTest(&argc, argv);
+  folly::Init init(&argc, &argv, false);
+  gflags::ParseCommandLineFlags(&argc, &argv, true); // Parse gflags
+  return RUN_ALL_TESTS();
+}