void validateDataFrameWithBeans(Bean bean, DataFrame df) {
StructType schema = df.schema();
Assert.assertEquals(
new StructField("a", DoubleType$.MODULE$, false, Metadata.empty()), schema.apply("a"));
Assert.assertEquals(
new StructField("b", new ArrayType(IntegerType$.MODULE$, true), true, Metadata.empty()),
schema.apply("b"));
ArrayType valueType = new ArrayType(DataTypes.IntegerType, false);
MapType mapType = new MapType(DataTypes.StringType, valueType, true);
Assert.assertEquals(new StructField("c", mapType, true, Metadata.empty()), schema.apply("c"));
Assert.assertEquals(
new StructField("d", new ArrayType(DataTypes.StringType, true), true, Metadata.empty()),
schema.apply("d"));
Row first = df.select("a", "b", "c", "d").first();
Assert.assertEquals(bean.getA(), first.getDouble(0), 0.0);
// Now Java lists and maps are converted to Scala Seq's and Map's. Once we get a Seq below,
// verify that it has the expected length, and contains expected elements.
Seq<Integer> result = first.getAs(1);
Assert.assertEquals(bean.getB().length, result.length());
for (int i = 0; i < result.length(); i++) {
Assert.assertEquals(bean.getB()[i], result.apply(i));
}
@SuppressWarnings("unchecked")
Seq<Integer> outputBuffer = (Seq<Integer>) first.getJavaMap(2).get("hello");
Assert.assertArrayEquals(
bean.getC().get("hello"),
Ints.toArray(JavaConverters.seqAsJavaListConverter(outputBuffer).asJava()));
Seq<String> d = first.getAs(3);
Assert.assertEquals(bean.getD().size(), d.length());
for (int i = 0; i < d.length(); i++) {
Assert.assertEquals(bean.getD().get(i), d.apply(i));
}
}
@Test
public void testCreateStructTypeFromList() {
List<StructField> fields1 = new ArrayList<>();
fields1.add(new StructField("id", DataTypes.StringType, true, Metadata.empty()));
StructType schema1 = StructType$.MODULE$.apply(fields1);
Assert.assertEquals(0, schema1.fieldIndex("id"));
List<StructField> fields2 =
Arrays.asList(new StructField("id", DataTypes.StringType, true, Metadata.empty()));
StructType schema2 = StructType$.MODULE$.apply(fields2);
Assert.assertEquals(0, schema2.fieldIndex("id"));
}
public JavaRDD<Row> toRows(StructType schema, JavaRDD<SolrDocument> docs) {
final String[] queryFields = schema.fieldNames();
JavaRDD<Row> rows =
docs.map(
new Function<SolrDocument, Row>() {
public Row call(SolrDocument doc) throws Exception {
List<Object> vals = new ArrayList<Object>(queryFields.length);
for (String field : queryFields) {
vals.add(doc.getFirstValue(field));
}
return RowFactory.create(vals.toArray());
}
});
return rows;
}