Source:
<code>public DataSet<ExecRow> writeAvroFile(DataSetProcessor dsp, int[] partitionBy, String location, String compression, OperationContext context) { try { StructType dataSchema = null; StructType tableSchema = context.getOperation().getExecRowDefinition().schema(); dataSchema = ExternalTableUtils.getDataSchema(dsp, tableSchema, partitionBy, location, "a"); if (dataSchema == null) dataSchema = tableSchema; Dataset<Row> insertDF = SpliceSpark.getSession().createDataFrame( rdd.map(new SparkSpliceFunctionWrapper<>(new CountWriteFunction(context))).map(new LocatedRowToRowAvroFunction()), dataSchema); List<String> partitionByCols = new ArrayList(); for (int i = 0; i < partitionBy.length; i++) { <START> partitionByCols.add(ValueRow.getNamedColumn(partitionBy[i])); <END> } if (partitionBy.length > 0) { List<Column> repartitionCols = new ArrayList(); for (int i = 0; i < partitionBy.length; i++) { repartitionCols.add(new Column(ValueRow.getNamedColumn(partitionBy[i]))); } insertDF = insertDF.repartition(scala.collection.JavaConversions.asScalaBuffer(repartitionCols).toList()); } insertDF.write().option(SPARK_COMPRESSION_OPTION,compression).partitionBy(partitionByCols.toArray(new String[partitionByCols.size()])) .mode(SaveMode.Append).format("com.databricks.spark.avro").save(location); ValueRow valueRow=new ValueRow(1); valueRow.setColumn(1,new SQLLongint(context.getRecordsWritten())); return new SparkDataSet<>(SpliceSpark.getContext().parallelize(Collections.singletonList(valueRow), 1)); } catch (Exception e) { throw new RuntimeException(e); } }</code><technical_language>I need handle partitioned columns in Orc Avro</technical_language>

Target:
public DataSet<ExecRow> writeAvroFile(DataSetProcessor dsp, int[] partitionBy, String location, String compression, OperationContext context) { try { StructType dataSchema = null; StructType tableSchema = context.getOperation().getExecRowDefinition().schema(); dataSchema = ExternalTableUtils.getDataSchema(dsp, tableSchema, partitionBy, location, "a"); if (dataSchema == null) dataSchema = tableSchema; Dataset<Row> insertDF = SpliceSpark.getSession().createDataFrame( rdd.map(new SparkSpliceFunctionWrapper<>(new CountWriteFunction(context))).map(new LocatedRowToRowAvroFunction()), dataSchema); List<String> partitionByCols = new ArrayList(); for (int i = 0; i < partitionBy.length; i++) { partitionByCols.add(dataSchema.fields()[partitionBy[i]].name()); } if (partitionBy.length > 0) { List<Column> repartitionCols = new ArrayList(); for (int i = 0; i < partitionBy.length; i++) { repartitionCols.add(new Column(dataSchema.fields()[partitionBy[i]].name())); } insertDF = insertDF.repartition(scala.collection.JavaConversions.asScalaBuffer(repartitionCols).toList()); } insertDF.write().option(SPARK_COMPRESSION_OPTION,compression).partitionBy(partitionByCols.toArray(new String[partitionByCols.size()])) .mode(SaveMode.Append).format("com.databricks.spark.avro").save(location); ValueRow valueRow=new ValueRow(1); valueRow.setColumn(1,new SQLLongint(context.getRecordsWritten())); return new SparkDataSet<>(SpliceSpark.getContext().parallelize(Collections.singletonList(valueRow), 1)); } catch (Exception e) { throw new RuntimeException(e); } }