Skip to content

Commit

Permalink
Add tests for array_sort with TimestampWithTImezone
Browse files Browse the repository at this point in the history 
Summary:
With #11136 array_sort just works with
TimestampWithTimezone.

Adding unit tests to verify this and ensure it doesn't break in the future.

Differential Revision: D63711984
  • Loading branch information
@facebook-github-bot
Kevin Wilfong authored and facebook-github-bot committed Oct 1, 2024
1 parent 131d175 commit 2bcdf6d
Showing 1 changed file with 111 additions and 0 deletions.
111 changes: 111 additions & 0 deletions velox/functions/prestosql/tests/ArraySortTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,10 @@
* limitations under the License.
*/
#include "velox/common/base/tests/GTestUtils.h"
#include "velox/functions/Macros.h"
#include "velox/functions/Registerer.h"
#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"

#include <fmt/format.h>
#include <cstdint>
Expand Down Expand Up @@ -722,6 +725,114 @@ TEST_F(ArraySortTest, failOnRowNullCompare) {
}
}

TEST_F(ArraySortTest, timestampWithTimezone) {
auto data = makeRowVector({makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1),
pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0),
pack(0, 1), pack(2, 2), std::nullopt, pack(1, 0), pack(0, 1),
pack(2, 2), std::nullopt, std::nullopt, pack(1, 2), pack(0, 1),
pack(2, 0), std::nullopt, pack(1, 1), pack(0, 2), std::nullopt,
pack(2, 0), pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2),
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()))});

auto expectedAsc = makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(0, 2), pack(1, 1), pack(2, 0), pack(0, 0), pack(1, 1),
pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(0, 1),
pack(1, 0), pack(2, 2), pack(0, 1), pack(1, 0), pack(2, 2),
std::nullopt, pack(0, 1), pack(1, 2), pack(2, 0), std::nullopt,
std::nullopt, pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt,
std::nullopt, pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt,
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));

auto expectedDesc = makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(2, 0), pack(1, 1), pack(0, 2), pack(2, 2), pack(1, 1),
pack(0, 0), pack(0, 0), pack(0, 1), pack(0, 2), pack(2, 2),
pack(1, 0), pack(0, 1), pack(2, 2), pack(1, 0), pack(0, 1),
std::nullopt, pack(2, 0), pack(1, 2), pack(0, 1), std::nullopt,
std::nullopt, pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt,
std::nullopt, pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt,
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));

auto resultAsc = evaluate("array_sort(c0)", data);
assertEqualVectors(expectedAsc, resultAsc);

auto resultDesc = evaluate("array_sort_desc(c0)", data);
assertEqualVectors(expectedDesc, resultDesc);
}

template <typename T>
struct TimeZoneFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);

FOLLY_ALWAYS_INLINE void call(
int64_t& result,
const arg_type<TimestampWithTimezone>& ts) {
result = unpackZoneKeyId(*ts);
}
};

TEST_F(ArraySortTest, timestampWithTimezoneWithLambda) {
registerFunction<TimeZoneFunction, int64_t, TimestampWithTimezone>(
{"timezone"});

auto data = makeRowVector({makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1),
pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0),
pack(0, 1), pack(2, 2), std::nullopt, pack(1, 0), pack(0, 1),
pack(2, 2), std::nullopt, std::nullopt, pack(1, 2), pack(0, 1),
pack(2, 0), std::nullopt, pack(1, 1), pack(0, 2), std::nullopt,
pack(2, 0), pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2),
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()))});

auto expectedAsc = makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(2, 0), pack(1, 1), pack(0, 2), pack(0, 0), pack(1, 1),
pack(2, 2), pack(0, 0), pack(0, 1), pack(0, 2), pack(1, 0),
pack(0, 1), pack(2, 2), pack(1, 0), pack(0, 1), pack(2, 2),
std::nullopt, pack(2, 0), pack(0, 1), pack(1, 2), std::nullopt,
std::nullopt, pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt,
std::nullopt, pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt,
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));

auto expectedDesc = makeArrayVector(
{0, 3, 6, 9, 12, 16, 21, 26, 31},
makeNullableFlatVector<int64_t>(
{pack(0, 2), pack(1, 1), pack(2, 0), pack(2, 2), pack(1, 1),
pack(0, 0), pack(0, 2), pack(0, 1), pack(0, 0), pack(2, 2),
pack(0, 1), pack(1, 0), pack(2, 2), pack(0, 1), pack(1, 0),
std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0), std::nullopt,
std::nullopt, pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt,
std::nullopt, pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt,
std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt,
std::nullopt},
TIMESTAMP_WITH_TIME_ZONE()));

auto resultAsc = evaluate("array_sort(c0, x -> timezone(x))", data);
assertEqualVectors(expectedAsc, resultAsc);

auto resultDesc = evaluate("array_sort_desc(c0, x -> timezone(x))", data);
assertEqualVectors(expectedDesc, resultDesc);
}

TEST_F(ArraySortTest, floatingPointExtremes) {
testFloatingPoint<float>();
testFloatingPoint<double>();
Expand Down

0 comments on commit 2bcdf6d

Please sign in to comment.