Skip to content

Commit 115a602

Browse files
committed
Support TIMESTAMP_NTZ type
1 parent 7e4a6bb commit 115a602

10 files changed

Lines changed: 112 additions & 8 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ object VeloxValidatorApi {
106106
private def isPrimitiveType(dataType: DataType): Boolean = {
107107
dataType match {
108108
case BooleanType | ByteType | ShortType | IntegerType | LongType | FloatType | DoubleType |
109-
StringType | BinaryType | _: DecimalType | DateType | TimestampType |
109+
StringType | BinaryType | _: DecimalType | DateType | TimestampType | TimestampNTZType |
110110
YearMonthIntervalType.DEFAULT | NullType =>
111111
true
112112
case _ => false

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 _: TimestampNTZType =>
5354
case _: DateType =>
5455
case _: BinaryType =>
5556
case _: DecimalType =>

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -505,4 +505,12 @@ abstract class DateFunctionsValidateSuite extends FunctionsValidateSuite {
505505
}
506506
}
507507
}
508+
509+
test("timestamp_ntz") {
510+
withSQLConf("spark.sql.ansi.enabled" -> "false") {
511+
runQueryAndCompare("select cast('2022-01-01 00:00:00' as timestamp_ntz)") {
512+
checkGlutenPlan[ProjectExecTransformer]
513+
}
514+
}
515+
}
508516
}

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);
@@ -194,6 +195,7 @@ void VeloxBackend::init(
194195
velox::orc::registerOrcReaderFactory();
195196
velox::exec::ExprToSubfieldFilterParser::registerParser(std::make_unique<SparkExprToSubfieldFilterParser>());
196197
velox::connector::hive::BufferedInputBuilder::registerBuilder(std::make_shared<GlutenBufferedInputBuilder>());
198+
velox::functions::sparksql::registerTimestampNTZType();
197199

198200
// Register Velox functions
199201
registerAllFunctions();
@@ -317,10 +319,10 @@ void VeloxBackend::initConnector(const std::shared_ptr<velox::config::ConfigBase
317319
}
318320
velox::connector::registerConnector(
319321
std::make_shared<velox::connector::hive::HiveConnector>(kHiveConnectorId, hiveConf, ioExecutor_.get()));
320-
322+
321323
// Register value-stream connector for runtime iterator-based inputs
322324
velox::connector::registerConnector(std::make_shared<ValueStreamConnector>(kIteratorConnectorId, hiveConf));
323-
325+
324326
#ifdef GLUTEN_ENABLE_GPU
325327
if (backendConf_->get<bool>(kCudfEnableTableScan, kCudfEnableTableScanDefault) &&
326328
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:
Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.gluten.substrait.expression;
18+
19+
import org.apache.gluten.substrait.type.TimestampTypeNode;
20+
import org.apache.gluten.substrait.type.TypeNode;
21+
22+
import io.substrait.proto.Expression.Literal.Builder;
23+
24+
public class TimestampNTZLiteralNode extends LiteralNodeWithValue<Long> {
25+
public TimestampNTZLiteralNode(Long value) {
26+
super(value, new TimestampTypeNode(true));
27+
}
28+
29+
public TimestampNTZLiteralNode(Long value, TypeNode typeNode) {
30+
super(value, typeNode);
31+
}
32+
33+
@Override
34+
protected void updateLiteralBuilder(Builder literalBuilder, Long value) {
35+
literalBuilder.setTimestamp(value);
36+
}
37+
}
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.gluten.substrait.type;
18+
19+
import io.substrait.proto.Type;
20+
21+
public class TimestampNTZTypeNode extends TypeNode {
22+
23+
public TimestampNTZTypeNode(Boolean nullable) {
24+
super(nullable);
25+
}
26+
27+
@Override
28+
public Type toProtobuf() {
29+
Type.Timestamp.Builder timestampBuilder = Type.Timestamp.newBuilder();
30+
if (nullable) {
31+
timestampBuilder.setNullability(Type.Nullability.NULLABILITY_NULLABLE);
32+
} else {
33+
timestampBuilder.setNullability(Type.Nullability.NULLABILITY_REQUIRED);
34+
}
35+
36+
Type.Builder builder = Type.newBuilder();
37+
builder.setTimestamp(timestampBuilder.build());
38+
return builder.build();
39+
}
40+
}

gluten-substrait/src/main/java/org/apache/gluten/substrait/type/TypeBuilder.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,10 @@ public static TypeNode makeTimestamp(Boolean nullable) {
8181
return new TimestampTypeNode(nullable);
8282
}
8383

84+
public static TypeNode makeTimestampNTZ(Boolean nullable) {
85+
return new TimestampNTZTypeNode(nullable);
86+
}
87+
8488
public static TypeNode makeStruct(Boolean nullable, List<TypeNode> types, List<String> names) {
8589
return new StructNode(nullable, types, names);
8690
}

gluten-substrait/src/main/scala/org/apache/gluten/expression/ConverterUtils.scala

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@ object ConverterUtils extends Logging {
160160
(StringType, isNullable(substraitType.getString.getNullability))
161161
case Type.KindCase.BINARY =>
162162
(BinaryType, isNullable(substraitType.getBinary.getNullability))
163+
case Type.KindCase.TIMESTAMP =>
164+
(TimestampNTZType, isNullable(substraitType.getTimestamp.getNullability))
163165
case Type.KindCase.TIMESTAMP_TZ =>
164166
(TimestampType, isNullable(substraitType.getTimestampTz.getNullability))
165167
case Type.KindCase.DATE =>
@@ -226,6 +228,8 @@ object ConverterUtils extends Logging {
226228
TypeBuilder.makeDecimal(nullable, precision, scale)
227229
case TimestampType =>
228230
TypeBuilder.makeTimestamp(nullable)
231+
case TimestampNTZType =>
232+
TypeBuilder.makeTimestampNTZ(nullable)
229233
case m: MapType =>
230234
TypeBuilder.makeMap(
231235
nullable,
@@ -399,6 +403,7 @@ object ConverterUtils extends Logging {
399403
case DoubleType => "fp64"
400404
case DateType => "date"
401405
case TimestampType => "ts"
406+
case TimestampNTZType => "ntz_ts"
402407
case StringType => "str"
403408
case BinaryType => "vbin"
404409
case DecimalType() =>

0 commit comments

Comments
 (0)