date done

substrait-io · Jan 16, 2025 · 122a2be · 122a2be
1 parent ea450f9
commit 122a2be
Show file tree

Hide file tree

Showing 6 changed files with 119 additions and 5 deletions.
diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
@@ -81,7 +81,8 @@ def Substrait_DateType : Substrait_Type<"Date", "date"> {
   }];
 }
 
-def Substrait_DateAttr : Substrait_Attr<"Date", "date"> {
+def Substrait_DateAttr : Substrait_Attr<"Date", "date",
+  [DeclareAttrInterfaceMethods<TypedAttrInterface>]> {
   let summary = "Substrait date type";
   let description = [{
     This type represents a substrait date attribute type.

diff --git a/lib/Dialect/Substrait/IR/Substrait.cpp b/lib/Dialect/Substrait/IR/Substrait.cpp
@@ -87,10 +87,17 @@ void printCountAsAll(OpAsmPrinter &printer, Operation *op, IntegerAttr count) {
 //===----------------------------------------------------------------------===//
 
 /// Implement the getType method for custom type `TimestampAttr`.
-::mlir::Type TimestampAttr::getType() const { return TimestampType::get(getContext()); }
+::mlir::Type TimestampAttr::getType() const {
+  return TimestampType::get(getContext());
+}
 
 /// Implement the getType method for custom type `TimestampTzAttr`.
-::mlir::Type TimestampTzAttr::getType() const { return TimestampTzType::get(getContext()); }
+::mlir::Type TimestampTzAttr::getType() const {
+  return TimestampTzType::get(getContext());
+}
+
+/// Implement the getType method for custom type `DateAttr`.
+::mlir::Type DateAttr::getType() const { return DateType::get(getContext()); }
 
 //===----------------------------------------------------------------------===//
 // Substrait operations

diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp
@@ -654,8 +654,9 @@ SubstraitExporter::exportOperation(LiteralOp op) {
     literal->set_timestamp(value.cast<TimestampAttr>().getValue());
   } else if (literalType.isa<TimestampTzType>()) {
     literal->set_timestamp_tz(value.cast<TimestampTzAttr>().getValue());
-  } // `DateType`.
-  else if (literalType.isa<DateType>()) {
+  } 
+  // `DateType`.
+  else if (auto binaryType = dyn_cast<DateType>(literalType)) {
     literal->set_date(value.cast<DateAttr>().getValue());
   } else
     op->emitOpError("has unsupported value");

diff --git a/test/Dialect/Substrait/literal.mlir b/test/Dialect/Substrait/literal.mlir
@@ -1,6 +1,30 @@
 // 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.date> {
+// CHECK-NEXT:    ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
+// CHECK-NEXT:      %[[V2:.*]] = literal #substrait.date<200000000> : !substrait.date
+// CHECK-NEXT:      yield %[[V2]] : !substrait.date
+// CHECK-NEXT:    }
+// CHECK-NEXT:    yield %[[V1]] : tuple<si1, !substrait.date>
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<si1>
+    %1 = project %0 : tuple<si1> -> tuple<si1, !substrait.date> {
+    ^bb0(%arg : tuple<si1>):
+      %date = literal #substrait.date<200000000> : !substrait.date
+      yield %date : !substrait.date
+    }
+    yield %1 : tuple<si1, !substrait.date> 
+  }
+}
+
+// -----
+
 // CHECK:      substrait.plan version 0 : 42 : 1 {
 // CHECK-NEXT:   relation
 // CHECK:         %[[V0:.*]] = named_table

diff --git a/test/Target/SubstraitPB/Export/literal.mlir b/test/Target/SubstraitPB/Export/literal.mlir
@@ -9,6 +9,33 @@
 // 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:            date: 200000000
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<si1>
+    %1 = project %0 : tuple<si1> -> tuple<si1, !substrait.date> {
+    ^bb0(%arg : tuple<si1>):
+      %date = literal #substrait.date<200000000> : !substrait.date
+      yield %date : !substrait.date
+    }
+    yield %1 : tuple<si1, !substrait.date> 
+  }
+}
+
+// -----
+
 // CHECK-LABEL: relations {
 // CHECK-NEXT:    rel {
 // CHECK-NEXT:      project {

diff --git a/test/Target/SubstraitPB/Import/literal.textpb b/test/Target/SubstraitPB/Import/literal.textpb
@@ -10,6 +10,60 @@
 # 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.date> {
+# CHECK-NEXT:    ^[[BB0:.*]](%[[ARG0:.*]]: tuple<si1>):
+# CHECK-NEXT:      %[[V2:.*]] = literal #substrait.date<200000000> : !substrait.date
+# CHECK-NEXT:      yield %[[V2]] : !substrait.date
+# CHECK-NEXT:    }
+# CHECK-NEXT:    yield %[[V1]] : tuple<si1, !substrait.date>
+
+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 {
+          date: 200000000
+        }
+      }
+    }
+  }
+}
+version {
+  minor_number: 42
+  patch_number: 1
+}
+
+# -----
+
 # CHECK:      substrait.plan version 0 : 42 : 1 {
 # CHECK-NEXT:   relation
 # CHECK:         %[[V0:.*]] = named_table