diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td index 3589c43..630a3ab 100644 --- a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td +++ b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td @@ -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 diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td index 49416ee..537849d 100644 --- a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td +++ b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td @@ -25,13 +25,30 @@ class Substrait_Attr 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 types = [ SI1, // Boolean - SI32 // I32 + SI32, // I32 + Substrait_DateType // Date ]; } @@ -39,7 +56,8 @@ def Substrait_AtomicTypes { def Substrait_AtomicAttributes { list attrs = [ SI1Attr, // Boolean - SI32Attr // I32 + SI32Attr, // I32 + Substrait_DateAttr // Date ]; } diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp index a6b1b4c..ab71d23 100644 --- a/lib/Target/SubstraitPB/Export.cpp +++ b/lib/Target/SubstraitPB/Export.cpp @@ -118,6 +118,18 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) { return std::move(type); } + // Handle date. + if (mlirType.isa()) { + auto dateType = std::make_unique(); + // TODO(ingomueller): support other nullability modes. + dateType->set_nullability( + Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED); + + auto type = std::make_unique(); + type->set_allocated_date(dateType.release()); + return std::move(type); + } + if (auto tupleType = llvm::dyn_cast(mlirType)) { auto structType = std::make_unique(); 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()) { + literal->set_date(value.cast().getValue()); } else op->emitOpError("has unsupported value"); diff --git a/lib/Target/SubstraitPB/Import.cpp b/lib/Target/SubstraitPB/Import.cpp index 11e688e..156ec84 100644 --- a/lib/Target/SubstraitPB/Import.cpp +++ b/lib/Target/SubstraitPB/Import.cpp @@ -97,6 +97,8 @@ static mlir::FailureOr 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 fieldTypes; @@ -266,6 +268,10 @@ importLiteral(ImplicitLocOpBuilder builder, IntegerType::get(context, 32, IntegerType::Signed), message.i32()); return builder.create(attr); } + case Expression::Literal::LiteralTypeCase::kDate: { + auto attr = DateAttr::get(context, message.date()); + return builder.create(attr); + } default: { const pb::FieldDescriptor *desc = Expression::Literal::GetDescriptor()->FindFieldByNumber(literalType); diff --git a/test/Target/SubstraitPB/Export/types.mlir b/test/Target/SubstraitPB/Export/types.mlir index 4d8c7ca..962de15 100644 --- a/test/Target/SubstraitPB/Export/types.mlir +++ b/test/Target/SubstraitPB/Export/types.mlir @@ -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 + yield %0 : tuple + } +} + +// ----- + // CHECK-LABEL: relations { // CHECK-NEXT: rel { // CHECK-NEXT: read { diff --git a/test/Target/SubstraitPB/Import/types.textpb b/test/Target/SubstraitPB/Import/types.textpb index d2d7c94..7043bec 100644 --- a/test/Target/SubstraitPB/Import/types.textpb +++ b/test/Target/SubstraitPB/Import/types.textpb @@ -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 + +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