Skip to content

Commit 78542bf

Browse files
rui-morui-mo
committed
Prepare time type extension for additional precisions and timezone support
1 parent 4d93698 commit 78542bf

File tree

10 files changed

+31
-13
lines changed

10 files changed

+31
-13
lines changed

‎velox/common/fuzzer/ConstrainedGenerators.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ class RandomInputGenerator : public AbstractInputGenerator {
6363
return variant(randDate(rng_));
6464
}
6565
if (type_->isTime()) {
66+
VELOX_DCHECK(type_->equivalent(*TIME()));
6667
return variant(randTime(rng_));
6768
}
6869
return variant(rand<T>(rng_));

‎velox/duckdb/conversion/DuckConversion.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ LogicalType fromVeloxType(const TypePtr& type) {
8080
return LogicalType::INTERVAL;
8181
}
8282
if (type->isTime()) {
83+
VELOX_DCHECK(type->equivalent(*TIME()));
8384
return LogicalType::TIME;
8485
}
8586
return LogicalType::BIGINT;

‎velox/exec/tests/utils/QueryAssertions.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ ::duckdb::Value duckValueAt<TypeKind::BIGINT>(
133133
}
134134

135135
if (type->isTime()) {
136+
VELOX_DCHECK(type->equivalent(*TIME()));
136137
// TIME is stored as milliseconds since midnight in Velox.
137138
// DuckDB TIME is stored as microseconds since midnight.
138139
const auto timeMillis = vector->as<SimpleVector<int64_t>>()->valueAt(index);
@@ -457,6 +458,7 @@ std::vector<MaterializedRow> materialize(
457458
dataChunk->GetValue(j, i).GetValue<::duckdb::date_t>()));
458459
row.push_back(value);
459460
} else if (type->isTime()) {
461+
VELOX_DCHECK(type->equivalent(*TIME()));
460462
// DuckDB TIME is in microseconds, Velox TIME is in milliseconds.
461463
auto value = variant(
462464
dataChunk->GetValue(j, i).GetValue<::duckdb::dtime_t>().micros /

‎velox/expression/CastExpr.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -844,8 +844,10 @@ void CastExpr::applyPeeled(
844844
fromType->toString(),
845845
toType->toString());
846846
} else if (fromType->isTime()) {
847+
VELOX_DCHECK(fromType->equivalent(*TIME()));
847848
result = castFromTime(rows, input, context, toType);
848849
} else if (toType->isTime()) {
850+
VELOX_DCHECK(toType->equivalent(*TIME()));
849851
result = castToTime(rows, input, context, fromType);
850852
} else if (toType->isShortDecimal()) {
851853
result = applyDecimal<int64_t>(rows, input, context, fromType, toType);

‎velox/functions/prestosql/DateTimeFunctions.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -602,6 +602,7 @@ class TimeIntervalYearMonthVectorFunction : public exec::VectorFunction {
602602
exec::EvalCtx& context,
603603
VectorPtr& result) const override {
604604
VectorPtr& timeVector = args[0]->type()->isTime() ? args[0] : args[1];
605+
VELOX_DCHECK(timeVector->type()->equivalent(*TIME()));
605606

606607
// Constant vector case
607608
// If time input is constant, create constant result - no iteration!

‎velox/functions/prestosql/types/TimeWithTimezoneRegistration.cpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ class TimeWithTimeZoneCastOperator final : public exec::CastOperator {
160160
bool isSupportedFromType(const TypePtr& other) const override {
161161
switch (other->kind()) {
162162
case TypeKind::BIGINT:
163-
return other->isTime();
163+
return other->equivalent(*TIME());
164164
case TypeKind::VARCHAR:
165165
return true;
166166
default:
@@ -171,7 +171,7 @@ class TimeWithTimeZoneCastOperator final : public exec::CastOperator {
171171
bool isSupportedToType(const TypePtr& other) const override {
172172
switch (other->kind()) {
173173
case TypeKind::BIGINT:
174-
return other->isTime();
174+
return other->equivalent(*TIME());
175175
case TypeKind::VARCHAR:
176176
return true;
177177
default:
@@ -193,7 +193,7 @@ class TimeWithTimeZoneCastOperator final : public exec::CastOperator {
193193
return;
194194
}
195195

196-
if (input.type()->isTime()) {
196+
if (input.type()->equivalent(*TIME())) {
197197
const auto& config = context.execCtx()->queryCtx()->queryConfig();
198198
const auto* sessionTimeZone = getTimeZoneFromConfig(config);
199199
const auto sessionStartTimeMs = config.sessionStartTimeMs();
@@ -246,7 +246,7 @@ class TimeWithTimeZoneCastOperator final : public exec::CastOperator {
246246

247247
switch (resultType->kind()) {
248248
case TypeKind::BIGINT:
249-
if (resultType->isTime()) {
249+
if (resultType->equivalent(*TIME())) {
250250
castToTime(input, context, rows, *result);
251251
return;
252252
}

‎velox/functions/prestosql/types/TimestampWithTimeZoneRegistration.cpp‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,7 @@ class TimestampWithTimeZoneCastOperator final : public exec::CastOperator {
256256
case TypeKind::INTEGER:
257257
return other->isDate();
258258
case TypeKind::BIGINT:
259-
return other->isTime();
259+
return other->equivalent(*TIME());
260260
default:
261261
return false;
262262
}
@@ -271,7 +271,7 @@ class TimestampWithTimeZoneCastOperator final : public exec::CastOperator {
271271
case TypeKind::INTEGER:
272272
return other->isDate();
273273
case TypeKind::BIGINT:
274-
return other->isTime();
274+
return other->equivalent(*TIME());
275275
default:
276276
return false;
277277
}
@@ -285,7 +285,8 @@ class TimestampWithTimeZoneCastOperator final : public exec::CastOperator {
285285
VectorPtr& result) const override {
286286
context.ensureWritable(rows, resultType, result);
287287

288-
if (input.typeKind() == TypeKind::BIGINT && input.type()->isTime()) {
288+
if (input.typeKind() == TypeKind::BIGINT &&
289+
input.type()->equivalent(*TIME())) {
289290
const auto& config = context.execCtx()->queryCtx()->queryConfig();
290291
const auto* sessionTimeZone = getTimeZoneFromConfig(config);
291292
const auto sessionStartTimeMs = config.sessionStartTimeMs();
@@ -368,7 +369,7 @@ class TimestampWithTimeZoneCastOperator final : public exec::CastOperator {
368369
VELOX_CHECK(resultType->isDate());
369370
castToDate(input, context, rows, *result);
370371
} else if (resultType->kind() == TypeKind::BIGINT) {
371-
VELOX_CHECK(resultType->isTime());
372+
VELOX_CHECK(resultType->equivalent(*TIME()));
372373
castToTime(input, context, rows, *result);
373374
} else {
374375
VELOX_UNSUPPORTED(

‎velox/vector/arrow/Bridge.cpp‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -294,6 +294,7 @@ const char* exportArrowFormatStr(
294294
return "i"; // int32
295295
case TypeKind::BIGINT:
296296
if (type->isTime()) {
297+
VELOX_DCHECK(type->equivalent(*TIME()));
297298
// TIME is stored as milliseconds since midnight in Velox.
298299
// Export as Arrow time32 with milliseconds unit.
299300
return "ttm";
@@ -728,12 +729,16 @@ bool isFlatScalarZeroCopy(const TypePtr& type, const ArrowOptions& options) {
728729
// - Velox's Timestamp representation (2x 64bit values) does not have an
729730
// equivalent in Arrow.
730731
// - Velox's TIME is in milliseconds, Arrow time64 is in microseconds.
732+
bool isTime = type->isTime();
733+
if (isTime) {
734+
VELOX_DCHECK(type->equivalent(*TIME()));
735+
}
731736
if (options.useDecimalTypeWidth) {
732737
// Short decimal is zero-copy.
733-
return !type->isTimestamp() && !type->isTime();
738+
return !type->isTimestamp() && !isTime;
734739
}
735740
// Short decimal requires conversion.
736-
return !type->isShortDecimal() && !type->isTimestamp() && !type->isTime();
741+
return !type->isShortDecimal() && !type->isTimestamp() && !isTime;
737742
}
738743

739744
// Returns the size of a single element of a given `type` in the target arrow
@@ -744,6 +749,7 @@ size_t getArrowElementSize(const TypePtr& type, const ArrowOptions& options) {
744749
} else if (type->isTimestamp()) {
745750
return sizeof(int64_t);
746751
} else if (type->isTime()) {
752+
VELOX_DCHECK(type->equivalent(*TIME()));
747753
// TIME is exported as Arrow time32 (int32_t).
748754
return sizeof(int32_t);
749755
}
@@ -778,6 +784,7 @@ void exportValues(
778784
if (type->kind() == TypeKind::TIMESTAMP) {
779785
gatherFromTimestampBuffer(vec, rows, options.timestampUnit, *values);
780786
} else if (type->kind() == TypeKind::BIGINT && type->isTime()) {
787+
VELOX_DCHECK(type->equivalent(*TIME()));
781788
gatherFromTimeBuffer(vec, rows, *values);
782789
} else {
783790
gatherFromBuffer(*type, *vec.values(), rows, options, *values);
@@ -2256,6 +2263,7 @@ VectorPtr importFromArrowImpl(
22562263
arrowArray.length,
22572264
arrowArray.null_count);
22582265
} else if (type->isTime()) {
2266+
VELOX_DCHECK(type->equivalent(*TIME()));
22592267
auto timeUnit = getTimeUnit(arrowSchema);
22602268
bool isTime32 =
22612269
(timeUnit == TimeUnit::kSecond || timeUnit == TimeUnit::kMilli);

‎velox/vector/fuzzer/VectorFuzzer.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ VectorPtr fuzzConstantPrimitiveImpl(
141141
return std::make_shared<ConstantVector<int128_t>>(
142142
pool, size, false, type, randLongDecimal(type, rng));
143143
} else if (type->isTime()) {
144+
VELOX_DCHECK(type->equivalent(*TIME()));
144145
return std::make_shared<ConstantVector<int64_t>>(
145146
pool, size, false, type, randTime(rng));
146147
} else {
@@ -177,6 +178,7 @@ void fuzzFlatPrimitiveImpl(
177178
} else if (vector->type()->isShortDecimal()) {
178179
flatVector->set(i, randShortDecimal(vector->type(), rng));
179180
} else if (vector->type()->isTime()) {
181+
VELOX_DCHECK(vector->type()->equivalent(*TIME()));
180182
flatVector->set(i, randTime(rng));
181183
} else {
182184
flatVector->set(i, rand<TCpp>(rng, opts.dataSpec));

‎velox/vector/fuzzer/tests/VectorFuzzerTest.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -695,7 +695,7 @@ TEST_F(VectorFuzzerTest, time) {
695695
// Test flat TIME vector.
696696
auto timeVector = fuzzer.fuzzFlat(TIME());
697697
ASSERT_EQ(VectorEncoding::Simple::FLAT, timeVector->encoding());
698-
ASSERT_TRUE(timeVector->type()->isTime());
698+
ASSERT_TRUE(timeVector->type()->equivalent(*TIME()));
699699
ASSERT_EQ(vectorSize, timeVector->size());
700700

701701
auto flatTimeVector = timeVector->as<FlatVector<int64_t>>();
@@ -704,7 +704,7 @@ TEST_F(VectorFuzzerTest, time) {
704704
// Test constant TIME vector.
705705
auto constTimeVector = fuzzer.fuzzConstant(TIME(), vectorSize);
706706
ASSERT_EQ(VectorEncoding::Simple::CONSTANT, constTimeVector->encoding());
707-
ASSERT_TRUE(constTimeVector->type()->isTime());
707+
ASSERT_TRUE(constTimeVector->type()->equivalent(*TIME()));
708708
ASSERT_EQ(vectorSize, constTimeVector->size());
709709

710710
// Verify constant TIME value is in valid range.
@@ -719,7 +719,7 @@ TEST_F(VectorFuzzerTest, time) {
719719
auto dictTimeVector =
720720
fuzzer.fuzzDictionary(fuzzer.fuzzFlat(TIME(), 100), 500);
721721
ASSERT_EQ(VectorEncoding::Simple::DICTIONARY, dictTimeVector->encoding());
722-
ASSERT_TRUE(dictTimeVector->type()->isTime());
722+
ASSERT_TRUE(dictTimeVector->type()->equivalent(*TIME()));
723723

724724
// Verify all TIME values in dictionary are in valid range.
725725
auto dictVector = dictTimeVector->as<DictionaryVector<int64_t>>();

0 commit comments

Comments
 (0)