Skip to content

Commit

Permalink
Add tests for Presto comparison functions for custom types with custo…
Browse files Browse the repository at this point in the history 
…m comparison functions (#11019)

Summary:
Pull Request resolved: #11019

#11015 added support for custom types to provide
custom comparison functions, motivated by the need to support TimestampWithTimezone's
comparison semantics.

I validated the UDFs provided in Comparisons.h on top of this change.

It does not look like any further changes are needed in those files, but I added explicit tests for
TimestampWithTimezone in the various comparison UDFs.  I also added tests for the UDFs that
support Generic arguments, to ensure that when the Type underlying the Generic provides custom
comparison functions, that is respected.

I considered adding checks to the UDFs that support templated arguments to ensure the type does
not provide custom comparison functions, as these are not supported in any of them (e.g. between,
<, >, and the SIMD versions of comparison functions). However, someone would have to explicitly
register these functions for the custom type, and I don't think it's unreasonable to expect the user to
test these functions work as expected when doing so, I don't think the per-batch check is worth the
(admittedly small) overhead.

Differential Revision: D62893143
  • Loading branch information
@facebook-github-bot
Kevin Wilfong authored and facebook-github-bot committed Sep 17, 2024
1 parent a34cad5 commit c7b46a9
Showing 1 changed file with 130 additions and 0 deletions.
130 changes: 130 additions & 0 deletions velox/functions/prestosql/tests/ComparisonsTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@
#include "velox/functions/Udf.h"
#include "velox/functions/lib/RegistrationHelpers.h"
#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
#include "velox/type/tests/utils/CustomTypesForTesting.h"
#include "velox/type/tz/TimeZoneMap.h"

using namespace facebook::velox;

Expand Down Expand Up @@ -738,6 +740,134 @@ TEST_F(ComparisonsTest, nanComparison) {
testNaN("", input, false);
}

TEST_F(ComparisonsTest, TimestampWithTimezone) {
auto makeTimestampWithTimezone = [](int64_t millis, const std::string& tz) {
return pack(millis, tz::getTimeZoneID(tz));
};

auto lhs = makeFlatVector<int64_t>(
{makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1549770072000, "+01:00"),
makeTimestampWithTimezone(1549770072000, "+01:00"),
makeTimestampWithTimezone(1549770072000, "+03:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+03:00")},
TIMESTAMP_WITH_TIME_ZONE());

auto rhs = makeFlatVector<int64_t>(
{makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1639426440000, "-14:00"),
makeTimestampWithTimezone(1639426440000, "-14:00"),
makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1639426440000, "-14:00"),
makeTimestampWithTimezone(1639426440000, "-14:00"),
makeTimestampWithTimezone(1549770072000, "+03:00"),
makeTimestampWithTimezone(1549770072000, "-14:00"),
makeTimestampWithTimezone(1549770072000, "-14:00")},
TIMESTAMP_WITH_TIME_ZONE());

auto input = makeRowVector({lhs, rhs});

auto eval = [&](const std::string& expr) {
return evaluate<SimpleVector<bool>>(fmt::format("c0 {} c1", expr), input);
};

test::assertEqualVectors(
eval("="),
makeFlatVector<bool>(
{true, true, true, false, false, false, false, false, false}));
test::assertEqualVectors(
eval("<>"),
makeFlatVector<bool>(
{false, false, false, true, true, true, true, true, true}));
test::assertEqualVectors(
eval("<"),
makeFlatVector<bool>(
{false, false, false, true, true, true, false, false, false}));
test::assertEqualVectors(
eval(">"),
makeFlatVector<bool>(
{false, false, false, false, false, false, true, true, true}));
test::assertEqualVectors(
eval("<="),
makeFlatVector<bool>(
{true, true, true, true, true, true, false, false, false}));
test::assertEqualVectors(
eval(">="),
makeFlatVector<bool>(
{true, true, true, false, false, false, true, true, true}));
test::assertEqualVectors(
eval("is distinct from"),
makeFlatVector<bool>(
{false, false, false, true, true, true, true, true, true}));

auto betweenInput = makeRowVector({
makeFlatVector<int64_t>(
{makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1539426440000, "+03:00"),
makeTimestampWithTimezone(1739426440000, "+01:00")},
TIMESTAMP_WITH_TIME_ZONE()),
makeFlatVector<int64_t>(
{
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1539426440000, "+01:00"),
makeTimestampWithTimezone(1539426440000, "+03:00"),
makeTimestampWithTimezone(1539426440000, "+03:00"),
makeTimestampWithTimezone(1539426440000, "+03:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1539426440000, "+03:00"),
},
TIMESTAMP_WITH_TIME_ZONE()),
makeFlatVector<int64_t>(
{
makeTimestampWithTimezone(1739426440000, "+01:00"),
makeTimestampWithTimezone(1739426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+03:00"),
makeTimestampWithTimezone(1739426440000, "+03:00"),
makeTimestampWithTimezone(1739426440000, "+03:00"),
makeTimestampWithTimezone(1739426440000, "+01:00"),
makeTimestampWithTimezone(1639426440000, "+01:00"),
},
TIMESTAMP_WITH_TIME_ZONE()),
});

test::assertEqualVectors(
evaluate<SimpleVector<bool>>(
fmt::format("c0 between c1 and c2"), betweenInput),
makeFlatVector<bool>({true, true, true, true, true, true, false, false}));
}

TEST_F(ComparisonsTest, CustomComparisonWithGenerics) {
// Tests that functions that support signatures with generics handle custom
// comparison correctly.
auto input = makeRowVector({
makeFlatVector<int64_t>(
{0, 1}, test::BIGINT_TYPE_WITH_CUSTOM_COMPARISON()),
makeFlatVector<int64_t>(
{256, 258}, test::BIGINT_TYPE_WITH_CUSTOM_COMPARISON()),
});

test::assertEqualVectors(
evaluate<SimpleVector<bool>>("c0 = c1", input),
makeFlatVector<bool>({true, false}));
test::assertEqualVectors(
evaluate<SimpleVector<bool>>("c0 <> c1", input),
makeFlatVector<bool>({false, true}));
test::assertEqualVectors(
evaluate<SimpleVector<bool>>("c0 is distinct from c1", input),
makeFlatVector<bool>({false, true}));
}

namespace {
template <typename Tp, typename Op, const char* fnName>
struct ComparisonTypeOp {
Expand Down

0 comments on commit c7b46a9

Please sign in to comment.