From 78d8129c13c569d8e2c06a35c182b29851176d63 Mon Sep 17 00:00:00 2001 From: Dalia Shaaban <144673861+dshaaban01@users.noreply.github.com> Date: Wed, 1 Jan 2025 16:25:21 +0100 Subject: [PATCH] introduce date type --- .../Dialect/Substrait/IR/SubstraitDialect.td | 2 ++ .../Dialect/Substrait/IR/SubstraitTypes.td | 18 ++++++++++ lib/Target/SubstraitPB/Export.cpp | 15 ++++++++ lib/Target/SubstraitPB/Import.cpp | 6 ++++ test/Target/SubstraitPB/Export/types.mlir | 25 +++++++++++++ test/Target/SubstraitPB/Import/types.textpb | 36 +++++++++++++++++++ 6 files changed, 102 insertions(+) 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 84db425..226078e 100644 --- a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td +++ b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td @@ -25,6 +25,22 @@ 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. @@ -37,6 +53,7 @@ def Substrait_AtomicTypes { SI64, // I64 F32, // FP32 F64, // FP64 + Substrait_DateType // Date ]; } @@ -50,6 +67,7 @@ def Substrait_AtomicAttributes { SI64Attr, // I64 F32Attr, // FP32 F64Attr, // FP64 + Substrait_DateAttr // Date ]; } diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp index e85e20e..a82ef67 100644 --- a/lib/Target/SubstraitPB/Export.cpp +++ b/lib/Target/SubstraitPB/Export.cpp @@ -195,6 +195,18 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) { return std::move(type); } + // Handle date. + if (mlirType.isa()) { + // TODO(ingomueller): support other nullability modes. + auto dateType = std::make_unique(); + 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()) { @@ -570,6 +582,9 @@ SubstraitExporter::exportOperation(LiteralOp op) { default: op->emitOpError("has float 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 75a77c4..e0d8bba 100644 --- a/lib/Target/SubstraitPB/Import.cpp +++ b/lib/Target/SubstraitPB/Import.cpp @@ -108,6 +108,8 @@ static mlir::FailureOr importType(MLIRContext *context, return FloatType::getF32(context); case proto::Type::kFp64: return FloatType::getF64(context); + case proto::Type::kDate: + return DateType::get(context); case proto::Type::kStruct: { const proto::Type::Struct &structType = type.struct_(); llvm::SmallVector fieldTypes; @@ -328,6 +330,10 @@ importLiteral(ImplicitLocOpBuilder builder, auto attr = FloatAttr::get(FloatType::getF64(context), message.fp64()); return builder.create(attr); } + case Expression::Literal::LiteralTypeCase::kDate: { + auto attr = DateAttr::get(context, message.date()); + return builder.create(attr); + } // TODO(ingomueller): Support more types. default: { const pb::FieldDescriptor *desc = diff --git a/test/Target/SubstraitPB/Export/types.mlir b/test/Target/SubstraitPB/Export/types.mlir index 700917b..bb91e91 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 7ecfc50..2cf87b2 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