diff --git a/velox/experimental/stateful/udf/CMakeLists.txt b/velox/experimental/stateful/udf/CMakeLists.txt index e9eba89c73b7..0f0169ce136c 100644 --- a/velox/experimental/stateful/udf/CMakeLists.txt +++ b/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() \ No newline at end of file diff --git a/velox/experimental/stateful/udf/Register.cpp b/velox/experimental/stateful/udf/Register.cpp index 27b473099534..fe0555a8564e 100644 --- a/velox/experimental/stateful/udf/Register.cpp +++ b/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( {prefix + "extract"}); + registerFunction( + {prefix + "split_index"}); } } // namespace facebook::velox::stateful::udf diff --git a/velox/experimental/stateful/udf/SplitIndex.h b/velox/experimental/stateful/udf/SplitIndex.h new file mode 100644 index 000000000000..ac7409f06cd8 --- /dev/null +++ b/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 +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& result, + const arg_type& input, + const arg_type& delimiter, + const int64_t& index) { + if (index <= 0) { + return false; + } + return functions::stringImpl::splitPart( + result, input, delimiter, index); + } + + FOLLY_ALWAYS_INLINE bool callAscii( + out_type& result, + const arg_type& input, + const arg_type& delimiter, + const int64_t& index) { + if (index <= 0) { + return false; + } + return functions::stringImpl::splitPart( + result, input, delimiter, index); + } +}; + +} // namespace facebook::velox::stateful::udf \ No newline at end of file diff --git a/velox/experimental/stateful/udf/tests/CMakeLists.txt b/velox/experimental/stateful/udf/tests/CMakeLists.txt new file mode 100644 index 000000000000..255edc09cfda --- /dev/null +++ b/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) diff --git a/velox/experimental/stateful/udf/tests/UDFTest.cpp b/velox/experimental/stateful/udf/tests/UDFTest.cpp new file mode 100644 index 000000000000..9173187dc218 --- /dev/null +++ b/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 +#include + +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& a, + const std::optional& d, + const std::optional& i) { + return evaluateOnce("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 res0 = splitIndex("a/b/c", "/", 0); + const std::optional res1 = splitIndex("a/b/c", "/", -1); + const std::optional 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(); +}