Skip to content

Commit

Permalink
Add trail function (#11265)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #11265

Add trail function as per specs here: https://prestodb.io/docs/current/functions/string.html#trail-string-N-varchar

We return the last N characters of the string up to at most the length of the input string.

Reviewed By: amitkdutta

Differential Revision: D64379575

fbshipit-source-id: dc0e766e54550b7b81bb72fad7ea25c3e3e77e36
  • Loading branch information
yuandagits authored and facebook-github-bot committed Oct 16, 2024
1 parent 1dd1e3e commit ad510db
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 0 deletions.
4 changes: 4 additions & 0 deletions velox/docs/functions/presto/string.rst
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,10 @@ String Functions

SELECT strrpos('aaa', 'aa', 2); -- 1

.. function:: trail(string, N) -> varchar

Returns the last ``N`` characters of the input ``string`` up to at most the length of ``string``.

.. function:: substr(string, start) -> varchar

Returns the rest of ``string`` from the starting position ``start``.
Expand Down
54 changes: 54 additions & 0 deletions velox/functions/prestosql/StringFunctions.h
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,60 @@ struct CodePointFunction {
}
};

/// trail(string, N) -> varchar
///
/// Returns the last N characters of the input string.
template <typename T>
struct TrailFunction {
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;

template <typename I>
FOLLY_ALWAYS_INLINE void callNullFree(
out_type<Varchar>& result,
const null_free_arg_type<Varchar>& input,
I N) {
doCall<false>(result, input, N);
}

template <typename I>
FOLLY_ALWAYS_INLINE void
callAscii(out_type<Varchar>& result, const arg_type<Varchar>& input, I N) {
doCall<true>(result, input, N);
}

private:
template <bool isAscii, typename I>
FOLLY_ALWAYS_INLINE void
doCall(out_type<Varchar>& result, const arg_type<Varchar>& input, I N) {
if (N <= 0) {
result.setEmpty();
return;
}

I numCharacters = stringImpl::length<isAscii>(input);

// Get the start position of the last N characters
// If N is greater than the number of characters, start at 1/
I start = N > numCharacters ? 1 : numCharacters - N + 1;

// Adjust length
I adjustedLength = std::min(N, numCharacters);

auto byteRange = stringCore::getByteRange<isAscii>(
input.data(), input.size(), start, adjustedLength);

// Generating output string
result.setNoCopy(StringView(
input.data() + byteRange.first, byteRange.second - byteRange.first));
}
};

/// substr(string, start) -> varchar
///
/// Returns the rest of string from the starting position start.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,9 @@ void registerSimpleFunctions(const std::string& prefix) {
registerFunction<EndsWithFunction, bool, Varchar, Varchar>(
{prefix + "ends_with"});

registerFunction<TrailFunction, Varchar, Varchar, int64_t>(
{prefix + "trail"});

registerFunction<SubstrFunction, Varchar, Varchar, int64_t>(
{prefix + "substr"});
registerFunction<SubstrFunction, Varchar, Varchar, int64_t, int64_t>(
Expand Down
16 changes: 16 additions & 0 deletions velox/functions/prestosql/tests/StringFunctionsTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2164,3 +2164,19 @@ TEST_F(StringFunctionsTest, normalize) {
normalizeWithForm("sch\u00f6n", "NFKE"),
"Normalization form must be one of [NFD, NFC, NFKD, NFKC]");
}

TEST_F(StringFunctionsTest, trail) {
auto trail = [&](std::optional<std::string> string,
std::optional<int64_t> N) {
return evaluateOnce<std::string>("trail(c0, c1)", string, N);
};

// Basic Test
EXPECT_EQ("bar", trail("foobar", 3));
EXPECT_EQ("foobar", trail("foobar", 7));
EXPECT_EQ("", trail("foobar", 0));
EXPECT_EQ("", trail("foobar", -1));

// Test empty
EXPECT_EQ("", trail("", 3));
}

0 comments on commit ad510db

Please sign in to comment.