Skip to content

Commit

Permalink
custom comparisons array min max
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevinwilfong @facebook-github-bot
kevinwilfong authored and facebook-github-bot committed Oct 1, 2024
1 parent c7560e3 commit 131d175
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 0 deletions.
57 changes: 57 additions & 0 deletions velox/functions/prestosql/tests/ArrayMaxTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

#include <optional>
#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"

using namespace facebook::velox;
using namespace facebook::velox::test;
Expand Down Expand Up @@ -143,6 +144,62 @@ TEST_F(ArrayMaxTest, longVarcharNoNulls) {
testArrayMax(input, expected);
}

TEST_F(ArrayMaxTest, timestampWithTimezone) {
// Test primitive type.
auto arrayVector = makeArrayVector(
{0, 6, 13, 18, 23},
makeFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), pack(105, 3)},
TIMESTAMP_WITH_TIME_ZONE()));
VectorPtr expected = makeNullableFlatVector<int64_t>(
{pack(4, 5), pack(4, 0), pack(-1, 4), pack(105, 3), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE());
testArrayMax(arrayVector, expected);

// Test primitive type with nulls.
arrayVector = makeArrayVector(
{0, 6, 13, 18, 23, 28, 28},
makeNullableFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), std::nullopt, std::nullopt, pack(-1, 4),
pack(-2, 5), pack(-3, 1), pack(-4, 0), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));
expected = makeNullableFlatVector<int64_t>(
{pack(4, 5),
pack(4, 0),
pack(-1, 4),
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE());
testArrayMax(arrayVector, expected);

// Test wrapped in complex type.
arrayVector = makeArrayVector(
{0, 6, 13, 18, 23},
makeRowVector({makeFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), pack(105, 3)},
TIMESTAMP_WITH_TIME_ZONE())}));
expected = makeRowVector(
{makeNullableFlatVector<int64_t>(
{pack(4, 5), pack(4, 0), pack(-1, 4), pack(105, 3), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE())},
[](auto row) { return row == 4; });
testArrayMax(arrayVector, expected);
}

// Test documented example.
TEST_F(ArrayMaxTest, docs) {
{
Expand Down
58 changes: 58 additions & 0 deletions velox/functions/prestosql/tests/ArrayMinTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#include <optional>
#include "velox/common/base/tests/GTestUtils.h"
#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"

using namespace facebook::velox;
using namespace facebook::velox::test;
Expand Down Expand Up @@ -284,3 +285,60 @@ TEST_F(ArrayMinTest, complexTypeElements) {
});
assertEqualVectors(expected, result);
}

TEST_F(ArrayMinTest, timestampWithTimezone) {
// Test primitive type.
auto arrayVector = makeArrayVector(
{0, 6, 13, 18, 23},
makeFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), pack(105, 3)},
TIMESTAMP_WITH_TIME_ZONE()));
VectorPtr expected = makeNullableFlatVector<int64_t>(
{pack(-1, 0), pack(-2, 6), pack(-5, 3), pack(101, 4), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE());
testExpr<int64_t>(expected, "array_min(C0)", {arrayVector});

// Test primitive type with nulls.
arrayVector = makeArrayVector(
{0, 6, 13, 18, 23, 28, 28},
makeNullableFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), std::nullopt, std::nullopt, pack(-1, 4),
pack(-2, 5), pack(-3, 1), pack(-4, 0), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));
expected = makeNullableFlatVector<int64_t>(
{pack(-1, 0),
pack(-2, 6),
pack(-5, 3),
std::nullopt,
std::nullopt,
std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE());
testExpr<int64_t>(expected, "array_min(C0)", {arrayVector});

// Test wrapped in complex type.
arrayVector = makeArrayVector(
{0, 6, 13, 18, 23},
makeRowVector({makeFlatVector<int64_t>(
{pack(-1, 0), pack(0, 1), pack(1, 2), pack(2, 3), pack(3, 4),
pack(4, 5), pack(4, 0), pack(3, 1), pack(2, 2), pack(1, 3),
pack(0, 4), pack(-1, 5), pack(-2, 6), pack(-5, 3), pack(-4, 2),
pack(-3, 1), pack(-2, 0), pack(-1, 4), pack(101, 4), pack(102, 0),
pack(103, 1), pack(104, 2), pack(105, 3)},
TIMESTAMP_WITH_TIME_ZONE())}));
expected = makeRowVector(
{makeNullableFlatVector<int64_t>(
{pack(-1, 0), pack(-2, 6), pack(-5, 3), pack(101, 4), std::nullopt},
TIMESTAMP_WITH_TIME_ZONE())},
[](auto row) { return row == 4; });
auto result = evaluate("array_min(c0)", makeRowVector({arrayVector}));
assertEqualVectors(expected, result);
}

0 comments on commit 131d175

Please sign in to comment.