Skip to content

Commit

Permalink
Convert TIMESTAMP_WITH_TIME_ZONE to primitive type
Browse files Browse the repository at this point in the history
  • Loading branch information
wypb committed Feb 22, 2024
1 parent 51007ea commit 4e02577
Show file tree
Hide file tree
Showing 16 changed files with 280 additions and 366 deletions.
8 changes: 2 additions & 6 deletions velox/core/tests/ConstantTypedExprTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,8 @@ TEST(ConstantTypedExprTest, null) {
EXPECT_FALSE(*makeNull(HYPERLOGLOG()) == *makeNull(VARBINARY()));
EXPECT_FALSE(*makeNull(VARBINARY()) == *makeNull(HYPERLOGLOG()));

EXPECT_FALSE(
*makeNull(TIMESTAMP_WITH_TIME_ZONE()) ==
*makeNull(ROW({BIGINT(), SMALLINT()})));
EXPECT_FALSE(
*makeNull(ROW({BIGINT(), SMALLINT()})) ==
*makeNull(TIMESTAMP_WITH_TIME_ZONE()));
EXPECT_FALSE(*makeNull(TIMESTAMP_WITH_TIME_ZONE()) == *makeNull(BIGINT()));
EXPECT_FALSE(*makeNull(BIGINT()) == *makeNull(TIMESTAMP_WITH_TIME_ZONE()));

EXPECT_TRUE(*makeNull(DOUBLE()) == *makeNull(DOUBLE()));
EXPECT_TRUE(*makeNull(ARRAY(DOUBLE())) == *makeNull(ARRAY(DOUBLE())));
Expand Down
2 changes: 0 additions & 2 deletions velox/functions/lib/RegistrationHelpers.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,6 @@ void registerNonSimdizableScalar(const std::vector<std::string>& aliases) {
registerFunction<T, TReturn, bool, bool>(aliases);
registerFunction<T, TReturn, Timestamp, Timestamp>(aliases);
registerFunction<T, TReturn, Date, Date>(aliases);
registerFunction<T, TReturn, TimestampWithTimezone, TimestampWithTimezone>(
aliases);
}

template <template <class> class T>
Expand Down
92 changes: 55 additions & 37 deletions velox/functions/prestosql/Comparisons.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,17 @@ struct TimestampWithTimezoneComparisonSupport {
FOLLY_ALWAYS_INLINE
int64_t toGMTMillis(
const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
const int64_t milliseconds = *timestampWithTimezone.template at<0>();
const int16_t timezone = *timestampWithTimezone.template at<1>();
Timestamp inputTimeStamp = Timestamp::fromMillis(milliseconds);
auto inputTimeStamp = unpackTimestampUtc(timestampWithTimezone);
const auto timezone = unpackZoneKeyId(timestampWithTimezone);
inputTimeStamp.toGMT(timezone);

return inputTimeStamp.toMillis();
}
};

} // namespace

#define VELOX_GEN_BINARY_EXPR(Name, Expr, tsExpr, TResult) \
#define VELOX_GEN_BINARY_EXPR(Name, Expr, TResult) \
template <typename T> \
struct Name : public TimestampWithTimezoneComparisonSupport<T> { \
VELOX_DEFINE_FUNCTION_TYPES(T); \
Expand All @@ -49,37 +49,45 @@ struct TimestampWithTimezoneComparisonSupport {
call(TResult& result, const TInput& lhs, const TInput& rhs) { \
result = (Expr); \
} \
\
FOLLY_ALWAYS_INLINE void call( \
bool& result, \
const arg_type<TimestampWithTimezone>& lhs, \
const arg_type<TimestampWithTimezone>& rhs) { \
result = (tsExpr); \
} \
};

VELOX_GEN_BINARY_EXPR(
#define VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE(Name, tsExpr, TResult) \
template <typename T> \
struct Name##TimestampWithTimezone \
: public TimestampWithTimezoneComparisonSupport<T> { \
VELOX_DEFINE_FUNCTION_TYPES(T); \
FOLLY_ALWAYS_INLINE void call( \
bool& result, \
const arg_type<TimestampWithTimezone>& lhs, \
const arg_type<TimestampWithTimezone>& rhs) { \
result = (tsExpr); \
} \
};

VELOX_GEN_BINARY_EXPR(LtFunction, lhs < rhs, bool);
VELOX_GEN_BINARY_EXPR(GtFunction, lhs > rhs, bool);
VELOX_GEN_BINARY_EXPR(LteFunction, lhs <= rhs, bool);
VELOX_GEN_BINARY_EXPR(GteFunction, lhs >= rhs, bool);

VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE(
LtFunction,
lhs < rhs,
this->toGMTMillis(lhs) < this->toGMTMillis(rhs),
bool);
VELOX_GEN_BINARY_EXPR(
VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE(
GtFunction,
lhs > rhs,
this->toGMTMillis(lhs) > this->toGMTMillis(rhs),
bool);
VELOX_GEN_BINARY_EXPR(
VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE(
LteFunction,
lhs <= rhs,
this->toGMTMillis(lhs) <= this->toGMTMillis(rhs),
bool);
VELOX_GEN_BINARY_EXPR(
VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE(
GteFunction,
lhs >= rhs,
this->toGMTMillis(lhs) >= this->toGMTMillis(rhs),
bool);

#undef VELOX_GEN_BINARY_EXPR
#undef VELOX_GEN_BINARY_EXPR_TIMESTAMP_WITH_TIME_ZONE

template <typename T>
struct DistinctFromFunction {
Expand All @@ -104,7 +112,7 @@ struct DistinctFromFunction {
};

template <typename T>
struct EqFunction : public TimestampWithTimezoneComparisonSupport<T> {
struct EqFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);

// Used for primitive inputs.
Expand All @@ -113,14 +121,6 @@ struct EqFunction : public TimestampWithTimezoneComparisonSupport<T> {
out = (lhs == rhs);
}

// For TimestampWithTimezone.
void call(
bool& result,
const arg_type<TimestampWithTimezone>& lhs,
const arg_type<TimestampWithTimezone>& rhs) {
result = this->toGMTMillis(lhs) == this->toGMTMillis(rhs);
}

// For arbitrary nested complex types. Can return null.
bool call(
bool& out,
Expand All @@ -139,21 +139,26 @@ struct EqFunction : public TimestampWithTimezoneComparisonSupport<T> {
};

template <typename T>
struct NeqFunction : public TimestampWithTimezoneComparisonSupport<T> {
struct EqFunctionTimestampWithTimezone
: public TimestampWithTimezoneComparisonSupport<T> {
VELOX_DEFINE_FUNCTION_TYPES(T);

// Used for primitive inputs.
template <typename TInput>
void call(bool& out, const TInput& lhs, const TInput& rhs) {
out = (lhs != rhs);
}

// For TimestampWithTimezone.
void call(
bool& result,
const arg_type<TimestampWithTimezone>& lhs,
const arg_type<TimestampWithTimezone>& rhs) {
result = this->toGMTMillis(lhs) != this->toGMTMillis(rhs);
result = this->toGMTMillis(lhs) == this->toGMTMillis(rhs);
}
};

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

// Used for primitive inputs.
template <typename TInput>
void call(bool& out, const TInput& lhs, const TInput& rhs) {
out = (lhs != rhs);
}

// For arbitrary nested complex types. Can return null.
Expand All @@ -170,6 +175,19 @@ struct NeqFunction : public TimestampWithTimezoneComparisonSupport<T> {
}
};

template <typename T>
struct NeqFunctionTimestampWithTimezone
: public TimestampWithTimezoneComparisonSupport<T> {
VELOX_DEFINE_FUNCTION_TYPES(T);

void call(
bool& result,
const arg_type<TimestampWithTimezone>& lhs,
const arg_type<TimestampWithTimezone>& rhs) {
result = this->toGMTMillis(lhs) != this->toGMTMillis(rhs);
}
};

template <typename T>
struct BetweenFunction {
template <typename TInput>
Expand Down
39 changes: 16 additions & 23 deletions velox/functions/prestosql/DateTimeFunctions.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ struct ToUnixtimeFunction {
FOLLY_ALWAYS_INLINE bool call(
double& result,
const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
const auto milliseconds = *timestampWithTimezone.template at<0>();
const auto milliseconds = unpackMillisUtc(timestampWithTimezone);
result = (double)milliseconds / kMillisecondsInSecond;
return true;
}
Expand Down Expand Up @@ -75,11 +75,9 @@ struct TimestampWithTimezoneSupport {
Timestamp toTimestamp(
const arg_type<TimestampWithTimezone>& timestampWithTimezone,
bool asGMT = false) {
const auto milliseconds = *timestampWithTimezone.template at<0>();
auto tz = *timestampWithTimezone.template at<1>();
Timestamp timestamp = Timestamp::fromMillis(milliseconds);
auto timestamp = unpackTimestampUtc(timestampWithTimezone);
if (!asGMT) {
timestamp.toTimezone(*timestampWithTimezone.template at<1>());
timestamp.toTimezone(unpackZoneKeyId(timestampWithTimezone));
}

return timestamp;
Expand All @@ -90,10 +88,9 @@ struct TimestampWithTimezoneSupport {
int64_t getGMTOffsetSec(
const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
Timestamp inputTimeStamp = this->toTimestamp(timestampWithTimezone);

// Create a copy of inputTimeStamp and convert it to GMT
auto gmtTimeStamp = inputTimeStamp;
gmtTimeStamp.toGMT(*timestampWithTimezone.template at<1>());
gmtTimeStamp.toGMT(unpackZoneKeyId(timestampWithTimezone));

// Get offset in seconds with GMT and convert to hour
return (inputTimeStamp.getSeconds() - gmtTimeStamp.getSeconds());
Expand Down Expand Up @@ -908,23 +905,20 @@ struct DateTruncFunction : public TimestampWithTimezoneSupport<T> {
}

if (unit == DateTimeUnit::kSecond) {
auto utcTimestamp =
Timestamp::fromMillis(*timestampWithTimezone.template at<0>());
result.template get_writer_at<0>() = utcTimestamp.getSeconds() * 1000;
result.template get_writer_at<1>() =
*timestampWithTimezone.template at<1>();
auto utcTimestamp = unpackTimestampUtc(timestampWithTimezone);
result = pack(
utcTimestamp.getSeconds() * 1000,
unpackZoneKeyId(timestampWithTimezone));
return;
}

auto timestamp = this->toTimestamp(timestampWithTimezone);
auto dateTime = getDateTime(timestamp, nullptr);
adjustDateTime(dateTime, unit);
timestamp = Timestamp::fromMillis(timegm(&dateTime) * 1000);
timestamp.toGMT(*timestampWithTimezone.template at<1>());
timestamp.toGMT(unpackZoneKeyId(timestampWithTimezone));

result.template get_writer_at<0>() = timestamp.toMillis();
result.template get_writer_at<1>() =
*timestampWithTimezone.template at<1>();
result = pack(timestamp.toMillis(), unpackZoneKeyId(timestampWithTimezone));
}
};

Expand Down Expand Up @@ -1009,9 +1003,9 @@ struct DateAddFunction : public TimestampWithTimezoneSupport<T> {

auto finalTimeStamp = addToTimestamp(
this->toTimestamp(timestampWithTimezone), unit, (int32_t)value);
finalTimeStamp.toGMT(*timestampWithTimezone.template at<1>());
result = std::make_tuple(
finalTimeStamp.toMillis(), *timestampWithTimezone.template at<1>());
auto tzID = unpackZoneKeyId(timestampWithTimezone);
finalTimeStamp.toGMT(tzID);
result = pack(finalTimeStamp.toMillis(), tzID);

return true;
}
Expand Down Expand Up @@ -1278,9 +1272,8 @@ struct FormatDateTimeFunction {
const arg_type<Varchar>& formatString) {
ensureFormatter(formatString);

const auto milliseconds = *timestampWithTimezone.template at<0>();
Timestamp timestamp = Timestamp::fromMillis(milliseconds);
int16_t timeZoneId = *timestampWithTimezone.template at<1>();
const auto timestamp = unpackTimestampUtc(timestampWithTimezone);
const auto timeZoneId = unpackZoneKeyId(timestampWithTimezone);
auto* timezonePtr = date::locate_zone(util::getTimeZoneName(timeZoneId));

auto maxResultSize = jodaDateTime_->maxResultSize(timezonePtr);
Expand Down Expand Up @@ -1345,7 +1338,7 @@ struct ParseDateTimeFunction {
? dateTimeResult.timezoneId
: sessionTzID_.value_or(0);
dateTimeResult.timestamp.toGMT(timezoneId);
result = std::make_tuple(dateTimeResult.timestamp.toMillis(), timezoneId);
result = pack(dateTimeResult.timestamp.toMillis(), timezoneId);
return true;
}
};
Expand Down
30 changes: 12 additions & 18 deletions velox/functions/prestosql/FromUnixTime.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
*/
#include "velox/expression/Expr.h"
#include "velox/expression/VectorFunction.h"
#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
#include "velox/type/tz/TimeZoneMap.h"

namespace facebook::velox::functions {
Expand Down Expand Up @@ -43,43 +44,36 @@ class FromUnixtimeFunction : public exec::VectorFunction {
const auto size = rows.end();
auto* pool = context.pool();

auto timestamps = BaseVector::create(BIGINT(), size, pool);
auto* rawTimestamps =
timestamps->asFlatVector<int64_t>()->mutableRawValues();
auto timestamps = AlignedBuffer::allocate<int64_t>(size, pool);
auto* rawTimestamps = timestamps->asMutable<int64_t>();

VectorPtr timezones;
if (timezoneNames->isConstantMapping()) {
auto timezoneName = timezoneNames->valueAt<StringView>(rows.begin());

int16_t timezoneId = util::getTimeZoneID(
std::string_view(timezoneName.data(), timezoneName.size()));
timezones = std::make_shared<ConstantVector<int16_t>>(
pool, size, false /*isNull*/, SMALLINT(), std::move(timezoneId));

rows.applyToSelected([&](auto row) {
rawTimestamps[row] = toMillis(unixtimes->valueAt<double>(row));
rawTimestamps[row] =
pack(toMillis(unixtimes->valueAt<double>(row)), timezoneId);
});
} else {
timezones = BaseVector::create(SMALLINT(), size, pool);
auto* rawTimezones =
timezones->asFlatVector<int16_t>()->mutableRawValues();

rows.applyToSelected([&](auto row) {
rawTimestamps[row] = toMillis(unixtimes->valueAt<double>(row));

auto timezoneName = timezoneNames->valueAt<StringView>(row);
rawTimezones[row] = util::getTimeZoneID(
auto timezone = util::getTimeZoneID(
std::string_view(timezoneName.data(), timezoneName.size()));
rawTimestamps[row] =
pack(toMillis(unixtimes->valueAt<double>(row)), timezone);
});
}

auto localResult = std::make_shared<RowVector>(
auto localResult = std::make_shared<FlatVector<int64_t>>(
pool,
outputType,
BufferPtr(nullptr),
rows.end(),
std::vector<VectorPtr>{timestamps, timezones},
0 /*nullCount*/);
size,
std::move(timestamps),
std::vector<BufferPtr>());

context.moveOrCopyResult(localResult, rows, result);
}
Expand Down
19 changes: 5 additions & 14 deletions velox/functions/prestosql/aggregates/PrestoHasher.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,11 @@ FOLLY_ALWAYS_INLINE void PrestoHasher::hash<TypeKind::BIGINT>(
// returns the corresponding value directly.
return vector_->valueAt<int64_t>(row);
});
} else if (isTimestampWithTimeZoneType(vector_->base()->type())) {
// Hash only timestamp value.
applyHashFunction(rows, *vector_.get(), hashes, [&](auto row) {
return hashInteger(unpackMillisUtc(vector_->valueAt<int64_t>(row)));
});
} else {
applyHashFunction(rows, *vector_.get(), hashes, [&](auto row) {
return hashInteger(vector_->valueAt<int64_t>(row));
Expand Down Expand Up @@ -319,20 +324,6 @@ void PrestoHasher::hash<TypeKind::ROW>(

auto rawHashes = hashes->asMutable<int64_t>();

if (isTimestampWithTimeZoneType(vector_->base()->type())) {
// Hash only timestamp value.
children_[0]->hash(baseRow->childAt(0), elementRows, childHashes);
auto rawChildHashes = childHashes->as<int64_t>();
rows.applyToSelected([&](auto row) {
if (!baseRow->isNullAt(indices[row])) {
rawHashes[row] = rawChildHashes[indices[row]];
} else {
rawHashes[row] = 0;
}
});
return;
}

BufferPtr combinedChildHashes =
AlignedBuffer::allocate<int64_t>(elementRows.end(), baseRow->pool());
auto* rawCombinedChildHashes = combinedChildHashes->asMutable<int64_t>();
Expand Down
Loading

0 comments on commit 4e02577

Please sign in to comment.