Skip to content

Commit

Permalink
Add tests for array_sort with TimestampWithTImezone (#11142)
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.

Reviewed By: bikramSingh91

Differential Revision: D63711984
  • Loading branch information
@facebook-github-bot
Kevin Wilfong authored and facebook-github-bot committed Oct 7, 2024
1 parent 0403db4 commit 8f06c64
Showing 1 changed file with 146 additions and 0 deletions.
146 changes: 146 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,149 @@ TEST_F(ArraySortTest, failOnRowNullCompare) {
}
}

TEST_F(ArraySortTest, timestampWithTimezone) {
auto testArraySort =
[this](
const std::vector<std::optional<int64_t>>& inputArray,
const std::vector<std::optional<int64_t>>& expectedAscArray,
const std::vector<std::optional<int64_t>>& expectedDescArray) {
const auto input = makeRowVector({makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
inputArray, TIMESTAMP_WITH_TIME_ZONE()))});
const auto expectedAsc = makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
expectedAscArray, TIMESTAMP_WITH_TIME_ZONE()));
const auto expectedDesc = makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
expectedDescArray, TIMESTAMP_WITH_TIME_ZONE()));

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

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

testArraySort(
{pack(2, 0), pack(1, 1), pack(0, 2)},
{pack(0, 2), pack(1, 1), pack(2, 0)},
{pack(2, 0), pack(1, 1), pack(0, 2)});
testArraySort(
{pack(0, 0), pack(1, 1), pack(2, 2)},
{pack(0, 0), pack(1, 1), pack(2, 2)},
{pack(2, 2), pack(1, 1), pack(0, 0)});
testArraySort(
{pack(0, 0), pack(0, 1), pack(0, 2)},
{pack(0, 0), pack(0, 1), pack(0, 2)},
{pack(0, 0), pack(0, 1), pack(0, 2)});
testArraySort(
{pack(1, 0), pack(0, 1), pack(2, 2)},
{pack(0, 1), pack(1, 0), pack(2, 2)},
{pack(2, 2), pack(1, 0), pack(0, 1)});
testArraySort(
{std::nullopt, pack(1, 0), pack(0, 1), pack(2, 2)},
{pack(0, 1), pack(1, 0), pack(2, 2), std::nullopt},
{pack(2, 2), pack(1, 0), pack(0, 1), std::nullopt});
testArraySort(
{std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0)},
{pack(0, 1), pack(1, 2), pack(2, 0), std::nullopt, std::nullopt},
{pack(2, 0), pack(1, 2), pack(0, 1), std::nullopt, std::nullopt});
testArraySort(
{std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, pack(2, 0)},
{pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, std::nullopt},
{pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, std::nullopt});
testArraySort(
{pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), std::nullopt},
{pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, std::nullopt},
{pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, std::nullopt});
testArraySort(
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt});
}

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 testArraySort =
[this](
const std::vector<std::optional<int64_t>>& inputArray,
const std::vector<std::optional<int64_t>>& expectedAscArray,
const std::vector<std::optional<int64_t>>& expectedDescArray) {
const auto input = makeRowVector({makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
inputArray, TIMESTAMP_WITH_TIME_ZONE()))});
const auto expectedAsc = makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
expectedAscArray, TIMESTAMP_WITH_TIME_ZONE()));
const auto expectedDesc = makeArrayVector(
{0},
makeNullableFlatVector<int64_t>(
expectedDescArray, TIMESTAMP_WITH_TIME_ZONE()));

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

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

testArraySort(
{pack(2, 0), pack(1, 1), pack(0, 2)},
{pack(2, 0), pack(1, 1), pack(0, 2)},
{pack(0, 2), pack(1, 1), pack(2, 0)});
testArraySort(
{pack(0, 0), pack(1, 1), pack(2, 2)},
{pack(0, 0), pack(1, 1), pack(2, 2)},
{pack(2, 2), pack(1, 1), pack(0, 0)});
testArraySort(
{pack(0, 0), pack(0, 1), pack(0, 2)},
{pack(0, 0), pack(0, 1), pack(0, 2)},
{pack(0, 2), pack(0, 1), pack(0, 0)});
testArraySort(
{pack(1, 0), pack(0, 1), pack(2, 2)},
{pack(1, 0), pack(0, 1), pack(2, 2)},
{pack(2, 2), pack(0, 1), pack(1, 0)});
testArraySort(
{std::nullopt, pack(1, 0), pack(0, 1), pack(2, 2)},
{pack(1, 0), pack(0, 1), pack(2, 2), std::nullopt},
{pack(2, 2), pack(0, 1), pack(1, 0), std::nullopt});
testArraySort(
{std::nullopt, std::nullopt, pack(1, 2), pack(0, 1), pack(2, 0)},
{pack(2, 0), pack(0, 1), pack(1, 2), std::nullopt, std::nullopt},
{pack(1, 2), pack(0, 1), pack(2, 0), std::nullopt, std::nullopt});
testArraySort(
{std::nullopt, pack(1, 1), pack(0, 2), std::nullopt, pack(2, 0)},
{pack(2, 0), pack(1, 1), pack(0, 2), std::nullopt, std::nullopt},
{pack(0, 2), pack(1, 1), pack(2, 0), std::nullopt, std::nullopt});
testArraySort(
{pack(1, 1), std::nullopt, pack(0, 0), pack(2, 2), std::nullopt},
{pack(0, 0), pack(1, 1), pack(2, 2), std::nullopt, std::nullopt},
{pack(2, 2), pack(1, 1), pack(0, 0), std::nullopt, std::nullopt});
testArraySort(
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt},
{std::nullopt, std::nullopt, std::nullopt, std::nullopt, std::nullopt});
}

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

0 comments on commit 8f06c64

Please sign in to comment.