Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
introduce date type
Browse files Browse the repository at this point in the history
dshaaban01 committed Jan 1, 2025

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
1 parent 584a68e commit 4f24fa3
Showing 6 changed files with 104 additions and 2 deletions.
Original file line number Diff line number Diff line change
@@ -34,6 +34,8 @@ def Substrait_Dialect : Dialect {
more natural in MLIR to represent several message types as a single op and
express message sub-types with interfaces instead.
}];
let useDefaultAttributePrinterParser = 1;
let useDefaultTypePrinterParser = 1;
}

#endif // SUBSTRAIT_DIALECT_SUBSTRAIT_IR_SUBSTRAITDIALECT
22 changes: 20 additions & 2 deletions include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
Original file line number Diff line number Diff line change
@@ -25,21 +25,39 @@ class Substrait_Attr<string name, string typeMnemonic, list<Trait> traits = []>
let mnemonic = typeMnemonic;
}

def Substrait_DateType : Substrait_Type<"Date", "date"> {
let summary = "Substrait date type";
let description = [{
This type represents a substrait date type.
}];
}

def Substrait_DateAttr : Substrait_Attr<"Date", "date"> {
let summary = "Substrait date type";
let description = [{
This type represents a substrait date attribute type.
}];
let parameters = (ins "int32_t":$value);
let assemblyFormat = [{ `<` $value `>` }];
}

/// Currently supported atomic types. These correspond directly to the types in
/// https://github.com/substrait-io/substrait/blob/main/proto/substrait/type.proto.
// TODO(ingomueller): Add the other low-hanging fruits here.
def Substrait_AtomicTypes {
list<Type> types = [
SI1, // Boolean
SI32 // I32
SI32, // I32
Substrait_DateType // Date
];
}

/// Attributes of currently supported atomic types.
def Substrait_AtomicAttributes {
list<Attr> attrs = [
SI1Attr, // Boolean
SI32Attr // I32
SI32Attr, // I32
Substrait_DateAttr // Date
];
}

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

// Handle date.
if (mlirType.isa<DateType>()) {
// TODO(ingomueller): support other nullability modes.
auto dateType = std::make_unique<proto::Type::Date>();
dateType->set_nullability(
Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);

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

if (auto tupleType = llvm::dyn_cast<TupleType>(mlirType)) {
auto structType = std::make_unique<proto::Type::Struct>();
for (mlir::Type fieldType : tupleType.getTypes()) {
@@ -428,6 +440,9 @@ SubstraitExporter::exportOperation(LiteralOp op) {
default:
op->emitOpError("has integer value with unsupported width");
}
} // `DateType`.
else if (literalType.isa<DateType>()) {
literal->set_date(value.cast<DateAttr>().getValue());
} else
op->emitOpError("has unsupported value");

6 changes: 6 additions & 0 deletions lib/Target/SubstraitPB/Import.cpp
Original file line number Diff line number Diff line change
@@ -97,6 +97,8 @@ static mlir::FailureOr<mlir::Type> importType(MLIRContext *context,
return IntegerType::get(context, 1, IntegerType::Signed);
case proto::Type::kI32:
return IntegerType::get(context, 32, IntegerType::Signed);
case proto::Type::kDate:
return DateType::get(context);
case proto::Type::kStruct: {
const proto::Type::Struct &structType = type.struct_();
llvm::SmallVector<mlir::Type> fieldTypes;
@@ -266,6 +268,10 @@ importLiteral(ImplicitLocOpBuilder builder,
IntegerType::get(context, 32, IntegerType::Signed), message.i32());
return builder.create<LiteralOp>(attr);
}
case Expression::Literal::LiteralTypeCase::kDate: {
auto attr = DateAttr::get(context, message.date());
return builder.create<LiteralOp>(attr);
}
default: {
const pb::FieldDescriptor *desc =
Expression::Literal::GetDescriptor()->FindFieldByNumber(literalType);
25 changes: 25 additions & 0 deletions test/Target/SubstraitPB/Export/types.mlir
Original file line number Diff line number Diff line change
@@ -9,6 +9,31 @@
// 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: struct {
// CHECK-NEXT: types {
// CHECK-NEXT: date {
// 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"] : tuple<!substrait.date>
yield %0 : tuple<!substrait.date>
}
}

// -----

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

# CHECK: substrait.plan
# CHECK-NEXT: relation
# CHECK-NEXT: named_table
# CHECK-SAME: : tuple<!substrait.date>

relations {
rel {
read {
common {
direct {
}
}
base_schema {
names: "a"
struct {
types {
date {
nullability: NULLABILITY_REQUIRED
}
}
nullability: NULLABILITY_REQUIRED
}
}
named_table {
names: "t1"
}
}
}
}
version {
minor_number: 42
patch_number: 1
}

# -----

# CHECK: substrait.plan
# CHECK-NEXT: relation
# CHECK-NEXT: named_table

0 comments on commit 4f24fa3

Please sign in to comment.