Skip to content

Commit

Permalink
introduce interval types
Browse files Browse the repository at this point in the history
  • Loading branch information
@dshaaban01
dshaaban01 committed Feb 6, 2025
1 parent 131caf5 commit 91b1ab0
Show file tree
Hide file tree
Showing 10 changed files with 320 additions and 2 deletions.
24 changes: 24 additions & 0 deletions include/substrait-mlir/Dialect/Substrait/IR/SubstraitAttrs.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,28 @@ def Substrait_TimestampTzAttr
let assemblyFormat = [{ `<` $value `` `us` `>` }];
}

def Substrait_IntervalYearMonthAttr
: Substrait_StaticallyTypedAttr<"IntervalYearMonth", "interval_year_month",
"IntervalYearMonthType"> {
let summary = "Substrait interval year to month type";
let description = [{
This type represents a substrait interval year to month attribute type.
}];
let parameters = (ins "int32_t":$years_value, "int32_t":$months_value);
let assemblyFormat = [{ `<` $years_value `` `y` $months_value `` `m` `>` }];
}

def Substrait_IntervalDaySecondAttr
: Substrait_StaticallyTypedAttr<"IntervalDaySecond", "interval_day_second",
"IntervalDaySecondType"> {
let summary = "Substrait interval day to second type";
let description = [{
This type represents a substrait interval day to second attribute type.
}];
let parameters = (ins "int32_t":$days_value, "int32_t":$seconds_value);
let assemblyFormat = [{ `<` $days_value `` `d` $seconds_value `` `s` `>` }];
}

def Substrait_VersionAttr : Substrait_Attr<"Version", "version"> {
let summary = "Substrait version";
let description = [{
Expand Down Expand Up @@ -153,6 +175,8 @@ def Substrait_AtomicAttributes {
Substrait_TimestampTzAttr, // TimestampTZ
Substrait_DateAttr, // Date
Substrait_TimeAttr, // Time
Substrait_IntervalYearMonthAttr, // IntervalYear
Substrait_IntervalDaySecondAttr, // IntervalDay
];
}

Expand Down
19 changes: 17 additions & 2 deletions include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,20 @@ def Substrait_DateType : Substrait_Type<"Date", "date"> {
}];
}

def Substrait_IntervalYearMonthType : Substrait_Type<"IntervalYearMonth", "interval_year_month"> {
let summary = "Substrait interval year to month type";
let description = [{
This type represents a substrait interval year to month type.
}];
}

def Substrait_IntervalDaySecondType : Substrait_Type<"IntervalDaySecond", "interval_day_second"> {
let summary = "Substrait interval day to second type";
let description = [{
This type represents a substrait interval day to second type.
}];
}

def Substrait_StringType : Substrait_Type<"String", "string"> {
let summary = "Substrait string type";
let description = [{
Expand Down Expand Up @@ -61,8 +75,7 @@ def Substrait_TimestampTzType : Substrait_Type<"TimestampTz", "timestamp_tz"> {
}];
}

/// Currently supported atomic types, listed in order of substrait specification.
/// These correspond directly to the types in
/// Currently supported atomic types. These correspond directly to the types in
/// https://github.com/substrait-io/substrait/blob/main/proto/substrait/type.proto.
def Substrait_AtomicTypes {
list<Type> types = [
Expand All @@ -79,6 +92,8 @@ def Substrait_AtomicTypes {
Substrait_TimestampTzType, // TimestampTZ
Substrait_DateType, // Date
Substrait_TimeType, // Time
Substrait_IntervalYearMonthType, // IntervalYear
Substrait_IntervalDaySecondType, // IntervalDay
];
}

Expand Down
46 changes: 46 additions & 0 deletions lib/Target/SubstraitPB/Export.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,29 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
return std::move(type);
}

// Handle interval_year.
if (mlir::isa<IntervalYearMonthType>(mlirType)) {
// TODO(ingomueller): support other nullability modes.
auto intervalYearType = std::make_unique<proto::Type::IntervalYear>();
intervalYearType->set_nullability(
Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);

auto type = std::make_unique<proto::Type>();
type->set_allocated_interval_year(intervalYearType.release());
return std::move(type);
}
// Handle interval_day.
if (mlir::isa<IntervalDaySecondType>(mlirType)) {
// TODO(ingomueller): support other nullability modes.
auto intervalDayType = std::make_unique<proto::Type::IntervalDay>();
intervalDayType->set_nullability(
Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);

auto type = std::make_unique<proto::Type>();
type->set_allocated_interval_day(intervalDayType.release());
return std::move(type);
}

// Handle tuple types.
if (auto tupleType = llvm::dyn_cast<TupleType>(mlirType)) {
auto structType = std::make_unique<proto::Type::Struct>();
Expand Down Expand Up @@ -833,6 +856,29 @@ SubstraitExporter::exportOperation(LiteralOp op) {
// `TimeType`.
else if (auto timeType = dyn_cast<TimeType>(literalType)) {
literal->set_time(mlir::cast<TimeAttr>(value).getValue());
}
// `IntervalType`'s.
else if (auto timeType = dyn_cast<IntervalYearMonthType>(literalType)) {
auto intervalYearToMonth = std::make_unique<
::substrait::proto::Expression_Literal_IntervalYearToMonth>();
auto intervalYear =
mlir::cast<IntervalYearMonthAttr>(value).getYearsValue();
auto intervalMonth =
mlir::cast<IntervalYearMonthAttr>(value).getMonthsValue();
intervalYearToMonth->set_years(intervalYear);
intervalYearToMonth->set_months(intervalMonth);
literal->set_allocated_interval_year_to_month(
intervalYearToMonth.release());
} else if (auto timeType = dyn_cast<IntervalDaySecondType>(literalType)) {
auto intervalDaytoSecond = std::make_unique<
::substrait::proto::Expression_Literal_IntervalDayToSecond>();
auto intervalDay = mlir::cast<IntervalDaySecondAttr>(value).getDaysValue();
auto intervalSecond =
mlir::cast<IntervalDaySecondAttr>(value).getSecondsValue();
intervalDaytoSecond->set_days(intervalDay);
intervalDaytoSecond->set_seconds(intervalSecond);
literal->set_allocated_interval_day_to_second(
intervalDaytoSecond.release());
} else
op->emitOpError("has unsupported value");

Expand Down
16 changes: 16 additions & 0 deletions lib/Target/SubstraitPB/Import.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,10 @@ static mlir::FailureOr<mlir::Type> importType(MLIRContext *context,
return DateType::get(context);
case proto::Type::kTime:
return TimeType::get(context);
case proto::Type::kIntervalYear:
return IntervalYearMonthType::get(context);
case proto::Type::kIntervalDay:
return IntervalDaySecondType::get(context);
case proto::Type::kStruct: {
const proto::Type::Struct &structType = type.struct_();
llvm::SmallVector<mlir::Type> fieldTypes;
Expand Down Expand Up @@ -592,6 +596,18 @@ importLiteral(ImplicitLocOpBuilder builder,
auto attr = TimeAttr::get(context, message.time());
return builder.create<LiteralOp>(attr);
}
case Expression::Literal::LiteralTypeCase::kIntervalYearToMonth: {
auto attr = IntervalYearMonthAttr::get(
context, message.interval_year_to_month().years(),
message.interval_year_to_month().months());
return builder.create<LiteralOp>(attr);
}
case Expression::Literal::LiteralTypeCase::kIntervalDayToSecond: {
auto attr = IntervalDaySecondAttr::get(
context, message.interval_day_to_second().days(),
message.interval_day_to_second().seconds());
return builder.create<LiteralOp>(attr);
}
// TODO(ingomueller): Support more types.
default: {
const pb::FieldDescriptor *desc =
Expand Down
26 changes: 26 additions & 0 deletions test/Dialect/Substrait/literal.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,32 @@
// RUN: substrait-opt -split-input-file %s \
// RUN: | FileCheck %s

// CHECK: substrait.plan version 0 : 42 : 1 {
// CHECK-NEXT: relation
// CHECK: %[[V0:.*]] = named_table
// CHECK-NEXT: %[[V1:.*]] = project %[[V0]] : tuple<si1> -> tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second> {
// CHECK-NEXT: ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
// CHECK-NEXT: %[[V2:.*]] = literal #substrait.interval_year_month<2024y 1m>{{$}}
// CHECK-NEXT: %[[V3:.*]] = literal #substrait.interval_day_second<9d 8000s>{{$}}
// CHECK-NEXT: yield %[[V2]], %[[V3]] : !substrait.interval_year_month, !substrait.interval_day_second
// CHECK-NEXT: }
// CHECK-NEXT: yield %[[V1]] : tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second>

substrait.plan version 0 : 42 : 1 {
relation {
%0 = named_table @t1 as ["a"] : tuple<si1>
%1 = project %0 : tuple<si1> -> tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second> {
^bb0(%arg : tuple<si1>):
%interval_year_month = literal #substrait.interval_year_month<2024y 1m>
%interval_day_second = literal #substrait.interval_day_second<9d 8000s>
yield %interval_year_month, %interval_day_second : !substrait.interval_year_month, !substrait.interval_day_second
}
yield %1 : tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second>
}
}

// -----

// CHECK: substrait.plan version 0 : 42 : 1 {
// CHECK-NEXT: relation
// CHECK: %[[V0:.*]] = named_table
Expand Down
14 changes: 14 additions & 0 deletions test/Dialect/Substrait/types.mlir
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,20 @@
// RUN: substrait-opt -split-input-file %s \
// RUN: | FileCheck %s

// CHECK-LABEL: substrait.plan
// CHECK: relation
// CHECK: %[[V0:.*]] = named_table @t1 as ["a", "b"] : tuple<!substrait.interval_year_month, !substrait.interval_day_second>
// CHECK-NEXT: yield %0 : tuple<!substrait.interval_year_month, !substrait.interval_day_second>

substrait.plan version 0 : 42 : 1 {
relation {
%0 = named_table @t1 as ["a", "b"] : tuple<!substrait.interval_year_month, !substrait.interval_day_second>
yield %0 : tuple<!substrait.interval_year_month, !substrait.interval_day_second>
}
}

// -----

// CHECK-LABEL: substrait.plan
// CHECK: relation
// CHECK: %[[V0:.*]] = named_table @t1 as ["a"] : tuple<!substrait.time>
Expand Down
38 changes: 38 additions & 0 deletions test/Target/SubstraitPB/Export/literal.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,44 @@
// RUN: --split-input-file --output-split-marker="# -----" \
// RUN: | FileCheck %s

// CHECK-LABEL: relations {
// CHECK-NEXT: rel {
// CHECK-NEXT: project {
// CHECK-NEXT: common {
// CHECK-NEXT: direct {
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: input {
// CHECK-NEXT: read {
// CHECK: expressions {
// CHECK-NEXT: literal {
// CHECK-NEXT: interval_year_to_month {
// CHECK-NEXT: years: 2024
// CHECK-NEXT: months: 1
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: expressions {
// CHECK-NEXT: literal {
// CHECK-NEXT: interval_day_to_second {
// CHECK-NEXT: days: 9
// CHECK-NEXT: seconds: 8000

substrait.plan version 0 : 42 : 1 {
relation {
%0 = named_table @t1 as ["a"] : tuple<si1>
%1 = project %0 : tuple<si1> -> tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second> {
^bb0(%arg : tuple<si1>):
%interval_year_month = literal #substrait.interval_year_month<2024y 1m>
%interval_day_second = literal #substrait.interval_day_second<9d 8000s>
yield %interval_year_month, %interval_day_second : !substrait.interval_year_month, !substrait.interval_day_second
}
yield %1 : tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second>
}
}

// -----

// CHECK-LABEL: relations {
// CHECK-NEXT: rel {
// CHECK-NEXT: project {
Expand Down
31 changes: 31 additions & 0 deletions test/Target/SubstraitPB/Export/types.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,37 @@
// RUN: --split-input-file --output-split-marker="# -----" \
// RUN: | FileCheck %s

// CHECK-LABEL: relations {
// CHECK-NEXT: rel {
// CHECK-NEXT: read {
// CHECK: base_schema {
// CHECK-NEXT: names: "a"
// CHECK-NEXT: names: "b"
// CHECK-NEXT: struct {
// CHECK-NEXT: types {
// CHECK-NEXT: interval_year {
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: types {
// CHECK-NEXT: interval_day {
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: nullability: NULLABILITY_REQUIRED
// CHECK-NEXT: }
// CHECK-NEXT: }
// CHECK-NEXT: named_table {

substrait.plan version 0 : 42 : 1 {
relation {
%0 = named_table @t1 as ["a", "b"] : tuple<!substrait.interval_year_month, !substrait.interval_day_second>
yield %0 : tuple<!substrait.interval_year_month, !substrait.interval_day_second>
}
}

// -----

// CHECK-LABEL: relations {
// CHECK-NEXT: rel {
// CHECK-NEXT: read {
Expand Down
66 changes: 66 additions & 0 deletions test/Target/SubstraitPB/Import/literal.textpb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,72 @@
# RUN: --split-input-file="# ""-----" --output-split-marker="// -----" \
# RUN: | FileCheck %s

# CHECK: substrait.plan version 0 : 42 : 1 {
# CHECK-NEXT: relation
# CHECK: %[[V0:.*]] = named_table
# CHECK-NEXT: %[[V1:.*]] = project %[[V0]] : tuple<si1> -> tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second> {
# CHECK-NEXT: ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
# CHECK-NEXT: %[[V2:.*]] = literal #substrait.interval_year_month<2024y 1m>
# CHECK-NEXT: %[[V3:.*]] = literal #substrait.interval_day_second<9d 8000s>
# CHECK-NEXT: yield %[[V2]], %[[V3]] : !substrait.interval_year_month, !substrait.interval_day_second
# CHECK-NEXT: }
# CHECK-NEXT: yield %[[V1]] : tuple<si1, !substrait.interval_year_month, !substrait.interval_day_second>

relations {
rel {
project {
common {
direct {
}
}
input {
read {
common {
direct {
}
}
base_schema {
names: "a"
struct {
types {
bool {
nullability: NULLABILITY_REQUIRED
}
}
nullability: NULLABILITY_REQUIRED
}
}
named_table {
names: "t1"
}
}
}
expressions {
literal {
interval_year_to_month {
years: 2024
months: 1
}
}
}
expressions {
literal {
interval_day_to_second {
days: 9
seconds: 8000
}
}
}
}
}
}
version {
minor_number: 42
patch_number: 1
}

# -----

# CHECK: substrait.plan version 0 : 42 : 1 {
# CHECK-NEXT: relation
# CHECK: %[[V0:.*]] = named_table
Expand Down
Loading

0 comments on commit 91b1ab0

Please sign in to comment.