This is an automated email from the ASF dual-hosted git repository.
rui-mo pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gluten.git
The following commit(s) were added to refs/heads/main by this push:
new e63d43ef4f [GLUTEN-12538][CORE] Preserve TIMESTAMP_NTZ literals
(#13041)
e63d43ef4f is described below
commit e63d43ef4f677c860b2a3c27f768572d46725a8c
Author: Felipe Pessoto <[email protected]>
AuthorDate: Thu Nov 5 14:22:30 2026 -0800
[GLUTEN-12538][CORE] Preserve TIMESTAMP_NTZ literals (#13041)
---
cpp/velox/substrait/SubstraitParser.cc | 2 +-
cpp/velox/substrait/VeloxToSubstraitExpr.cc | 119 ++++++------
cpp/velox/substrait/VeloxToSubstraitExpr.h | 5 +-
cpp/velox/tests/SubstraitVeloxExprConverterTest.cc | 95 ++++++++++
cpp/velox/tests/VeloxSubstraitRoundTripTest.cc | 106 +++++++++++
.../substrait/expression/ExpressionBuilder.java | 3 +
.../substrait/expression/StructLiteralNode.java | 2 +-
.../apache/gluten/expression/ConverterUtils.scala | 2 +
.../extension/columnar/validator/Validators.scala | 27 ++-
.../gluten/utils/velox/VeloxTestSettings.scala | 1 +
.../spark/sql/GlutenTimestampNtzLiteralSuite.scala | 137 ++++++++++++++
.../gluten/utils/velox/VeloxTestSettings.scala | 1 +
.../spark/sql/GlutenTimestampNtzLiteralSuite.scala | 1 +
.../GlutenTimestampNtzValidatorSuite.scala | 201 +++++++++++++++++++++
14 files changed, 638 insertions(+), 64 deletions(-)
diff --git a/cpp/velox/substrait/SubstraitParser.cc
b/cpp/velox/substrait/SubstraitParser.cc
index 3333810e01..031840fbf1 100644
--- a/cpp/velox/substrait/SubstraitParser.cc
+++ b/cpp/velox/substrait/SubstraitParser.cc
@@ -426,7 +426,7 @@ bool SubstraitParser::getLiteralValue(const
::substrait::Expression::Literal& li
template <>
Timestamp SubstraitParser::getLiteralValue(const
::substrait::Expression::Literal& literal) {
- return Timestamp::fromMicros(literal.timestamp_tz());
+ return Timestamp::fromMicros(literal.has_timestamp() ? literal.timestamp() :
literal.timestamp_tz());
}
template <>
diff --git a/cpp/velox/substrait/VeloxToSubstraitExpr.cc
b/cpp/velox/substrait/VeloxToSubstraitExpr.cc
index f9cd1d2880..d335316331 100644
--- a/cpp/velox/substrait/VeloxToSubstraitExpr.cc
+++ b/cpp/velox/substrait/VeloxToSubstraitExpr.cc
@@ -23,10 +23,21 @@ namespace gluten {
namespace {
const ::substrait::Expression_Literal& toSubstraitNullLiteral(
google::protobuf::Arena& arena,
- const velox::TypeKind& typeKind) {
+ const velox::TypePtr& type) {
::substrait::Expression_Literal* substraitField =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
- switch (typeKind) {
+ substraitField->set_nullable(true);
+ auto nestedType = type;
+ while (nestedType->isArray()) {
+ nestedType = nestedType->asArray().elementType();
+ }
+ if (nestedType->kind() == TypeKind::TIMESTAMP) {
+ VeloxToSubstraitTypeConvertor typeConvertor;
+
substraitField->mutable_null()->MergeFrom(typeConvertor.toSubstraitType(arena,
type));
+ return *substraitField;
+ }
+ // Preserve physical-kind encoding for non-timestamp nulls.
+ switch (type->kind()) {
case velox::TypeKind::BOOLEAN: {
::substrait::Type_Boolean* nullValue =
google::protobuf::Arena::CreateMessage<::substrait::Type_Boolean>(&arena);
nullValue->set_nullability(::substrait::Type_Nullability_NULLABILITY_NULLABLE);
@@ -92,16 +103,32 @@ const ::substrait::Expression_Literal&
toSubstraitNullLiteral(
break;
}
default: {
- VELOX_UNSUPPORTED("Unsupported type '{}'",
TypeKindName::toName(typeKind));
+ VELOX_UNSUPPORTED("Unsupported type '{}'",
TypeKindName::toName(type->kind()));
}
}
- substraitField->set_nullable(true);
return *substraitField;
}
+const ::substrait::Expression_Literal&
+toSubstraitTimestampLiteral(google::protobuf::Arena& arena, const Timestamp&
value, const TypePtr& type) {
+ auto* literal =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
+ const auto micros = value.toMicros();
+ if (type->equivalent(*TIMESTAMP_UTC())) {
+ literal->set_timestamp(micros);
+ } else {
+ literal->set_timestamp_tz(micros);
+ }
+ literal->set_nullable(false);
+ return *literal;
+}
+
const ::substrait::Expression_Literal& toSubstraitNotNullLiteral(
google::protobuf::Arena& arena,
- const velox::variant& variantValue) {
+ const velox::variant& variantValue,
+ const TypePtr& logicalType) {
+ if (variantValue.kind() == TypeKind::TIMESTAMP) {
+ return toSubstraitTimestampLiteral(arena,
variantValue.value<TypeKind::TIMESTAMP>(), logicalType);
+ }
::substrait::Expression_Literal* literalExpr =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
switch (variantValue.kind()) {
@@ -133,12 +160,6 @@ const ::substrait::Expression_Literal&
toSubstraitNotNullLiteral(
literalExpr->set_fp64(variantValue.value<TypeKind::DOUBLE>());
break;
}
- case velox::TypeKind::TIMESTAMP: {
- auto vTimeStamp = variantValue.value<TypeKind::TIMESTAMP>();
- auto micros = vTimeStamp.getSeconds() * 1000000 + vTimeStamp.getNanos()
/ 1000;
- literalExpr->set_timestamp_tz(micros);
- break;
- }
case velox::TypeKind::VARCHAR: {
auto vCharValue = variantValue.value<StringView>();
auto* sVarChar =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal::VarChar>(&arena);
@@ -243,18 +264,6 @@ const ::substrait::Expression_Literal&
toSubstraitNotNullLiteral<TypeKind::DOUBL
return *literalExpr;
}
-template <>
-const ::substrait::Expression_Literal&
toSubstraitNotNullLiteral<TypeKind::TIMESTAMP>(
- google::protobuf::Arena& arena,
- const Timestamp& value) {
- ::substrait::Expression_Literal* literalExpr =
-
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
- auto micros = value.getSeconds() * 1000000 + value.getNanos() / 1000;
- literalExpr->set_timestamp_tz(micros);
- literalExpr->set_nullable(false);
- return *literalExpr;
-}
-
template <>
const ::substrait::Expression_Literal&
toSubstraitNotNullLiteral<TypeKind::VARCHAR>(
google::protobuf::Arena& arena,
@@ -269,6 +278,21 @@ const ::substrait::Expression_Literal&
toSubstraitNotNullLiteral<TypeKind::VARCH
return *literalExpr;
}
+template <TypeKind kind>
+const ::substrait::Expression_Literal&
+scalarVectorToLiteral(google::protobuf::Arena& arena, const VectorPtr& vector,
vector_size_t row) {
+ if (vector->isNullAt(row)) {
+ return toSubstraitNullLiteral(arena, vector->type());
+ }
+ using T = typename TypeTraits<kind>::NativeType;
+ const auto value = vector->as<SimpleVector<T>>()->valueAt(row);
+ if constexpr (kind == TypeKind::TIMESTAMP) {
+ return toSubstraitTimestampLiteral(arena, value, vector->type());
+ } else {
+ return toSubstraitNotNullLiteral<kind>(arena, value);
+ }
+}
+
template <TypeKind Kind>
void arrayVectorToLiteral(
google::protobuf::Arena& arena,
@@ -276,15 +300,9 @@ void arrayVectorToLiteral(
::substrait::Expression_Literal_List* listLiteral,
vector_size_t offset,
vector_size_t size) {
- using T = typename TypeTraits<Kind>::NativeType;
- auto elements = arrayVector->elements()->as<SimpleVector<T>>();
for (auto i = offset; i < offset + size; ++i) {
::substrait::Expression_Literal* childLiteral = listLiteral->add_values();
- if (elements->isNullAt(i)) {
- childLiteral->MergeFrom(toSubstraitNullLiteral(arena, Kind));
- } else {
- childLiteral->MergeFrom(toSubstraitNotNullLiteral<Kind>(arena,
elements->valueAt(i)));
- }
+ childLiteral->MergeFrom(scalarVectorToLiteral<Kind>(arena,
arrayVector->elements(), i));
}
}
@@ -292,24 +310,9 @@ template <TypeKind kind>
void convertVectorValue(
google::protobuf::Arena& arena,
const velox::VectorPtr& vectorValue,
- ::substrait::Expression_Nested_Struct* litValue,
- ::substrait::Expression_Literal* substraitField) {
- const TypePtr& childType = vectorValue->type();
-
- using T = typename TypeTraits<kind>::NativeType;
-
- auto childToFlatVec = vectorValue->as<SimpleVector<T>>();
-
- // Get the batchSize and convert each value in it.
- vector_size_t flatVecSize = childToFlatVec->size();
- for (int64_t i = 0; i < flatVecSize; i++) {
- substraitField = litValue->add_fields()->mutable_literal();
- if (childToFlatVec->isNullAt(i)) {
- // Process the null value.
- substraitField->MergeFrom(toSubstraitNullLiteral(arena,
childType->kind()));
- } else {
- substraitField->MergeFrom(toSubstraitNotNullLiteral<kind>(arena,
childToFlatVec->valueAt(i)));
- }
+ ::substrait::Expression_Nested_Struct* litValue) {
+ for (vector_size_t i = 0; i < vectorValue->size(); ++i) {
+
litValue->add_fields()->mutable_literal()->MergeFrom(scalarVectorToLiteral<kind>(arena,
vectorValue, i));
}
}
@@ -516,20 +519,21 @@ const ::substrait::Expression_Literal&
VeloxToSubstraitExprConvertor::toSubstrai
if (constExpr->hasValueVector()) {
return toSubstraitLiteral(arena, constExpr->valueVector(), litValue);
} else {
- return toSubstraitLiteral(arena, constExpr->value());
+ return toSubstraitLiteral(arena, constExpr->value(), constExpr->type());
}
}
const ::substrait::Expression_Literal&
VeloxToSubstraitExprConvertor::toSubstraitLiteral(
google::protobuf::Arena& arena,
- const velox::variant& variantValue) {
+ const velox::variant& variantValue,
+ const TypePtr& logicalType) {
::substrait::Expression_Literal* literalExpr =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
if (variantValue.isNull()) {
- literalExpr->MergeFrom(toSubstraitNullLiteral(arena, variantValue.kind()));
+ literalExpr->MergeFrom(toSubstraitNullLiteral(arena, logicalType));
} else {
- literalExpr->MergeFrom(toSubstraitNotNullLiteral(arena, variantValue));
+ literalExpr->MergeFrom(toSubstraitNotNullLiteral(arena, variantValue,
logicalType));
}
return *literalExpr;
}
@@ -572,7 +576,7 @@ const ::substrait::Expression_Literal&
VeloxToSubstraitExprConvertor::toSubstrai
if (constantVector->typeKind() == TypeKind::ARRAY) {
if (constantVector->isNullAt(0)) {
// Process the null value.
- substraitField->MergeFrom(toSubstraitNullLiteral(arena,
constantVector->typeKind()));
+ substraitField->MergeFrom(toSubstraitNullLiteral(arena,
constantVector->type()));
return *substraitField;
}
auto encoding = constantVector->valueVector()->encoding();
@@ -599,8 +603,13 @@ const ::substrait::Expression_Literal&
VeloxToSubstraitExprConvertor::toSubstrai
::substrait::Expression_Literal* substraitField =
google::protobuf::Arena::CreateMessage<::substrait::Expression_Literal>(&arena);
if (vectorValue->isScalar()) {
- VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(
- convertVectorValue, vectorValue->type()->kind(), arena, vectorValue,
litValue, substraitField);
+ if (litValue == nullptr) {
+ VELOX_CHECK_GT(vectorValue->size(), 0, "A scalar constant must contain a
value");
+ substraitField->MergeFrom(
+ VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(scalarVectorToLiteral,
vectorValue->typeKind(), arena, vectorValue, 0));
+ return *substraitField;
+ }
+ VELOX_DYNAMIC_SCALAR_TYPE_DISPATCH(convertVectorValue,
vectorValue->type()->kind(), arena, vectorValue, litValue);
return *substraitField;
}
diff --git a/cpp/velox/substrait/VeloxToSubstraitExpr.h
b/cpp/velox/substrait/VeloxToSubstraitExpr.h
index 83b9a5260f..dafbbc2e8f 100644
--- a/cpp/velox/substrait/VeloxToSubstraitExpr.h
+++ b/cpp/velox/substrait/VeloxToSubstraitExpr.h
@@ -47,6 +47,7 @@ class VeloxToSubstraitExprConvertor {
/// @param constExpr Velox Constant expression needed to be converted.
/// @param litValue The Nested.Struct the converted literals are appended
to, each wrapped in an
/// Expression. Substrait models struct-valued expressions as
Expression.Nested.Struct.
+ /// If null, converts only the constant value at index zero.
/// @return A pointer to Substrait Literal expression object allocated on
/// the arena and representing the input Velox Constant expression.
const ::substrait::Expression_Literal& toSubstraitExpr(
@@ -67,6 +68,7 @@ class VeloxToSubstraitExprConvertor {
/// Convert Velox vector to Substrait literal. One literal per row is
appended to litValue,
/// each wrapped in an Expression so it fits Substrait's
Expression.Nested.Struct container.
+ /// If litValue is null, converts only the scalar value at index zero.
const ::substrait::Expression_Literal& toSubstraitLiteral(
google::protobuf::Arena& arena,
const velox::VectorPtr& vectorValue,
@@ -88,7 +90,8 @@ class VeloxToSubstraitExprConvertor {
/// Convert Velox variant to Substrait Literal Expression.
const ::substrait::Expression_Literal& toSubstraitLiteral(
google::protobuf::Arena& arena,
- const velox::variant& variantValue);
+ const velox::variant& variantValue,
+ const velox::TypePtr& logicalType);
/// Convert values in Velox array vector to Substrait Literal List.
const ::substrait::Expression_Literal_List&
diff --git a/cpp/velox/tests/SubstraitVeloxExprConverterTest.cc
b/cpp/velox/tests/SubstraitVeloxExprConverterTest.cc
index 32b94da2b5..c3e02b24cb 100644
--- a/cpp/velox/tests/SubstraitVeloxExprConverterTest.cc
+++ b/cpp/velox/tests/SubstraitVeloxExprConverterTest.cc
@@ -16,6 +16,7 @@
*/
#include "substrait/SubstraitToVeloxExpr.h"
+#include "substrait/VeloxToSubstraitExpr.h"
#include "velox/common/base/tests/GTestUtils.h"
#include "velox/core/QueryConfig.h"
@@ -30,6 +31,100 @@ namespace gluten {
class SubstraitVeloxExprConverterExecutionTest : public
exec::test::OperatorTestBase {};
+TEST_F(SubstraitVeloxExprConverterExecutionTest, timestampLiterals) {
+ const std::unordered_map<uint64_t, std::string> functionMap;
+ SubstraitVeloxExprConverter converter(pool(), functionMap);
+ for (const auto& type : std::vector<TypePtr>{TIMESTAMP_UTC(), TIMESTAMP()}) {
+ SCOPED_TRACE(type->toString());
+ for (const int64_t micros : {-1LL, 0LL, 1LL, 1'704'067'200'123'456LL}) {
+ SCOPED_TRACE(micros);
+ ::substrait::Expression expression;
+ auto* literal = expression.mutable_literal();
+ if (type->equivalent(*TIMESTAMP_UTC())) {
+ literal->set_timestamp(micros);
+ } else {
+ literal->set_timestamp_tz(micros);
+ }
+ auto result =
+ std::dynamic_pointer_cast<const
core::ConstantTypedExpr>(converter.toVeloxExpr(expression, ROW({}, {})));
+ ASSERT_NE(result, nullptr);
+ EXPECT_TRUE(type->equivalent(*result->type()));
+ EXPECT_EQ(result->value().value<Timestamp>().toMicros(), micros);
+ }
+ }
+}
+
+TEST_F(SubstraitVeloxExprConverterExecutionTest, timestampLiteralRoundTrip) {
+ const std::unordered_map<uint64_t, std::string> functionMap;
+ SubstraitVeloxExprConverter reader(pool(), functionMap);
+ auto extensions = std::make_shared<SubstraitExtensionCollector>();
+ VeloxToSubstraitExprConvertor writer(extensions);
+ const std::vector<std::optional<Timestamp>> values = {
+ Timestamp(-1, 999'999'000), Timestamp(1'704'067'200, 123'456'000),
std::nullopt};
+
+ for (const auto& type : std::vector<TypePtr>{TIMESTAMP_UTC(), TIMESTAMP()}) {
+ SCOPED_TRACE(type->toString());
+ for (const auto& value : values) {
+ SCOPED_TRACE(value.has_value() ? value->toString() : "null");
+ for (const bool vectorBacked : {false, true}) {
+ SCOPED_TRACE(vectorBacked);
+ auto constant = vectorBacked
+ ? std::make_shared<const
core::ConstantTypedExpr>(makeConstant<Timestamp>(value, 1, type))
+ : std::make_shared<const core::ConstantTypedExpr>(
+ type, value.has_value() ? variant(value.value()) :
variant::null(TypeKind::TIMESTAMP));
+ google::protobuf::Arena arena;
+ ::substrait::Expression expression;
+ expression.mutable_literal()->MergeFrom(writer.toSubstraitExpr(arena,
constant));
+ const auto& literal = expression.literal();
+ if (value.has_value()) {
+ if (type->equivalent(*TIMESTAMP_UTC())) {
+ ASSERT_TRUE(literal.has_timestamp());
+ EXPECT_EQ(literal.timestamp(), value->toMicros());
+ } else {
+ ASSERT_TRUE(literal.has_timestamp_tz());
+ EXPECT_EQ(literal.timestamp_tz(), value->toMicros());
+ }
+ } else {
+ ASSERT_TRUE(literal.has_null());
+
EXPECT_TRUE(type->equivalent(*SubstraitParser::parseType(literal.null())));
+ }
+ auto result =
+ std::dynamic_pointer_cast<const
core::ConstantTypedExpr>(reader.toVeloxExpr(expression, ROW({}, {})));
+ ASSERT_NE(result, nullptr);
+ EXPECT_TRUE(type->equivalent(*result->type()));
+ ASSERT_EQ(result->isNull(), !value.has_value());
+ if (value.has_value()) {
+ EXPECT_EQ(result->value().value<Timestamp>(), value.value());
+ }
+ }
+ }
+ }
+}
+
+TEST_F(SubstraitVeloxExprConverterExecutionTest, dateLiteralEncoding) {
+ auto extensions = std::make_shared<SubstraitExtensionCollector>();
+ VeloxToSubstraitExprConvertor writer(extensions);
+ for (const auto& value : std::vector<std::optional<int32_t>>{std::nullopt,
19'723}) {
+ SCOPED_TRACE(value.has_value() ? std::to_string(value.value()) : "null");
+ for (const bool vectorBacked : {false, true}) {
+ SCOPED_TRACE(vectorBacked);
+ auto constant = vectorBacked
+ ? std::make_shared<const
core::ConstantTypedExpr>(makeConstant<int32_t>(value, 1, DATE()))
+ : std::make_shared<const core::ConstantTypedExpr>(
+ DATE(), value.has_value() ? variant(value.value()) :
variant::null(TypeKind::INTEGER));
+ google::protobuf::Arena arena;
+ const auto& literal = writer.toSubstraitExpr(arena, constant);
+ if (value.has_value()) {
+ ASSERT_TRUE(literal.has_i32());
+ EXPECT_EQ(literal.i32(), value.value());
+ } else {
+ ASSERT_TRUE(literal.has_null());
+ EXPECT_TRUE(literal.null().has_i32());
+ }
+ }
+ }
+}
+
// Regression test for a SIGSEGV in
// SubstraitVeloxExprConverter::toVeloxExpr(Expression::FieldReference, ...).
// The direct-reference loop descends one nested struct_field at a time with
diff --git a/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc
b/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc
index 6a9b2447d0..ae1922c329 100644
--- a/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc
+++ b/cpp/velox/tests/VeloxSubstraitRoundTripTest.cc
@@ -122,6 +122,21 @@ class VeloxSubstraitRoundTripTest : public
OperatorTestBase {
}
}
+ void assertTypedPlanConversion(const ::substrait::Plan& plan, const
RowVectorPtr& expected) {
+ auto config =
std::make_shared<facebook::velox::config::ConfigBase>(std::unordered_map<std::string,
std::string>());
+ SubstraitToVeloxPlanConverter converter(
+ pool_.get(),
+ config.get(),
+ std::vector<std::shared_ptr<ResultIterator>>{},
+ VeloxConnectorIds{},
+ std::nullopt,
+ std::nullopt,
+ true);
+ auto result = converter.toVeloxPlan(plan);
+ ASSERT_TRUE(expected->type()->equivalent(*result->outputType())) <<
result->outputType()->toString();
+ assertQuery(result, expected);
+ }
+
std::shared_ptr<VeloxToSubstraitPlanConvertor> veloxConvertor_ =
std::make_shared<VeloxToSubstraitPlanConvertor>();
};
@@ -431,6 +446,78 @@ TEST_F(VeloxSubstraitRoundTripTest, notNullLiteral) {
assertPlanConversion(plan, "SELECT true, 23, 45, 678, 910, 1.23, 4.56,
'789'");
}
+TEST_F(VeloxSubstraitRoundTripTest, timestampLiteralProjection) {
+ const auto beforeEpoch = Timestamp(-1, 999'999'000);
+ const auto afterEpoch = Timestamp(1'704'067'200, 123'456'000);
+ auto input = makeRowVector(ROW({}, {}), 1);
+ for (const auto& type : std::vector<TypePtr>{TIMESTAMP_UTC(), TIMESTAMP()}) {
+ SCOPED_TRACE(type->toString());
+ auto expected = makeRowVector(
+ {makeFlatVector<Timestamp>({beforeEpoch}, type),
+ makeFlatVector<Timestamp>({afterEpoch}, type),
+ makeNullableFlatVector<Timestamp>({std::nullopt}, type)});
+ auto plan = PlanBuilder(pool_.get())
+ .values({input})
+ .projectExpressions(
+ {makeConstantExpr(type, variant(beforeEpoch)),
+ makeConstantExpr(type, variant(afterEpoch)),
+ makeConstantExpr(type,
variant::null(TypeKind::TIMESTAMP))})
+ .planNode();
+ assertQuery(plan, expected);
+ google::protobuf::Arena arena;
+ const auto& substraitPlan = veloxConvertor_->toSubstrait(arena, plan);
+ ASSERT_NO_FATAL_FAILURE(assertTypedPlanConversion(substraitPlan,
expected));
+ }
+}
+
+TEST_F(VeloxSubstraitRoundTripTest, timestampValues) {
+ const auto beforeEpoch = Timestamp(-1, 999'999'000);
+ const auto afterEpoch = Timestamp(1'704'067'200, 123'456'000);
+ auto input = makeRowVector(
+ {makeNullableFlatVector<Timestamp>({beforeEpoch, afterEpoch,
std::nullopt}, TIMESTAMP_UTC()),
+ makeNullableFlatVector<Timestamp>({afterEpoch, std::nullopt,
beforeEpoch}, TIMESTAMP())});
+ auto plan = PlanBuilder(pool_.get()).values({input}).planNode();
+ google::protobuf::Arena arena;
+ const auto& substraitPlan = veloxConvertor_->toSubstrait(arena, plan);
+ const auto& fields =
substraitPlan.relations(0).root().input().read().virtual_table().expressions(0).fields();
+ ASSERT_EQ(fields.size(), 6);
+ ASSERT_TRUE(fields.Get(0).literal().has_timestamp());
+ EXPECT_EQ(fields.Get(0).literal().timestamp(), beforeEpoch.toMicros());
+ ASSERT_TRUE(fields.Get(2).literal().has_null());
+ EXPECT_TRUE(fields.Get(2).literal().null().has_precision_timestamp());
+ ASSERT_TRUE(fields.Get(3).literal().has_timestamp_tz());
+ EXPECT_EQ(fields.Get(3).literal().timestamp_tz(), afterEpoch.toMicros());
+ ASSERT_TRUE(fields.Get(4).literal().has_null());
+ EXPECT_TRUE(fields.Get(4).literal().null().has_precision_timestamp_tz());
+ assertTypedPlanConversion(substraitPlan, input);
+}
+
+TEST_F(VeloxSubstraitRoundTripTest, timestampArrayLiteral) {
+ const auto beforeEpoch = Timestamp(-1, 999'999'000);
+ const auto afterEpoch = Timestamp(1'704'067'200, 123'456'000);
+ auto input = makeRowVector(ROW({}, {}), 1);
+ for (const auto& type : std::vector<TypePtr>{TIMESTAMP_UTC(), TIMESTAMP()}) {
+ SCOPED_TRACE(type->toString());
+ auto values = makeNullableArrayVector<Timestamp>({{beforeEpoch,
std::nullopt, afterEpoch}}, ARRAY(type));
+ auto empty = makeArrayVector<Timestamp>({{}}, type);
+ auto nullArray = BaseVector::createNullConstant(ARRAY(type), 1,
pool_.get());
+ auto nullNestedArray = BaseVector::createNullConstant(ARRAY(ARRAY(type)),
1, pool_.get());
+ auto expected = makeRowVector({values, empty, nullArray, nullNestedArray});
+ auto plan = PlanBuilder(pool_.get())
+ .values({input})
+ .projectExpressions(
+ {makeConstantExpr(values),
+ makeConstantExpr(empty),
+ makeConstantExpr(nullArray),
+ makeConstantExpr(nullNestedArray)})
+ .planNode();
+ assertQuery(plan, expected);
+ google::protobuf::Arena arena;
+ const auto& substraitPlan = veloxConvertor_->toSubstrait(arena, plan);
+ ASSERT_NO_FATAL_FAILURE(assertTypedPlanConversion(substraitPlan,
expected));
+ }
+}
+
TEST_F(VeloxSubstraitRoundTripTest, arrayLiteral) {
auto vectors = makeRowVector(ROW({}), 1);
auto plan = PlanBuilder(pool_.get())
@@ -470,6 +557,25 @@ TEST_F(VeloxSubstraitRoundTripTest, arrayLiteral) {
"array[], array[array[1,2,3], array[4,5]]");
}
+TEST_F(VeloxSubstraitRoundTripTest, dateArrayNullPositions) {
+ auto leadingNull = makeNullableArrayVector<int32_t>({{std::nullopt,
19'723}}, ARRAY(DATE()));
+ auto trailingNull = makeNullableArrayVector<int32_t>({{19'723,
std::nullopt}}, ARRAY(DATE()));
+ auto comparison = std::make_shared<const core::CallTypedExpr>(
+ BOOLEAN(),
+ std::vector<core::TypedExprPtr>{makeConstantExpr(leadingNull),
makeConstantExpr(trailingNull)},
+ "equalto");
+ auto input = makeRowVector(ROW({}, {}), 1);
+ auto expected = makeRowVector({makeFlatVector<bool>({false})});
+ auto plan = PlanBuilder(pool_.get())
+ .values({input})
+
.projectExpressions(std::vector<core::TypedExprPtr>{comparison})
+ .planNode();
+ assertQuery(plan, expected);
+ google::protobuf::Arena arena;
+ const auto& substraitPlan = veloxConvertor_->toSubstrait(arena, plan);
+ assertTypedPlanConversion(substraitPlan, expected);
+}
+
// Disabled due to https://github.com/facebookincubator/velox/pull/17318.
TEST_F(VeloxSubstraitRoundTripTest, DISABLED_dateType) {
auto a = makeFlatVector<int32_t>({0, 1});
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/ExpressionBuilder.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/ExpressionBuilder.java
index 9d5b5c9a10..9f41fce803 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/ExpressionBuilder.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/ExpressionBuilder.java
@@ -177,6 +177,9 @@ public class ExpressionBuilder {
if (typeNode instanceof TimestampTypeNode) {
return makeTimestampLiteral((Long) obj, typeNode);
}
+ if (typeNode instanceof TimestampNTZTypeNode) {
+ return new TimestampNTZLiteralNode((Long) obj, typeNode);
+ }
if (typeNode instanceof StringTypeNode) {
return makeStringLiteral(obj.toString(), typeNode);
}
diff --git
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/StructLiteralNode.java
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/StructLiteralNode.java
index 981fa36f4c..aadeb0f035 100644
---
a/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/StructLiteralNode.java
+++
b/gluten-substrait/src/main/java/org/apache/gluten/substrait/expression/StructLiteralNode.java
@@ -58,7 +58,7 @@ public class StructLiteralNode extends
LiteralNodeWithValue<InternalRow> {
if (type instanceof DateTypeNode) {
return ExpressionBuilder.makeLiteral(value.getInt(index), type);
}
- if (type instanceof TimestampTypeNode) {
+ if (type instanceof TimestampTypeNode || type instanceof
TimestampNTZTypeNode) {
return ExpressionBuilder.makeLiteral(value.getLong(index), type);
}
if (type instanceof StringTypeNode) {
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala
index 2cd0fa9c06..284f21dbb5 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala
@@ -328,6 +328,8 @@ object ConverterUtils extends Logging {
DecimalType(d.precision, d.scale)
case _: TimestampTypeNode =>
TimestampType
+ case _: TimestampNTZTypeNode =>
+ TimestampNTZType
case m: MapNode =>
val keyType = parseFromTypeNode(m.getKeyType)
val valueType = parseFromTypeNode(m.getValueType)
diff --git
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala
index 07f847f70a..792e1c0edd 100644
---
a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala
+++
b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala
@@ -25,7 +25,8 @@ import
org.apache.gluten.extension.columnar.offload.OffloadSingleNode
import org.apache.gluten.sql.shims.SparkShimLoader
import org.apache.spark.internal.Logging
-import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Cast,
ConvertTimezone, Expression, GetStructField, Hour, IsNotNull, IsNull,
MakeTimestamp, Minute, Second, TimestampAdd}
+import org.apache.spark.sql.catalyst.analysis.PartitionSpec
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Cast,
ConvertTimezone, Expression, GetStructField, Hour, IsNotNull, IsNull,
MakeTimestamp, Minute, Second, TimestampAdd, WindowFrame, WindowSpecDefinition}
import org.apache.spark.sql.execution._
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec,
ObjectHashAggregateExec, SortAggregateExec}
import org.apache.spark.sql.execution.datasources.WriteFilesExec
@@ -93,9 +94,12 @@ object Validators {
builder.add(new FallbackByTestInjects())
}
- /** Fails validation if a plan node's input or output schema contains
TimestampNTZType. */
+ /**
+ * Fails validation for unsupported TimestampNTZType schemas, and for NTZ
expressions on
+ * backends without TimestampNTZ support.
+ */
def fallbackByTimestampNTZ(): Validator.Builder = {
- builder.add(new FallbackByTimestampNTZ(veloxConf))
+ builder.add(new FallbackByTimestampNTZ(veloxConf,
settings.supportTimestampNtz))
}
/**
@@ -232,7 +236,10 @@ object Validators {
}
}
- private class FallbackByTimestampNTZ(veloxConf: Option[Any]) extends
Validator {
+ private[validator] class FallbackByTimestampNTZ(
+ veloxConf: Option[Any],
+ backendSupportsTimestampNtz: Boolean)
+ extends Validator {
// Check if TimestampNTZ validation is enabled via VeloxConfig
// Default to true (enabled) if VeloxConfig is not available or method
call fails
private val enableValidation: Boolean = veloxConf
@@ -247,8 +254,6 @@ object Validators {
}
.getOrElse(true)
- private val backendSupportsTimestampNtz =
BackendsApiManager.getSettings.supportTimestampNtz
-
override def validate(plan: SparkPlan): Validator.OutCome = {
def containsNTZ(dataType: DataType): Boolean = dataType match {
case TimestampNTZType => true
@@ -267,6 +272,16 @@ object Validators {
val hasNTZ = plan.output.exists(a => containsNTZ(a.dataType)) ||
plan.children.exists(_.output.exists(a => containsNTZ(a.dataType)))
if (!hasNTZ) {
+ if (
+ !backendSupportsTimestampNtz &&
+ plan.expressions.exists(_.exists {
+ // These metadata nodes have no dataType; exists still visits
their children.
+ case _: WindowSpecDefinition | _: WindowFrame | _: PartitionSpec
=> false
+ case expr => containsNTZ(expr.dataType)
+ })
+ ) {
+ return fail(s"${plan.nodeName} has TimestampNTZType in expressions")
+ }
return pass()
}
diff --git
a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
index d96a2b0bd8..1e29452c78 100644
---
a/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
+++
b/gluten-ut/spark41/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
@@ -205,6 +205,7 @@ class VeloxTestSettings extends BackendTestSettings {
.exclude("default")
// FIXME(yma11): ObjectType is not covered in RowEncoder/Serializer in
vanilla spark
.exclude("SPARK-37967: Literal.create support ObjectType")
+ enableSuite[GlutenTimestampNtzLiteralSuite]
enableSuite[GlutenMathExpressionsSuite]
// Spark round UT for round(3.1415,3) is not correct.
.exclude("round/bround/floor/ceil")
diff --git
a/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
new file mode 100644
index 0000000000..28c1168b15
--- /dev/null
+++
b/gluten-ut/spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql
+
+import org.apache.gluten.config.GlutenConfig
+import org.apache.gluten.execution.ProjectExecTransformer
+import org.apache.gluten.substrait.expression.ExpressionBuilder
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, StructField, StructType,
TimestampNTZType, TimestampType}
+
+import java.time.LocalDateTime
+
+class GlutenTimestampNtzLiteralSuite extends GlutenSQLTestsTrait {
+
+ testGluten("timestamp literal serialization preserves logical types") {
+ Seq(TimestampNTZType, TimestampType).foreach {
+ dataType =>
+ Seq(-1L, 1704067200123456L).foreach {
+ micros =>
+ val scalar =
+ ExpressionBuilder.makeLiteral(Long.box(micros), dataType,
true).toProtobuf.getLiteral
+ val array = ExpressionBuilder
+ .makeLiteral(
+ new GenericArrayData(Array[Any](micros, null)),
+ ArrayType(dataType),
+ true)
+ .toProtobuf.getLiteral
+ val struct = ExpressionBuilder
+ .makeLiteral(
+ InternalRow(micros, null),
+ StructType(Seq(StructField("value", dataType),
StructField("missing", dataType))),
+ true)
+ .toProtobuf.getLiteral
+ Seq(scalar, array.getList.getValues(0),
struct.getStruct.getFields(0)).foreach {
+ literal =>
+ if (dataType == TimestampNTZType) {
+ assert(literal.hasTimestamp)
+ assert(literal.getTimestamp == micros)
+ } else {
+ assert(literal.hasTimestampTz)
+ assert(literal.getTimestampTz == micros)
+ }
+ }
+ Seq(
+ ExpressionBuilder.makeLiteral(null, dataType,
true).toProtobuf.getLiteral,
+ array.getList.getValues(1),
+ struct.getStruct.getFields(1)).foreach {
+ literal =>
+ assert(literal.hasNull)
+ if (dataType == TimestampNTZType) {
+ assert(literal.getNull.hasPrecisionTimestamp)
+ } else {
+ assert(literal.getNull.hasPrecisionTimestampTz)
+ }
+ }
+ }
+ }
+ }
+
+ testGluten("timestamp_ntz literals in native expressions") {
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "false",
+ GlutenConfig.GLUTEN_ANSI_FALLBACK_ENABLED.key -> "false",
+ "spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation"
-> "false"
+ ) {
+ val result = spark.range(2).selectExpr(
+ "timestampadd(MICROSECOND, id, cast('1969-12-31 23:59:59.999999' as
timestamp_ntz))",
+ "timestampadd(MICROSECOND, id, cast('2024-01-01 00:00:00.123456' as
timestamp_ntz))"
+ )
+ checkAnswer(
+ result,
+ Seq(
+ Row(
+ LocalDateTime.parse("1969-12-31T23:59:59.999999"),
+ LocalDateTime.parse("2024-01-01T00:00:00.123456")),
+ Row(
+ LocalDateTime.parse("1970-01-01T00:00:00"),
+ LocalDateTime.parse("2024-01-01T00:00:00.123457"))
+ )
+ )
+ val resultOutput = result.queryExecution.executedPlan.outputSet
+ assert(
+ getExecutedPlan(result).exists {
+ case project: ProjectExecTransformer => project.outputSet ==
resultOutput
+ case _ => false
+ },
+ result.queryExecution.executedPlan.treeString
+ )
+ }
+ }
+
+ testGluten("timestamp_ntz literals nested in native expressions ignore the
session timezone") {
+ Seq("UTC", "America/Los_Angeles").foreach {
+ timeZone =>
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "false",
+ SQLConf.SESSION_LOCAL_TIMEZONE.key -> timeZone,
+ GlutenConfig.GLUTEN_ANSI_FALLBACK_ENABLED.key -> "false",
+
"spark.gluten.sql.columnar.backend.velox.enableTimestampNtzValidation" ->
"false"
+ ) {
+ val result = spark.range(2).selectExpr(
+ """hour(element_at(
+ | array(
+ | TIMESTAMP_NTZ '2024-01-01 00:00:00',
+ | TIMESTAMP_NTZ '2024-01-01 01:00:00'),
+ | CAST(id + 1 AS INT))) AS hour
+ |""".stripMargin)
+ checkAnswer(result, Seq(Row(0), Row(1)))
+ val resultOutput = result.queryExecution.executedPlan.outputSet
+ assert(
+ getExecutedPlan(result).exists {
+ case project: ProjectExecTransformer => project.outputSet ==
resultOutput
+ case _ => false
+ },
+ result.queryExecution.executedPlan.treeString
+ )
+ }
+ }
+ }
+}
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
b/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
index f39b80e634..36c9d7c327 100644
---
a/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
+++
b/gluten-ut/spark42/src/test/scala/org/apache/gluten/utils/velox/VeloxTestSettings.scala
@@ -230,6 +230,7 @@ class VeloxTestSettings extends BackendTestSettings {
.exclude("default")
// FIXME(yma11): ObjectType is not covered in RowEncoder/Serializer in
vanilla spark
.exclude("SPARK-37967: Literal.create support ObjectType")
+ enableSuite[GlutenTimestampNtzLiteralSuite]
enableSuite[GlutenMathExpressionsSuite]
// Spark round UT for round(3.1415,3) is not correct.
.exclude("round/bround/floor/ceil")
diff --git
a/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
new file mode 120000
index 0000000000..59e14dd497
--- /dev/null
+++
b/gluten-ut/spark42/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
@@ -0,0 +1 @@
+../../../../../../../../spark41/src/test/scala/org/apache/spark/sql/GlutenTimestampNtzLiteralSuite.scala
\ No newline at end of file
diff --git
a/gluten-ut/test/src/test/scala/org/apache/gluten/extension/columnar/validator/GlutenTimestampNtzValidatorSuite.scala
b/gluten-ut/test/src/test/scala/org/apache/gluten/extension/columnar/validator/GlutenTimestampNtzValidatorSuite.scala
new file mode 100644
index 0000000000..2927a45d31
--- /dev/null
+++
b/gluten-ut/test/src/test/scala/org/apache/gluten/extension/columnar/validator/GlutenTimestampNtzValidatorSuite.scala
@@ -0,0 +1,201 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.extension.columnar.validator
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.ResolvedPartitionSpec
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+import org.apache.spark.sql.catalyst.util.GenericArrayData
+import org.apache.spark.sql.connector.catalog.SupportsPartitionManagement
+import org.apache.spark.sql.execution.{FilterExec, LeafExecNode, ProjectExec,
SparkPlan}
+import org.apache.spark.sql.execution.datasources.v2.DropPartitionExec
+import org.apache.spark.sql.execution.window.WindowExec
+import org.apache.spark.sql.types._
+
+import org.mockito.Mockito.mock
+
+class GlutenTimestampNtzValidatorSuite extends SparkFunSuite {
+ import GlutenTimestampNtzValidatorSuite.InputPlan
+
+ private val id = AttributeReference("id", IntegerType, nullable = false)()
+ private val input = InputPlan(Seq(id))
+
+ private def validator(supportsNtz: Boolean, enableValidation: Boolean):
Validator = {
+ new Validators.FallbackByTimestampNTZ(
+ Some(new GlutenTimestampNtzValidatorSuite.TestConfig(enableValidation)),
+ supportsNtz)
+ }
+
+ private def windowPlan(spec: WindowSpecDefinition): WindowExec = {
+ val sum = Alias(WindowExpression(Sum(id).toAggregateExpression(), spec),
"sum")()
+ WindowExec(Seq(sum), spec.partitionSpec, spec.orderSpec, input)
+ }
+
+ test("timestamp literals inside projections and filters respect backend
support") {
+ for {
+ dataType <- Seq(TimestampNTZType, TimestampType)
+ supportsNtz <- Seq(false, true)
+ enableValidation <- Seq(false, true)
+ } {
+ val timestamps = Literal(
+ new GenericArrayData(Array(1704067200000000L, 1704070800000000L)),
+ ArrayType(dataType, containsNull = false))
+ val hour = Hour(
+ ElementAt(timestamps, Add(id, Literal(1)), failOnError = false),
+ Some("America/Los_Angeles"))
+ val plans = Seq[SparkPlan](
+ ProjectExec(Seq(Alias(hour, "hour")()), input),
+ FilterExec(EqualTo(hour, Literal(0)), input))
+ plans.foreach {
+ plan =>
+ withClue(s"$dataType, supportsNtz=$supportsNtz,
validation=$enableValidation, $plan") {
+ assert(plan.output.forall(_.dataType == IntegerType))
+ assert(plan.children.forall(_.output.forall(_.dataType ==
IntegerType)))
+ val expected = if (dataType == TimestampNTZType && !supportsNtz) {
+ Validator.Failed(s"${plan.nodeName} has TimestampNTZType in
expressions")
+ } else {
+ Validator.Passed
+ }
+ assert(validator(supportsNtz, enableValidation).validate(plan) ==
expected)
+ }
+ }
+ }
+ }
+
+ test("unsupported backends reject NTZ nested in literal types") {
+ val structType = StructType(Seq(StructField("timestamps",
ArrayType(TimestampNTZType))))
+ val dataTypes = Seq(
+ TimestampNTZType,
+ ArrayType(TimestampNTZType),
+ structType,
+ MapType(TimestampNTZType, StringType),
+ MapType(StringType, ArrayType(structType)))
+ for {
+ dataType <- dataTypes
+ literal <- Seq(Literal.default(dataType), Literal.create(null, dataType))
+ enableValidation <- Seq(false, true)
+ } {
+ val plan = ProjectExec(Seq(Alias(IsNull(literal), "missing")()), input)
+ withClue(s"$literal, validation=$enableValidation") {
+ assert(plan.output.map(_.dataType) == Seq(BooleanType))
+ assert(
+ validator(supportsNtz = false, enableValidation =
enableValidation).validate(plan) ==
+ Validator.Failed("Project has TimestampNTZType in expressions"))
+ }
+ }
+ }
+
+ test("window metadata without a data type does not fail NTZ validation") {
+ for {
+ frameType <- Seq(RowFrame, RangeFrame)
+ supportsNtz <- Seq(false, true)
+ enableValidation <- Seq(false, true)
+ } {
+ val spec = WindowSpecDefinition(
+ Seq(id),
+ Seq(SortOrder(id, Ascending)),
+ SpecifiedWindowFrame(frameType, UnboundedPreceding, CurrentRow))
+ withClue(s"$frameType, supportsNtz=$supportsNtz,
validation=$enableValidation") {
+ assert(validator(
+ supportsNtz,
+ enableValidation).validate(windowPlan(spec)) == Validator.Passed)
+ }
+ }
+ }
+
+ test("NTZ expressions in window specifications and frame bounds still
trigger fallback") {
+ for {
+ dataType <- Seq(TimestampNTZType, TimestampType)
+ supportsNtz <- Seq(false, true)
+ enableValidation <- Seq(false, true)
+ } {
+ val hour = Hour(Literal(1704070800000000L, dataType), Some("UTC"))
+ val spec = WindowSpecDefinition(
+ Seq(id),
+ Seq(SortOrder(id, Ascending)),
+ SpecifiedWindowFrame(RowFrame, UnboundedPreceding, CurrentRow))
+ val specs = Seq(
+ spec.copy(partitionSpec = Seq(hour)),
+ spec.copy(orderSpec = Seq(SortOrder(hour, Ascending))),
+ spec.copy(frameSpecification = SpecifiedWindowFrame(RowFrame,
CurrentRow, hour))
+ )
+ specs.foreach {
+ windowSpec =>
+ val plan = windowPlan(windowSpec)
+ withClue(
+ s"$dataType, supportsNtz=$supportsNtz,
validation=$enableValidation, $windowSpec") {
+ assert(plan.output.map(_.dataType) == Seq(IntegerType, LongType))
+ val expected = if (dataType == TimestampNTZType && !supportsNtz) {
+ Validator.Failed("Window has TimestampNTZType in expressions")
+ } else {
+ Validator.Passed
+ }
+ assert(validator(supportsNtz, enableValidation).validate(plan) ==
expected)
+ }
+ }
+ }
+ }
+
+ test("partition command metadata without a data type does not fail NTZ
validation") {
+ val plan = DropPartitionExec(
+ mock(classOf[SupportsPartitionManagement]),
+ Seq(ResolvedPartitionSpec(Seq("id"), InternalRow(1))),
+ ignoreIfNotExists = false,
+ purge = false,
+ refreshCache = () => ()
+ )
+ for {
+ supportsNtz <- Seq(false, true)
+ enableValidation <- Seq(false, true)
+ } {
+ withClue(s"supportsNtz=$supportsNtz, validation=$enableValidation") {
+ assert(plan.output.isEmpty)
+ assert(validator(supportsNtz, enableValidation).validate(plan) ==
Validator.Passed)
+ }
+ }
+ }
+
+ test("existing NTZ input schema validation is unchanged") {
+ val timestamp = AttributeReference("ts", TimestampNTZType)()
+ val plan = ProjectExec(
+ Seq(Alias(Hour(timestamp, Some("UTC")), "hour")()),
+ InputPlan(Seq(timestamp)))
+ for {
+ supportsNtz <- Seq(false, true)
+ enableValidation <- Seq(false, true)
+ } {
+ val expected = if (supportsNtz && !enableValidation) {
+ Validator.Passed
+ } else {
+ Validator.Failed("Project has TimestampNTZType in input/output schema")
+ }
+ assert(validator(supportsNtz, enableValidation).validate(plan) ==
expected)
+ }
+ }
+}
+
+object GlutenTimestampNtzValidatorSuite {
+ private case class InputPlan(override val output: Seq[Attribute]) extends
LeafExecNode {
+ override protected def doExecute(): RDD[InternalRow] =
+ throw new UnsupportedOperationException("This plan is only used for
schema validation.")
+ }
+
+ class TestConfig(val enableTimestampNtzValidation: Boolean)
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
