in client-spark/spark-3/src/main/scala/org/apache/spark/sql/execution/columnar/CelebornColumnarBatchCodeGenBuild.scala [94:293]
def genCode(schema: StructType, batchSize: Int): (
mutable.StringBuilder,
mutable.StringBuilder,
mutable.StringBuilder,
mutable.StringBuilder) = {
val initCode = new mutable.StringBuilder()
val buildCode = new mutable.StringBuilder()
val writeCode = new mutable.StringBuilder()
val writeRowCode = new mutable.StringBuilder()
for (index <- schema.indices) {
schema.fields(index).dataType match {
case NullType =>
initCode.append(
s"""
| ${classOf[CelebornNullColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornNullColumnBuilder].getName}();
| builder.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case ByteType =>
initCode.append(
s"""
| ${classOf[CelebornByteCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornByteCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case BooleanType =>
initCode.append(
s"""
| ${classOf[CelebornBooleanCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornBooleanCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case ShortType =>
initCode.append(
s"""
| ${classOf[CelebornShortCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornShortCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case IntegerType =>
initCode.append(
s"""
| ${classOf[CelebornIntCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornIntCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case LongType =>
initCode.append(
s"""
| ${classOf[CelebornLongCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornLongCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case FloatType =>
initCode.append(
s"""
| ${classOf[CelebornFloatCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornFloatCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case DoubleType =>
initCode.append(
s"""
| ${classOf[CelebornDoubleCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornDoubleCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case StringType =>
initCode.append(
s"""
| ${classOf[CelebornStringCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornStringCodeGenColumnBuilder].getName}();
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case dt: DecimalType if dt.precision <= Decimal.MAX_INT_DIGITS =>
initCode.append(
s"""
| ${classOf[CelebornCompactMiniDecimalCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index =
| new ${classOf[CelebornCompactMiniDecimalCodeGenColumnBuilder].getName}(
| new ${classOf[DecimalType].getName}(${dt.precision}, ${dt.scale}));
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case dt: DecimalType if dt.precision <= Decimal.MAX_LONG_DIGITS =>
initCode.append(
s"""
| ${classOf[CelebornCompactDecimalCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index =
| new ${classOf[CelebornCompactDecimalCodeGenColumnBuilder].getName}
| (new ${classOf[DecimalType].getName}(${dt.precision}, ${dt.scale}));
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
case dt: DecimalType =>
initCode.append(
s"""
| ${classOf[CelebornDecimalCodeGenColumnBuilder].getName} b$index;
""".stripMargin)
buildCode.append(
s"""
| b$index = new ${classOf[CelebornDecimalCodeGenColumnBuilder].getName}
| (new ${classOf[DecimalType].getName}(${dt.precision}, ${dt.scale}));
| b$index.initialize($batchSize, "${schema.fields(index).name}", false);
""".stripMargin)
writeCode.append(genWriteCode(index))
writeRowCode.append(
s"""
| b$index.appendFrom(row, $index);
""".stripMargin)
}
}
(initCode, buildCode, writeCode, writeRowCode)
}