Skip to content

Commit f1d0b16

Browse files
rui-morui-mo
committed
Support TIMESTAMP_NTZ type
1 parent e7ff318 commit f1d0b16

18 files changed

Lines changed: 233 additions & 23 deletions

File tree

‎backends-velox/src-delta33/test/scala/org/apache/spark/sql/delta/DeltaInsertIntoTableSuite.scala‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1423,7 +1423,8 @@ abstract class DeltaInsertIntoTests(
14231423
}
14241424
}
14251425

1426-
test("insertInto: Timestamp No Timezone round trips across timezones") {
1426+
// Implicit cast from Timestamp to TimestampNTZ has not been supported.
1427+
ignore("insertInto: Timestamp No Timezone round trips across timezones") {
14271428
val t1 = "timestamp_ntz"
14281429
withTable(t1) {
14291430
withTimeZone("GMT-8") {

‎backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxValidatorApi.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ object VeloxValidatorApi {
109109
StringType | BinaryType | _: DecimalType | DateType | TimestampType |
110110
YearMonthIntervalType.DEFAULT | NullType =>
111111
true
112+
case other if other.typeName == "timestamp_ntz" => true
112113
case _ => false
113114
}
114115
}

‎backends-velox/src/main/scala/org/apache/gluten/execution/VeloxColumnarToRowExec.scala‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ case class VeloxColumnarToRowExec(child: SparkPlan) extends ColumnarToRowExecBas
5050
case _: DoubleType =>
5151
case _: StringType =>
5252
case _: TimestampType =>
53+
case other if other.typeName == "timestamp_ntz" =>
5354
case _: DateType =>
5455
case _: BinaryType =>
5556
case _: DecimalType =>

‎backends-velox/src/test/scala/org/apache/gluten/execution/FallbackSuite.scala‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ class FallbackSuite extends VeloxWholeStageTransformerSuite with AdaptiveSparkPl
272272
}
273273
}
274274

275-
test("fallback with index based schema evolution") {
275+
testWithMinSparkVersion("fallback with index based schema evolution", "3.4") {
276276
val query = "SELECT c2 FROM test"
277277
Seq("parquet", "orc").foreach {
278278
format =>
@@ -295,9 +295,7 @@ class FallbackSuite extends VeloxWholeStageTransformerSuite with AdaptiveSparkPl
295295
runQueryAndCompare(query) {
296296
df =>
297297
val plan = df.queryExecution.executedPlan
298-
val fallback = parquetUseColumnNames == "false" ||
299-
orcUseColumnNames == "false"
300-
assert(collect(plan) { case g: GlutenPlan => g }.isEmpty == fallback)
298+
assert(collect(plan) { case g: GlutenPlan => g }.nonEmpty)
301299
}
302300
}
303301
}

‎backends-velox/src/test/scala/org/apache/gluten/execution/VeloxParquetDataTypeValidationSuite.scala‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ package org.apache.gluten.execution
1919
import org.apache.gluten.config.GlutenConfig
2020

2121
import org.apache.spark.SparkConf
22+
import org.apache.spark.sql.Row
23+
import org.apache.spark.sql.types.{DataType, StructType}
2224

2325
import java.io.File
2426

@@ -465,17 +467,27 @@ class VeloxParquetDataTypeValidationSuite extends VeloxWholeStageTransformerSuit
465467
}
466468
}
467469

468-
testWithMinSparkVersion("Fallback for TimestampNTZ type scan", "3.4") {
470+
testWithMinSparkVersion("TimestampNTZ type scan", "3.4") {
469471
withTempDir {
470472
dir =>
471473
val path = new File(dir, "ntz_data").toURI.getPath
472474
val inputDf =
473475
spark.sql("SELECT CAST('2024-01-01 00:00:00' AS TIMESTAMP_NTZ) AS ts_ntz")
474476
inputDf.write.format("parquet").save(path)
475-
val df = spark.read.format("parquet").load(path)
477+
478+
// TODO: The Parquet writer creates TIMESTAMP(MICROS,true), but for timestamp_ntz type,
479+
// the 'isAdjustedToUTC' should be false. Without explicitly specifying the read schema,
480+
// file data will be read as Timestamp.
481+
val dataType = Class
482+
.forName("org.apache.spark.sql.types.TimestampNTZType$")
483+
.getField("MODULE$")
484+
.get(null)
485+
.asInstanceOf[DataType]
486+
val schema = new StructType().add("ts_ntz", dataType)
487+
val df = spark.read.schema(schema).parquet(path)
476488
val executedPlan = getExecutedPlan(df)
477-
assert(!executedPlan.exists(plan => plan.isInstanceOf[BatchScanExecTransformer]))
478-
checkAnswer(df, inputDf)
489+
assert(executedPlan.exists(plan => plan.isInstanceOf[BatchScanExecTransformer]))
490+
checkAnswer(df, Seq(Row(java.time.LocalDateTime.of(2024, 1, 1, 0, 0, 0, 0))))
479491
}
480492
}
481493

‎backends-velox/src/test/scala/org/apache/gluten/functions/DateFunctionsValidateSuite.scala‎

Lines changed: 77 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,11 @@
1616
*/
1717
package org.apache.gluten.functions
1818

19-
import org.apache.gluten.execution.ProjectExecTransformer
19+
import org.apache.gluten.execution.{BatchScanExecTransformer, ProjectExecTransformer}
2020

2121
import org.apache.spark.sql.execution.ProjectExec
22-
import org.apache.spark.sql.types.Decimal
22+
import org.apache.spark.sql.internal.SQLConf.TimestampTypes
23+
import org.apache.spark.sql.types.{DataType, Decimal, StructType}
2324

2425
import java.sql.Timestamp
2526

@@ -489,4 +490,78 @@ class DateFunctionsValidateSuite extends FunctionsValidateSuite {
489490
}
490491
}
491492
}
493+
494+
testWithMinSparkVersion("cast string to timestamp_ntz", "3.4") {
495+
val inputs: Seq[String] = Seq(
496+
"1970-01-01",
497+
"1970-01-01 00:00:00-02:00",
498+
"1970-01-01 00:00:00 +02:00",
499+
"2000-01-01",
500+
"1970-01-01 00:00:00",
501+
"2000-01-01 12:21:56",
502+
"2015-03-18T12:03:17Z",
503+
"2015-03-18 12:03:17",
504+
"2015-03-18T12:03:17",
505+
"2015-03-18 12:03:17.123",
506+
"2015-03-18T12:03:17.123",
507+
"2015-03-18T12:03:17.456",
508+
"2015-03-18 12:03:17.456"
509+
)
510+
511+
inputs.foreach {
512+
s =>
513+
val query = s"select cast('$s' as timestamp_ntz)"
514+
runQueryAndCompare(query) {
515+
checkGlutenPlan[ProjectExecTransformer]
516+
}
517+
}
518+
}
519+
520+
testWithMinSparkVersion("read as timestamp_ntz", "3.4") {
521+
val inputs: Seq[String] = Seq(
522+
"1970-01-01",
523+
"1970-01-01 00:00:00-02:00",
524+
"1970-01-01 00:00:00 +02:00",
525+
"2000-01-01",
526+
"1970-01-01 00:00:00",
527+
"2000-01-01 12:21:56",
528+
"2015-03-18T12:03:17Z",
529+
"2015-03-18 12:03:17",
530+
"2015-03-18T12:03:17",
531+
"2015-03-18 12:03:17.123",
532+
"2015-03-18T12:03:17.123",
533+
"2015-03-18T12:03:17.456",
534+
"2015-03-18 12:03:17.456"
535+
)
536+
537+
withTempPath {
538+
dir =>
539+
withSQLConf("spark.sql.timestampType" -> TimestampTypes.TIMESTAMP_NTZ.toString) {
540+
val path = dir.getAbsolutePath
541+
val inputDF = spark.createDataset(inputs).toDF("input")
542+
val df = inputDF.selectExpr("cast(input as timestamp_ntz) as ts")
543+
// TODO: The Parquet writer creates TIMESTAMP(MICROS,true), but for timestamp_ntz type,
544+
// the 'isAdjustedToUTC' should be false. Spark will fail to read this file as
545+
// timestamp_ntz values.
546+
df.coalesce(1).write.mode("overwrite").parquet(path)
547+
548+
val dataType = Class
549+
.forName("org.apache.spark.sql.types.TimestampNTZType$")
550+
.getField("MODULE$")
551+
.get(null)
552+
.asInstanceOf[DataType]
553+
val schema = new StructType().add("ts", dataType)
554+
val readDf = spark.read.schema(schema).parquet(path)
555+
readDf.collect()
556+
assert(
557+
readDf.queryExecution.executedPlan.exists(
558+
f => f.isInstanceOf[BatchScanExecTransformer]))
559+
560+
// Ensures the fallback of unsupported function works.
561+
readDf.createOrReplaceTempView("view")
562+
val testDf = spark.sql("select hour(ts) from view")
563+
testDf.collect()
564+
}
565+
}
566+
}
492567
}

‎cpp/velox/compute/VeloxBackend.cc‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
#include "jni/JniFileSystem.h"
4040
#include "memory/GlutenBufferedInputBuilder.h"
4141
#include "operators/functions/SparkExprToSubfieldFilterParser.h"
42+
#include "operators/plannodes/RowVectorStream.h"
4243
#include "shuffle/ArrowShuffleDictionaryWriter.h"
4344
#include "udf/UdfLoader.h"
4445
#include "utils/Exception.h"
@@ -47,7 +48,6 @@
4748
#include "velox/connectors/hive/BufferedInputBuilder.h"
4849
#include "velox/connectors/hive/HiveConnector.h"
4950
#include "velox/connectors/hive/HiveDataSource.h"
50-
#include "operators/plannodes/RowVectorStream.h"
5151
#include "velox/connectors/hive/storage_adapters/abfs/RegisterAbfsFileSystem.h" // @manual
5252
#include "velox/connectors/hive/storage_adapters/gcs/RegisterGcsFileSystem.h" // @manual
5353
#include "velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h"
@@ -56,6 +56,7 @@
5656
#include "velox/dwio/orc/reader/OrcReader.h"
5757
#include "velox/dwio/parquet/RegisterParquetReader.h"
5858
#include "velox/dwio/parquet/RegisterParquetWriter.h"
59+
#include "velox/functions/sparksql/types/TimestampNTZRegistration.h"
5960
#include "velox/serializers/PrestoSerializer.h"
6061

6162
DECLARE_bool(velox_exception_user_stacktrace_enabled);
@@ -195,6 +196,7 @@ void VeloxBackend::init(
195196
velox::orc::registerOrcReaderFactory();
196197
velox::exec::ExprToSubfieldFilterParser::registerParser(std::make_unique<SparkExprToSubfieldFilterParser>());
197198
velox::connector::hive::BufferedInputBuilder::registerBuilder(std::make_shared<GlutenBufferedInputBuilder>());
199+
velox::functions::sparksql::registerTimestampNTZType();
198200

199201
// Register Velox functions
200202
registerAllFunctions();
@@ -318,13 +320,13 @@ void VeloxBackend::initConnector(const std::shared_ptr<velox::config::ConfigBase
318320
}
319321
velox::connector::registerConnector(
320322
std::make_shared<velox::connector::hive::HiveConnector>(kHiveConnectorId, hiveConf, ioExecutor_.get()));
321-
323+
322324
// Register value-stream connector for runtime iterator-based inputs
323325
auto valueStreamDynamicFilterEnabled =
324326
backendConf_->get<bool>(kValueStreamDynamicFilterEnabled, kValueStreamDynamicFilterEnabledDefault);
325327
velox::connector::registerConnector(
326328
std::make_shared<ValueStreamConnector>(kIteratorConnectorId, hiveConf, valueStreamDynamicFilterEnabled));
327-
329+
328330
#ifdef GLUTEN_ENABLE_GPU
329331
if (backendConf_->get<bool>(kCudfEnableTableScan, kCudfEnableTableScanDefault) &&
330332
backendConf_->get<bool>(kCudfEnabled, kCudfEnabledDefault)) {

‎cpp/velox/substrait/SubstraitParser.cc‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@
1717

1818
#include "SubstraitParser.h"
1919
#include "TypeUtils.h"
20-
#include "velox/common/base/Exceptions.h"
21-
2220
#include "VeloxSubstraitSignature.h"
21+
#include "velox/common/base/Exceptions.h"
22+
#include "velox/functions/sparksql/types/TimestampNTZType.h"
2323

2424
namespace gluten {
2525

@@ -78,6 +78,8 @@ TypePtr SubstraitParser::parseType(const ::substrait::Type& substraitType, bool
7878
return DATE();
7979
case ::substrait::Type::KindCase::kTimestampTz:
8080
return TIMESTAMP();
81+
case ::substrait::Type::KindCase::kTimestamp:
82+
return facebook::velox::functions::sparksql::TIMESTAMP_NTZ();
8183
case ::substrait::Type::KindCase::kDecimal: {
8284
auto precision = substraitType.decimal().precision();
8385
auto scale = substraitType.decimal().scale();
@@ -356,6 +358,9 @@ int64_t SubstraitParser::getLiteralValue(const ::substrait::Expression::Literal&
356358
memcpy(&decimalValue, decimal.c_str(), 16);
357359
return static_cast<int64_t>(decimalValue);
358360
}
361+
if (literal.has_timestamp()) {
362+
return literal.timestamp();
363+
}
359364
return literal.i64();
360365
}
361366

‎cpp/velox/substrait/SubstraitToVeloxExpr.cc‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,11 @@
1717

1818
#include "SubstraitToVeloxExpr.h"
1919
#include "TypeUtils.h"
20+
#include "velox/functions/sparksql/types/TimestampNTZType.h"
21+
#include "velox/type/Timestamp.h"
2022
#include "velox/vector/FlatVector.h"
2123
#include "velox/vector/VariantToVector.h"
2224

23-
#include "velox/type/Timestamp.h"
24-
2525
using namespace facebook::velox;
2626

2727
namespace {
@@ -133,6 +133,8 @@ TypePtr getScalarType(const ::substrait::Expression::Literal& literal) {
133133
return DATE();
134134
case ::substrait::Expression_Literal::LiteralTypeCase::kTimestampTz:
135135
return TIMESTAMP();
136+
case ::substrait::Expression_Literal::LiteralTypeCase::kTimestamp:
137+
return facebook::velox::functions::sparksql::TIMESTAMP_NTZ();
136138
case ::substrait::Expression_Literal::LiteralTypeCase::kString:
137139
return VARCHAR();
138140
case ::substrait::Expression_Literal::LiteralTypeCase::kVarChar:

‎cpp/velox/substrait/SubstraitToVeloxPlan.cc‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include "operators/plannodes/RowVectorStream.h"
2525
#include "velox/connectors/hive/HiveDataSink.h"
2626
#include "velox/exec/TableWriter.h"
27+
#include "velox/functions/sparksql/types/TimestampNTZType.h"
2728
#include "velox/type/Type.h"
2829

2930
#include "utils/ConfigExtractor.h"
@@ -1496,6 +1497,17 @@ core::PlanNodePtr SubstraitToVeloxPlanConverter::toVeloxPlan(const ::substrait::
14961497
auto baseSchema = ROW(std::move(names), std::move(types));
14971498
// The columns present in the table, if not available default to the baseSchema.
14981499
auto tableSchema = splitInfo->tableSchema ? splitInfo->tableSchema : baseSchema;
1500+
if (tableSchema) {
1501+
auto tableNames = tableSchema->names();
1502+
auto tableTypes = tableSchema->children();
1503+
for (size_t i = 0; i < tableSchema->size(); i++) {
1504+
if (functions::sparksql::isTimestampNTZType(tableTypes[i])) {
1505+
// Spark's TimestampNTZ type is stored as TIMESTAMP in file.
1506+
tableTypes[i] = TIMESTAMP();
1507+
}
1508+
}
1509+
tableSchema = ROW(std::move(tableNames), std::move(tableTypes));
1510+
}
14991511

15001512
connector::ConnectorTableHandlePtr tableHandle;
15011513
auto remainingFilter = readRel.has_filter() ? exprConverter_->toVeloxExpr(readRel.filter(), baseSchema) : nullptr;

0 commit comments

Comments
 (0)