feat: introduce ListType<T> type and attribute

include/substrait-mlir/Dialect/Substrait/IR/SubstraitAttrs.td

@@ -290,10 +290,30 @@ def Substrait_ParametrizedAttributes {
 /// Attribute of one of the currently supported atomic types.
 def Substrait_AtomicAttribute : AnyAttrOf<Substrait_SimpleAttributes.attrs#Substrait_ParametrizedAttributes.attrs>;
 
+def Substrait_ListAttr : Substrait_Attr<"List", "list", [TypedAttrInterface]> {
+  let summary = "Substrait list attribute";
+  let description = [{
+    This attribute represents a list of atomic attributes, parameterized with a ListType.
+  }];
+  // TODO: Extend to list of container attributes.
+  // TODO: Implement custom parser and printer.
+  let parameters = (ins "mlir::ArrayAttr":$value, "ListType":$type);
+  let assemblyFormat = [{ `<` $value `,` $type `>` }];
+  let genVerifyDecl = 1;
+}
+
+def Substrait_ContainerAttributes {
+  list<Attr> attrs = [
+    Substrait_ListAttr, // List
+  ];
+}
+
+def Substrait_ContainerAttribute : AnyAttrOf<Substrait_ContainerAttributes.attrs>;
+
 /// Attribute of one of the currently supported atomic or container types.
 def Substrait_ExpressionAttribute :
   AnyAttrOf<[
-    // TODO: add container attributes here once we have them.
+    Substrait_ContainerAttribute,
     Substrait_AtomicAttribute,
   ]>;
 

include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td

@@ -180,14 +180,33 @@ def Substrait_AtomicTypes {
 /// One of the currently supported atomic types.
 def Substrait_AtomicType :
   AnyTypeOf<Substrait_SimpleTypes.types#Substrait_ParametrizedTypes.types>;
+
+def Substrait_ListType : Substrait_Type<"List", "list"> {
+  let summary = "Substrait list type";
+  let description = [{
+    This type represents a list type, where type is atomic.
+  }];
+  // TODO: Extend to list of container attributes.
+  let parameters = (ins Substrait_AtomicType:$type);
+  let assemblyFormat = [{ `<` $type `>` }];
+  let genVerifyDecl = 1;
+}
+
+def Substrait_ContainerTypes {
+  list<Type> types = [
+    Substrait_ListType, // List
+  ];
+}
 
 /// Any container type, i.e., structs, maps, lists, and nestings thereof.
-def Substrait_ContainerType : NestedTupleOf<Substrait_AtomicTypes.types>;
+def Substrait_ContainerType : AnyTypeOf<
+  Substrait_ContainerTypes.types # [NestedTupleOf<Substrait_AtomicTypes.types>]
+>;
 
 /// Currently supported expression types.
 def Substrait_ExpressionTypes {
   list<Type> types =
-  Substrait_AtomicTypes.types;
+  Substrait_AtomicTypes.types#Substrait_ContainerTypes.types;
   // TODO: add container types here once we have them.
 }
 

lib/Dialect/Substrait/IR/Substrait.cpp

@@ -116,6 +116,29 @@ LogicalResult mlir::substrait::IntervalDaySecondAttr::verify(
   return success();
 }
 
+LogicalResult mlir::substrait::ListAttr::verify(
+    llvm::function_ref<mlir::InFlightDiagnostic()> emitError, ArrayAttr value,
+    ListType listType) {
+
+  mlir::Type expectedType =
+      listType.getType(); // The atomic type the list is parameterized with
+
+  for (mlir::Attribute attr : value) {
+    auto typedAttr = mlir::dyn_cast<mlir::TypedAttr>(attr);
+    if (!typedAttr) {
+      return emitError()
+             << "'ListAttr' values must be typed attributes, but got: " << attr;
+    }
+
+    if (typedAttr.getType() != expectedType) {
+      return emitError() << "mismatched element type in 'ListAttr': expected "
+                         << expectedType << ", but got " << typedAttr.getType();
+    }
+  }
+
+  return success();
+}
+
 LogicalResult mlir::substrait::VarCharAttr::verify(
     llvm::function_ref<mlir::InFlightDiagnostic()> emitError, StringAttr value,
     VarCharType type) {
@@ -144,6 +167,24 @@ LogicalResult mlir::substrait::DecimalType::verify(
   return success();
 }
 
+LogicalResult mlir::substrait::ListType::verify(
+    llvm::function_ref<mlir::InFlightDiagnostic()> emitError, mlir::Type type) {
+  // ListType must be parameterized with an atomic type.
+  if (!mlir::isa<mlir::IntegerType, mlir::Float32Type, mlir::Float64Type,
+                 substrait::StringType, substrait::BinaryType,
+                 substrait::TimestampType, substrait::TimestampTzType,
+                 substrait::DateType, substrait::TimeType,
+                 substrait::IntervalYearMonthType,
+                 substrait::IntervalDaySecondType, substrait::UUIDType,
+                 substrait::FixedCharType, substrait::VarCharType,
+                 substrait::FixedBinaryType, substrait::DecimalType>(type))
+    return emitError()
+           << "'ListType' must be parameterized with an atomic type, but got: "
+           << type;
+
+  return mlir::success();
+}
+
 // Count the number of digits in an APInt in base 10.
 static size_t countDigits(const APInt &value) {
   llvm::SmallVector<char> buffer;

lib/Target/SubstraitPB/Export.cpp

@@ -90,6 +90,9 @@ class SubstraitExporter {
   std::unique_ptr<pb::Any> exportAny(StringAttr attr);
   FailureOr<std::unique_ptr<NamedStruct>>
   exportNamedStruct(Location loc, ArrayAttr fieldNames, TupleType relationType);
+  FailureOr<std::unique_ptr<Expression::Literal>>
+  exportAttribute(Attribute value,
+                  std::function<InFlightDiagnostic()> emitError);
   FailureOr<std::unique_ptr<pb::Message>> exportOperation(Operation *op);
   FailureOr<std::unique_ptr<proto::Type>> exportType(Location loc,
                                                      mlir::Type mlirType);
@@ -431,6 +434,20 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
     return std::move(type);
   }
 
+  // Handle list type.
+  if (auto listType = llvm::dyn_cast<ListType>(mlirType)) {
+    // TODO(ingomueller): support other nullability modes.
+    auto listTypeProto = std::make_unique<proto::Type::List>();
+    listTypeProto->set_nullability(
+        Type_Nullability::Type_Nullability_NULLABILITY_REQUIRED);
+    listTypeProto->set_allocated_type(
+        exportType(loc, listType.getType()).value().release());
+
+    auto type = std::make_unique<proto::Type>();
+    type->set_allocated_list(listTypeProto.release());
+    return std::move(type);
+  }
+
   // Handle tuple types.
   if (auto tupleType = llvm::dyn_cast<TupleType>(mlirType)) {
     auto structType = std::make_unique<proto::Type::Struct>();
@@ -947,17 +964,18 @@ SubstraitExporter::exportOperation(FilterOp op) {
   return rel;
 }
 
-FailureOr<std::unique_ptr<Expression>>
-SubstraitExporter::exportOperation(LiteralOp op) {
+FailureOr<std::unique_ptr<Expression::Literal>>
+SubstraitExporter::exportAttribute(
+    Attribute value, std::function<InFlightDiagnostic()> emitError) {
   // Build `Literal` message depending on type.
-  Attribute value = op.getValue();
   mlir::Type literalType = getAttrType(value);
   auto literal = std::make_unique<Expression::Literal>();
 
   // `IntegerType`s.
   if (auto intType = dyn_cast<IntegerType>(literalType)) {
     if (!intType.isSigned())
-      op->emitOpError("has integer value with unsupported signedness");
+      // has integer value with unsupported signedness
+      return emitError() << " has integer value with unsupported signedness";
     switch (intType.getWidth()) {
     case 1:
       literal->set_boolean(mlir::cast<IntegerAttr>(value).getSInt());
@@ -976,7 +994,8 @@ SubstraitExporter::exportOperation(LiteralOp op) {
       literal->set_i64(mlir::cast<IntegerAttr>(value).getSInt());
       break;
     default:
-      op->emitOpError("has integer value with unsupported width");
+      // has integer value with unsupported width
+      return emitError() << " has integer value with unsupported width";
     }
   }
   // `FloatType`s.
@@ -990,7 +1009,7 @@ SubstraitExporter::exportOperation(LiteralOp op) {
       literal->set_fp64(mlir::cast<FloatAttr>(value).getValueAsDouble());
       break;
     default:
-      op->emitOpError("has float value with unsupported width");
+      return emitError() << " float value with unsupported width";
     }
   }
   // `StringType`.
@@ -1067,12 +1086,45 @@ SubstraitExporter::exportOperation(LiteralOp op) {
     decimal->set_precision(decimalType.getPrecision());
     decimal->set_value(res);
     literal->set_allocated_decimal(decimal.release());
-  } else
-    op->emitOpError("has unsupported value");
+  } // `ListType`.
+  else if (auto listType = dyn_cast<ListType>(literalType)) {
+    auto list = std::make_unique<::substrait::proto::Expression_Literal_List>();
+    auto listAttr = mlir::cast<ListAttr>(value);
+    for (mlir::Attribute attr : listAttr.getValue()) {
+
+      FailureOr<std::unique_ptr<Expression::Literal>> listElement =
+          exportAttribute(attr, emitError);
+      if (failed(listElement))
+        return emitError() << "Failed to export list element attribute: "
+                           << attr;
+
+      auto *literalExpr = dyn_cast<Expression_Literal>(listElement->get());
+      if (!literalExpr)
+        return emitError()
+               << "Exported list element is not a valid Expression_Literal: "
+               << attr;
+
+      list->add_values()->CopyFrom(*literalExpr);
+    }
+    literal->set_allocated_list(list.release());
+  } else {
+    return emitError() << " has unsupported value";
+  }
+  return literal;
+}
+
+FailureOr<std::unique_ptr<Expression>>
+SubstraitExporter::exportOperation(LiteralOp op) {
+  // Build `Literal` message depending on type.
+  Attribute value = op.getValue();
+  auto literal = exportAttribute(value, [&]() { return op->emitOpError(); });
+
+  if (failed(literal))
+    return op->emitOpError("has unsupported value");
 
   // Build `Expression` message.
   auto expression = std::make_unique<Expression>();
-  expression->set_allocated_literal(literal.release());
+  expression->set_allocated_literal(literal->release());
 
   return expression;
 }