introduce binary type

substrait-io · Dec 27, 2024 · c99c873 · c99c873
1 parent abbd35d
commit c99c873
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 2 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,21 +25,39 @@ class Substrait_Attr<string name, string typeMnemonic, list<Trait> traits = []>
   let mnemonic = typeMnemonic;
 }
 
+def Substrait_BinaryType : Substrait_Type<"Binary", "binary"> {
+  let summary = "Substrait binary type";
+  let description = [{
+    This type represents a substrait binary type.
+  }];
+}
+
+def Substrait_BinaryAttr : Substrait_Attr<"Binary", "binary"> {
+  let summary = "Substrait binary type";
+  let description = [{
+    This type represents a substrait binary attribute type.
+  }];
+  let parameters = (ins ArrayRefParameter<"uint8_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<Type> types = [
     SI1, // Boolean
-    SI32 // I32
+    SI32, // I32
+    Substrait_BinaryType // Binary
   ];
 }
 
 /// Attributes of currently supported atomic types.
 def Substrait_AtomicAttributes {
   list<Attr> attrs = [
     SI1Attr, // Boolean
-    SI32Attr // I32
+    SI32Attr, // I32
+    Substrait_BinaryAttr // Binary
   ];
 }
 

diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp
@@ -118,6 +118,14 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
     return std::move(type);
   }
 
+  // Handle binary type.
+  if (auto binaryType = mlirType.dyn_cast<BinaryType>()) {
+    auto binary = std::make_unique<proto::Type::Binary>();
+    auto type = std::make_unique<proto::Type>();
+    type->set_allocated_binary(binary.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()) {
@@ -428,6 +436,10 @@ SubstraitExporter::exportOperation(LiteralOp op) {
     default:
       op->emitOpError("has integer value with unsupported width");
     }
+  } // `BinaryType`.
+  else if (auto binaryType = dyn_cast<BinaryType>(literalType)) {
+    ArrayRef<uint8_t> ref = value.cast<BinaryAttr>().getValue();
+    literal->set_binary(std::string(ref.begin(), ref.end()));
   } else
     op->emitOpError("has unsupported value");
 

diff --git a/lib/Target/SubstraitPB/Import.cpp b/lib/Target/SubstraitPB/Import.cpp
@@ -97,6 +97,8 @@ static mlir::FailureOr<mlir::Type> importType(MLIRContext *context,
     return IntegerType::get(context, 1, IntegerType::Signed);
   case proto::Type::kI32:
     return IntegerType::get(context, 32, IntegerType::Signed);
+  case proto::Type::kBinary:
+    return BinaryType::get(context);
   case proto::Type::kStruct: {
     const proto::Type::Struct &structType = type.struct_();
     llvm::SmallVector<mlir::Type> fieldTypes;
@@ -266,6 +268,13 @@ importLiteral(ImplicitLocOpBuilder builder,
         IntegerType::get(context, 32, IntegerType::Signed), message.i32());
     return builder.create<LiteralOp>(attr);
   }
+  case Expression::Literal::LiteralTypeCase::kBinary: {
+    auto attr = BinaryAttr::get(
+        context, ArrayRef<uint8_t>(
+                     reinterpret_cast<const uint8_t *>(message.binary().data()),
+                     message.binary().size()));
+    return builder.create<LiteralOp>(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
@@ -9,6 +9,29 @@
 // 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:              binary {
+// CHECK:                     nullability: NULLABILITY_REQUIRED
+// CHECK-NEXT:              }
+// CHECK-NEXT:            }
+// CHECK-NEXT:            named_table {
+// CHECK-NEXT:              names: "t1"
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<!substrait.binary>
+    yield %0 : tuple<!substrait.binary>
+  }
+}
+
+// -----
+
 // 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,41 @@
 # RUN:   --split-input-file="# ""-----" --output-split-marker="// -----" \
 # RUN: | FileCheck %s
 
+# CHECK:      substrait.plan
+# CHECK-NEXT:   relation
+# CHECK-NEXT:     named_table
+# CHECK-SAME:       : tuple<!substrait.binary>
+
+relations {
+  rel {
+    read {
+      common {
+        direct {
+        }
+      }
+      base_schema {
+        names: "a"
+        struct {
+          types {
+            binary {
+            }
+          }
+          nullability: NULLABILITY_REQUIRED
+        }
+      }
+      named_table {
+        names: "t1"
+      }
+    }
+  }
+}
+version {
+  minor_number: 42
+  patch_number: 1
+}
+
+# -----
+
 # CHECK:      substrait.plan
 # CHECK-NEXT:   relation
 # CHECK-NEXT:     named_table