feat: add a custom assembly for union distinct operation (#17)

Add a custom assembly with explicit return type and adapt tests.
substrait-io · Nov 26, 2024 · e3437ea · e3437ea
1 parent cffb42c
commit e3437ea
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitOps.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitOps.td
@@ -394,6 +394,9 @@ def Substrait_UnionDistinct_Op : Substrait_RelOp<"union_distinct"> {
   );
 
   let results = (outs Substrait_Relation:$result);
+  let assemblyFormat = [{
+    $left `u` $right attr-dict `:` type($left) `u` type($right) `->` type($result)
+  }];
 }
 
 def Substrait_EmitOp : Substrait_RelOp<"emit", [

diff --git a/test/Dialect/Substrait/union-distinct.mlir b/test/Dialect/Substrait/union-distinct.mlir
@@ -5,15 +5,16 @@
 // CHECK:         relation
 // CHECK:           %[[V0:.*]] = named_table
 // CHECK:           %[[V1:.*]] = named_table
-// CHECK-NEXT:      %[[V2:.*]] = "substrait.union_distinct"(%[[V0]], %[[V1]]) 
-// CHECK-SAME:        : (tuple<si32>, tuple<si32>) -> tuple<si32>
+// CHECK-NEXT:      %[[V2:.*]] = union_distinct %[[V0]] u %[[V1]] 
+// CHECK-SAME:        : tuple<si32> u tuple<si32> -> tuple<si32>
 // CHECK-NEXT:      yield %[[V2]] : tuple<si32>
 
+
 substrait.plan version 0 : 42 : 1 {
   relation {
     %0 = named_table @t1 as ["a"] : tuple<si32>
     %1 = named_table @t2 as ["b"] : tuple<si32>
-    %2 = "substrait.union_distinct"(%0 , %1) : (tuple<si32>,tuple<si32>) -> tuple<si32>
+    %2 = union_distinct %0 u %1 : tuple<si32> u tuple<si32> -> tuple<si32>
     yield %2 : tuple<si32>
   }
 }