introduce string type and type tests

substrait-io · Jan 9, 2025 · 434cbeb · 434cbeb
1 parent fe77d82
commit 434cbeb
Show file tree

Hide file tree

Showing 7 changed files with 107 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,7 @@ 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 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
@@ -12,6 +12,7 @@
 include "substrait-mlir/Dialect/Substrait/IR/SubstraitDialect.td"
 include "mlir/IR/CommonTypeConstraints.td"
 include "mlir/IR/OpBase.td"
+include "mlir/IR/CommonAttrConstraints.td"
 
 // Base class for Substrait dialect types.
 class Substrait_Type<string name, string typeMnemonic, list<Trait> traits = []>
@@ -25,6 +26,13 @@ class Substrait_Attr<string name, string typeMnemonic, list<Trait> traits = []>
   let mnemonic = typeMnemonic;
 }
 
+def Substrait_String : Substrait_Type<"String", "string"> {
+  let summary = "Substrait string type";
+  let description = [{
+    This type represents a substrait string type.
+  }];
+}
+
 /// 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 +45,7 @@ def Substrait_AtomicTypes {
     SI64, // I64
     F32, // FP32
     F64, // FP64
+    Substrait_String, // String
   ];
 }
 
@@ -50,6 +59,7 @@ def Substrait_AtomicAttributes {
     SI64Attr, // I64
     F32Attr, // FP32
     F64Attr, // FP64
+    StrAttr, // String
   ];
 }
 

diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp
@@ -194,6 +194,17 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
     type->set_allocated_fp64(fp64Type.release());
     return std::move(type);
   }
+  // Handle String.
+  if (mlirType.isa<StringType>()) {
+    // TODO(ingomueller): support other nullability modes.
+    auto stringType = std::make_unique<proto::Type::String>();
+    stringType->set_nullability(
+        Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);
+
+    auto type = std::make_unique<proto::Type>();
+    type->set_allocated_string(stringType.release());
+    return std::move(type);
+  }
 
   if (auto tupleType = llvm::dyn_cast<TupleType>(mlirType)) {
     auto structType = std::make_unique<proto::Type::Struct>();
@@ -570,6 +581,10 @@ SubstraitExporter::exportOperation(LiteralOp op) {
     default:
       op->emitOpError("has float value with unsupported width");
     }
+  }
+  // `StringType`.
+  else if (auto stringType = dyn_cast<StringType>(literalType)) {
+    literal->set_string(value.cast<StringAttr>().getValue().str());
   } 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::kString:
+    return StringType::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::kString: {
+    auto attr = StringAttr::get(context, message.string());
+    return builder.create<LiteralOp>(attr);
+  }
   // TODO(ingomueller): Support more types.
   default: {
     const pb::FieldDescriptor *desc =

diff --git a/test/Dialect/Substrait/types.mlir b/test/Dialect/Substrait/types.mlir
@@ -1,6 +1,20 @@
 // RUN: substrait-opt -split-input-file %s \
 // RUN: | FileCheck %s
 
+// CHECK-LABEL: substrait.plan
+// CHECK:         relation
+// CHECK:         %[[V0:.*]] = named_table @t1 as ["a"] : tuple<!substrait.string>
+// CHECK-NEXT:    yield %0 : tuple<!substrait.string>
+
+substrait.plan version 0 : 42 : 1 {
+  relation {
+    %0 = named_table @t1 as ["a"] : tuple<!substrait.string>
+    yield %0 : tuple<!substrait.string>
+  }
+}
+
+// -----
+
 // CHECK-LABEL: substrait.plan
 // CHECK:         relation
 // CHECK:         %[[V0:.*]] = named_table @t1 as ["a", "b"] : tuple<f32, f64>

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:              string {
+// 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.string>
+    yield %0 : tuple<!substrait.string>
+  }
+}
+
+// -----
+
 // 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.string>
+
+relations {
+  rel {
+    read {
+      common {
+        direct {
+        }
+      }
+      base_schema {
+        names: "a"
+        struct {
+          types {
+            string {
+              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