in spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala [87:181]
def serializeDataType(dt: DataType): Option[ExprOuterClass.DataType] = {
val typeId = dt match {
case _: BooleanType => 0
case _: ByteType => 1
case _: ShortType => 2
case _: IntegerType => 3
case _: LongType => 4
case _: FloatType => 5
case _: DoubleType => 6
case _: StringType => 7
case _: BinaryType => 8
case _: TimestampType => 9
case _: DecimalType => 10
case _: TimestampNTZType => 11
case _: DateType => 12
case _: NullType => 13
case _: ArrayType => 14
case _: MapType => 15
case _: StructType => 16
case dt =>
emitWarning(s"Cannot serialize Spark data type: $dt")
return None
}
val builder = ProtoDataType.newBuilder()
builder.setTypeIdValue(typeId)
// Decimal
val dataType = dt match {
case t: DecimalType =>
val info = DataTypeInfo.newBuilder()
val decimal = DecimalInfo.newBuilder()
decimal.setPrecision(t.precision)
decimal.setScale(t.scale)
info.setDecimal(decimal)
builder.setTypeInfo(info.build()).build()
case a: ArrayType =>
val elementType = serializeDataType(a.elementType)
if (elementType.isEmpty) {
return None
}
val info = DataTypeInfo.newBuilder()
val list = ListInfo.newBuilder()
list.setElementType(elementType.get)
list.setContainsNull(a.containsNull)
info.setList(list)
builder.setTypeInfo(info.build()).build()
case m: MapType =>
val keyType = serializeDataType(m.keyType)
if (keyType.isEmpty) {
return None
}
val valueType = serializeDataType(m.valueType)
if (valueType.isEmpty) {
return None
}
val info = DataTypeInfo.newBuilder()
val map = MapInfo.newBuilder()
map.setKeyType(keyType.get)
map.setValueType(valueType.get)
map.setValueContainsNull(m.valueContainsNull)
info.setMap(map)
builder.setTypeInfo(info.build()).build()
case s: StructType =>
val info = DataTypeInfo.newBuilder()
val struct = StructInfo.newBuilder()
val fieldNames = s.fields.map(_.name).toIterable.asJava
val fieldDatatypes = s.fields.map(f => serializeDataType(f.dataType)).toSeq
val fieldNullable = s.fields.map(f => Boolean.box(f.nullable)).toIterable.asJava
if (fieldDatatypes.exists(_.isEmpty)) {
return None
}
struct.addAllFieldNames(fieldNames)
struct.addAllFieldDatatypes(fieldDatatypes.map(_.get).asJava)
struct.addAllFieldNullable(fieldNullable)
info.setStruct(struct)
builder.setTypeInfo(info.build()).build()
case _ => builder.build()
}
Some(dataType)
}