Apache Spark SQL StructType 和 UDF

Question

Spark 1.6 / Java-7

Initial dataframe with a new column

# adding new column for the UDF computation:
df = df.withColumn("TEMP_COLUMN", lit(null));

What is the correct format for the UDF function to create new StructType and put it into the cell?

public static DataFrame compute(SQLContext sqlContext, DataFrame df) {
    sqlContext.udf().register("compute", new MyUdf(), new ArrayType(new StructType(), true));
    return df.withColumn("TEMP_COLUMN", functions.callUDF("compute"));
}

class MyUdf implements UDF0<List<StructType>> {
@Override
public  List<StructType> call() {
    ...
    return ? // what must be returned here? List<StructType> or List<String> or anything else?
}


+-------------------------+
|TEMP_COLUMN              |
+-------------------------+
|[A[1, 2, 3], B[4, 5, 6]] |
+-------------------------+

I want to have a structure with the array of elements with several fields for each element.
I don't understand, is registration with the type new ArrayType(new StructType(), true) correct and the same for the return type of the UDF function List<StructType> .
How is the data should be returned? Is it like new StructType(new StructField[]{new StructField(...)) ?

Answer 1

Answering my own question since we were lucky to find out how to do it:

Let's say that we have a 'complex' structure for our needs:

MapType CLIENTS_INFO_DATA_TYPE = DataTypes.createMapType(
  DataTypes.StringType,
  DataTypes.createStructType(
    new StructField[] {
        DataTypes.createStructField("NAME_1", DataTypes.DoubleType, false),
        DataTypes.createStructField("NAME_2", DataTypes.DoubleType, false),
        DataTypes.createStructField("NAME_3", DataTypes.DoubleType, false)
  ),
  true
);


StructType COMPLEX_DATA_TYPE = DataTypes.createStructType(new StructField[] {
  DataTypes.createStructField("CLIENTS_INFO", CLIENTS_INFO_DATA_TYPE, true),
  DataTypes.createStructField("COMMENT", DataTypes.StringType, true)
}

And it's schema:

dataFrame.printSchema()

|-- COMPLEX_DATA_TYPE: struct (nullable = true)
|    |-- CLIENTS_INFO: map (nullable = true)
|    |    |-- key: string
|    |    |-- value: struct (valueContainsNull = true)
|    |    |    |-- NAME_1: double (nullable = false)
|    |    |    |-- NAME_2: double (nullable = false)
|    |    |    |-- NAME_3: double (nullable = false)
|    |-- COMMENT: string (nullable = true)

Next we have to register the UDF function that operates with our structure:

DataFrame compute(SQLContext sqlContext, DataFrame df) {
sqlContext.udf().register(
        "computeUDF",
        new MyUDF(),
        COMPLEX_DATA_TYPE);

  return df.withColumn("TEMP_FIELD_NAME", functions.callUDF("computeUDF", field_1.getColumn(), field_2.getColumn()));
}

And the final step is the UDF function by itself that returns a Row object (that will be converted into our structure):

public final class MyUDF implements UDF2<Double, Double, Row> {
@Override
public Row call(Double value1, Double value2) {
    Map<String, Row> clientsInfoMap = new HashMap<>();
    ...
    for (Map.Entry<String, ClientInfo> clientInfoEntry : clientsInfo.entrySet()) {
        final String client = clientInfoEntry.getKey();
        final ClientInfo clientInfo = clientInfoEntry.getValue();

        final Double[] clientInfoValues = {10.0, 20.0, 30.0};
        
        Row clientInfoRow = new GenericRow(clientInfoValues);
        clientsInfoMap.put(client, clientInfoRow);
    }

    Object[] fullClientsInfo = new Object[] {clientsInfoMap, "string-as-a-comment"};
    return new GenericRow(fullClientsInfo);
  }
}

And now, since it is a structure, we can select by using the TEMP_FIELD_NAME.CLIENTS_INFO and anything else by the namings.