Skip to content

Commit a48be6c

Browse files
dshaaban01dshaaban01
authored
feat: introduce time type (#47)
1 parent b5a9eaa commit a48be6c

File tree

9 files changed

+209
-12
lines changed

9 files changed

+209
-12
lines changed

‎include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,28 @@ def Substrait_StringType : Substrait_Type<"String", "string"> {
6262
}];
6363
}
6464

65+
def Substrait_TimeType : Substrait_Type<"Time", "time"> {
66+
let summary = "Substrait time type";
67+
let description = [{
68+
This type represents a substrait time type.
69+
}];
70+
}
71+
72+
def Substrait_TimeAttr : Substrait_Attr<"Time", "time",
73+
[TypedAttrInterface]> {
74+
let summary = "Substrait time type";
75+
let description = [{
76+
This type represents a substrait time attribute type.
77+
}];
78+
let parameters = (ins "int64_t":$value);
79+
let assemblyFormat = [{ `<` $value `` `us` `>` }];
80+
let extraClassDeclaration = [{
81+
::mlir::Type getType() const {
82+
return TimeType::get(getContext());
83+
}
84+
}];
85+
}
86+
6587
def Substrait_TimestampType : Substrait_Type<"Timestamp", "timestamp"> {
6688
let summary = "Substrait timezone-unaware timestamp type";
6789
let description = [{
@@ -124,6 +146,7 @@ def Substrait_AtomicTypes {
124146
Substrait_TimestampType, // Timestamp
125147
Substrait_TimestampTzType, // TimestampTZ
126148
Substrait_DateType, // Date
149+
Substrait_TimeType, // Time
127150
];
128151
}
129152

@@ -143,6 +166,7 @@ def Substrait_AtomicAttributes {
143166
Substrait_TimestampAttr, // Timestamp
144167
Substrait_TimestampTzAttr, // TimestampTZ
145168
Substrait_DateAttr, // Date
169+
Substrait_TimeAttr, // Time
146170
];
147171
}
148172

‎lib/Target/SubstraitPB/Export.cpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,18 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
280280
return std::move(type);
281281
}
282282

283+
// Handle time.
284+
if (mlir::isa<TimeType>(mlirType)) {
285+
// TODO(ingomueller): support other nullability modes.
286+
auto timeType = std::make_unique<proto::Type::Time>();
287+
timeType->set_nullability(
288+
Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);
289+
290+
auto type = std::make_unique<proto::Type>();
291+
type->set_allocated_time(timeType.release());
292+
return std::move(type);
293+
}
294+
283295
// Handle tuple types.
284296
if (auto tupleType = llvm::dyn_cast<TupleType>(mlirType)) {
285297
auto structType = std::make_unique<proto::Type::Struct>();
@@ -674,6 +686,10 @@ SubstraitExporter::exportOperation(LiteralOp op) {
674686
// `DateType`.
675687
else if (auto dateType = dyn_cast<DateType>(literalType)) {
676688
literal->set_date(mlir::cast<DateAttr>(value).getValue());
689+
}
690+
// `TimeType`.
691+
else if (auto timeType = dyn_cast<TimeType>(literalType)) {
692+
literal->set_time(value.cast<TimeAttr>().getValue());
677693
} else
678694
op->emitOpError("has unsupported value");
679695

‎lib/Target/SubstraitPB/Import.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,8 @@ static mlir::FailureOr<mlir::Type> importType(MLIRContext *context,
127127
return TimestampTzType::get(context);
128128
case proto::Type::kDate:
129129
return DateType::get(context);
130+
case proto::Type::kTime:
131+
return TimeType::get(context);
130132
case proto::Type::kStruct: {
131133
const proto::Type::Struct &structType = type.struct_();
132134
llvm::SmallVector<mlir::Type> fieldTypes;
@@ -367,6 +369,10 @@ importLiteral(ImplicitLocOpBuilder builder,
367369
auto attr = DateAttr::get(context, message.date());
368370
return builder.create<LiteralOp>(attr);
369371
}
372+
case Expression::Literal::LiteralTypeCase::kTime: {
373+
auto attr = TimeAttr::get(context, message.time());
374+
return builder.create<LiteralOp>(attr);
375+
}
370376
// TODO(ingomueller): Support more types.
371377
default: {
372378
const pb::FieldDescriptor *desc =

‎test/Dialect/Substrait/literal.mlir‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,30 @@
11
// RUN: substrait-opt -split-input-file %s \
22
// RUN: | FileCheck %s
33

4+
// CHECK: substrait.plan version 0 : 42 : 1 {
5+
// CHECK-NEXT: relation
6+
// CHECK: %[[V0:.*]] = named_table
7+
// CHECK-NEXT: %[[V1:.*]] = project %[[V0]] : tuple<si1> -> tuple<si1, !substrait.time> {
8+
// CHECK-NEXT: ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
9+
// CHECK-NEXT: %[[V2:.*]] = literal #substrait.time<200000000us> : !substrait.time
10+
// CHECK-NEXT: yield %[[V2]] : !substrait.time
11+
// CHECK-NEXT: }
12+
// CHECK-NEXT: yield %[[V1]] : tuple<si1, !substrait.time>
13+
14+
substrait.plan version 0 : 42 : 1 {
15+
relation {
16+
%0 = named_table @t1 as ["a"] : tuple<si1>
17+
%1 = project %0 : tuple<si1> -> tuple<si1, !substrait.time> {
18+
^bb0(%arg : tuple<si1>):
19+
%time = literal #substrait.time<200000000us> : !substrait.time
20+
yield %time : !substrait.time
21+
}
22+
yield %1 : tuple<si1, !substrait.time>
23+
}
24+
}
25+
26+
// -----
27+
428
// CHECK: substrait.plan version 0 : 42 : 1 {
529
// CHECK-NEXT: relation
630
// CHECK: %[[V0:.*]] = named_table

‎test/Dialect/Substrait/types.mlir‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,20 @@
11
// RUN: substrait-opt -split-input-file %s \
22
// RUN: | FileCheck %s
33

4+
// CHECK-LABEL: substrait.plan
5+
// CHECK: relation
6+
// CHECK: %[[V0:.*]] = named_table @t1 as ["a"] : tuple<!substrait.time>
7+
// CHECK-NEXT: yield %0 : tuple<!substrait.time>
8+
9+
substrait.plan version 0 : 42 : 1 {
10+
relation {
11+
%0 = named_table @t1 as ["a"] : tuple<!substrait.time>
12+
yield %0 : tuple<!substrait.time>
13+
}
14+
}
15+
16+
// -----
17+
418
// CHECK-LABEL: substrait.plan
519
// CHECK: relation
620
// CHECK: %[[V0:.*]] = named_table @t1 as ["a"] : tuple<!substrait.date>

‎test/Target/SubstraitPB/Export/literal.mlir‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,35 @@
99
// RUN: --split-input-file --output-split-marker="# -----" \
1010
// RUN: | FileCheck %s
1111

12+
// CHECK-LABEL: relations {
13+
// CHECK-NEXT: rel {
14+
// CHECK-NEXT: project {
15+
// CHECK-NEXT: common {
16+
// CHECK-NEXT: direct {
17+
// CHECK-NEXT: }
18+
// CHECK-NEXT: }
19+
// CHECK-NEXT: input {
20+
// CHECK-NEXT: read {
21+
// CHECK: expressions {
22+
// CHECK-NEXT: literal {
23+
// CHECK-NEXT: time: 200000000
24+
// CHECK-NEXT: }
25+
// CHECK-NEXT: }
26+
27+
substrait.plan version 0 : 42 : 1 {
28+
relation {
29+
%0 = named_table @t1 as ["a"] : tuple<si1>
30+
%1 = project %0 : tuple<si1> -> tuple<si1, !substrait.time> {
31+
^bb0(%arg : tuple<si1>):
32+
%time = literal #substrait.time<200000000us> : !substrait.time
33+
yield %time : !substrait.time
34+
}
35+
yield %1 : tuple<si1, !substrait.time>
36+
}
37+
}
38+
39+
// -----
40+
1241
// CHECK-LABEL: relations {
1342
// CHECK-NEXT: rel {
1443
// CHECK-NEXT: project {

‎test/Target/SubstraitPB/Export/types.mlir‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
// CHECK-NEXT: names: "a"
1717
// CHECK-NEXT: struct {
1818
// CHECK-NEXT: types {
19-
// CHECK-NEXT: date {
19+
// CHECK-NEXT: time {
2020
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
2121
// CHECK-NEXT: }
2222
// CHECK-NEXT: }
@@ -27,8 +27,8 @@
2727

2828
substrait.plan version 0 : 42 : 1 {
2929
relation {
30-
%0 = named_table @t1 as ["a"] : tuple<!substrait.date>
31-
yield %0 : tuple<!substrait.date>
30+
%0 = named_table @t1 as ["a"] : tuple<!substrait.time>
31+
yield %0 : tuple<!substrait.time>
3232
}
3333
}
3434

@@ -39,15 +39,9 @@ substrait.plan version 0 : 42 : 1 {
3939
// CHECK-NEXT: read {
4040
// CHECK: base_schema {
4141
// CHECK-NEXT: names: "a"
42-
// CHECK-NEXT: names: "b"
4342
// CHECK-NEXT: struct {
4443
// CHECK-NEXT: types {
45-
// CHECK-NEXT: timestamp {
46-
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
47-
// CHECK-NEXT: }
48-
// CHECK-NEXT: }
49-
// CHECK-NEXT: types {
50-
// CHECK-NEXT: timestamp_tz {
44+
// CHECK-NEXT: date {
5145
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
5246
// CHECK-NEXT: }
5347
// CHECK-NEXT: }
@@ -58,8 +52,8 @@ substrait.plan version 0 : 42 : 1 {
5852

5953
substrait.plan version 0 : 42 : 1 {
6054
relation {
61-
%0 = named_table @t1 as ["a", "b"] : tuple<!substrait.timestamp, !substrait.timestamp_tz>
62-
yield %0 : tuple<!substrait.timestamp, !substrait.timestamp_tz>
55+
%0 = named_table @t1 as ["a"] : tuple<!substrait.date>
56+
yield %0 : tuple<!substrait.date>
6357
}
6458
}
6559

‎test/Target/SubstraitPB/Import/literal.textpb‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,60 @@
1010
# RUN: --split-input-file="# ""-----" --output-split-marker="// -----" \
1111
# RUN: | FileCheck %s
1212

13+
# CHECK: substrait.plan version 0 : 42 : 1 {
14+
# CHECK-NEXT: relation
15+
# CHECK: %[[V0:.*]] = named_table
16+
# CHECK-NEXT: %[[V1:.*]] = project %[[V0]] : tuple<si1> -> tuple<si1, !substrait.time> {
17+
# CHECK-NEXT: ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
18+
# CHECK-NEXT: %[[V2:.*]] = literal #substrait.time<200000000us> : !substrait.time
19+
# CHECK-NEXT: yield %[[V2]] : !substrait.time
20+
# CHECK-NEXT: }
21+
# CHECK-NEXT: yield %[[V1]] : tuple<si1, !substrait.time>
22+
23+
relations {
24+
rel {
25+
project {
26+
common {
27+
direct {
28+
}
29+
}
30+
input {
31+
read {
32+
common {
33+
direct {
34+
}
35+
}
36+
base_schema {
37+
names: "a"
38+
struct {
39+
types {
40+
bool {
41+
nullability: NULLABILITY_REQUIRED
42+
}
43+
}
44+
nullability: NULLABILITY_REQUIRED
45+
}
46+
}
47+
named_table {
48+
names: "t1"
49+
}
50+
}
51+
}
52+
expressions {
53+
literal {
54+
time: 200000000
55+
}
56+
}
57+
}
58+
}
59+
}
60+
version {
61+
minor_number: 42
62+
patch_number: 1
63+
}
64+
65+
# -----
66+
1367
# CHECK: substrait.plan version 0 : 42 : 1 {
1468
# CHECK-NEXT: relation
1569
# CHECK: %[[V0:.*]] = named_table

‎test/Target/SubstraitPB/Import/types.textpb‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,42 @@
1010
# RUN: --split-input-file="# ""-----" --output-split-marker="// -----" \
1111
# RUN: | FileCheck %s
1212

13+
# CHECK: substrait.plan
14+
# CHECK-NEXT: relation
15+
# CHECK-NEXT: named_table
16+
# CHECK-SAME: : tuple<!substrait.time>
17+
18+
relations {
19+
rel {
20+
read {
21+
common {
22+
direct {
23+
}
24+
}
25+
base_schema {
26+
names: "a"
27+
struct {
28+
types {
29+
time {
30+
nullability: NULLABILITY_REQUIRED
31+
}
32+
}
33+
nullability: NULLABILITY_REQUIRED
34+
}
35+
}
36+
named_table {
37+
names: "t1"
38+
}
39+
}
40+
}
41+
}
42+
version {
43+
minor_number: 42
44+
patch_number: 1
45+
}
46+
47+
# -----
48+
1349
# CHECK: substrait.plan
1450
# CHECK-NEXT: relation
1551
# CHECK-NEXT: named_table

0 commit comments

Comments
 (0)