introduce date type

substrait-io · Jan 9, 2025 · 78d8129 · 78d8129
1 parent fe77d82
commit 78d8129
Show file tree

Hide file tree

Showing 6 changed files with 102 additions and 0 deletions.
diff --git 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
@@ -25,6 +25,22 @@ 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.
@@ -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
@@ -195,6 +195,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()) {
@@ -570,6 +582,9 @@ SubstraitExporter::exportOperation(LiteralOp op) {
     default:
       op->emitOpError("has float value with unsupported width");
     }
+  } // `DateType`.
+  else if (literalType.isa<DateType>()) {
+    literal->set_date(value.cast<DateAttr>().getValue());
   } else
     op->emitOpError("has unsupported value");
 

diff --git a/lib/Target/SubstraitPB/Import.cpp b/lib/Target/SubstraitPB/Import.cpp
@@ -108,6 +108,8 @@ static mlir::FailureOr<mlir::Type> 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<mlir::Type> fieldTypes;
@@ -328,6 +330,10 @@ importLiteral(ImplicitLocOpBuilder builder,
     auto attr = FloatAttr::get(FloatType::getF64(context), message.fp64());
     return builder.create<LiteralOp>(attr);
   }
+  case Expression::Literal::LiteralTypeCase::kDate: {
+    auto attr = DateAttr::get(context, message.date());
+    return builder.create<LiteralOp>(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
@@ -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 {

diff --git 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<!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