Java Code Examples for org.apache.spark.sql.Dataset#collectAsList()

private void assertCorrectness(Dataset<Row> rowDataset, Transformer transformer) {
	List<Row> sparkOutput = rowDataset.collectAsList();

	for (Row row : sparkOutput) {
		Map<String, Object> data = new HashMap<>();
		data.put("mergedAddress", row.get(0));

		List<Object> list = row.getList(1);
		String[] sanitizedAddress = new String[list.size()];
		for (int j = 0; j < sanitizedAddress.length; j++) {
			sanitizedAddress[j] = (String) list.get(j);
		}

		data.put("sanitizedAddress", sanitizedAddress);
		transformer.transform(data);

		assertEquals("number of words should be equals", row.get(2), data.get("numWords"));
		assertEquals("number of commas should be equals", row.get(3), data.get("numCommas"));
		assertEquals("numericPresent should be equals", row.get(4), data.get("numericPresent"));
		assertEquals("addressLength should be equals", row.get(5), data.get("addressLength"));
		assertEquals("favouredStart should be equals", row.get(6), data.get("favouredStart"));
		assertEquals("unfavouredStart should be equals", row.get(7), data.get("unfavouredStart"));
	}
}
 

@Test
public void testStringDataFrameToVectorDataFrameNull() {
	List<String> list = new ArrayList<>();
	list.add("[1.2, 3.4]");
	list.add(null);
	JavaRDD<String> javaRddString = sc.parallelize(list);
	JavaRDD<Row> javaRddRow = javaRddString.map(new StringToRow());
	SparkSession sparkSession = SparkSession.builder().sparkContext(sc.sc()).getOrCreate();
	List<StructField> fields = new ArrayList<>();
	fields.add(DataTypes.createStructField("C1", DataTypes.StringType, true));
	StructType schema = DataTypes.createStructType(fields);
	Dataset<Row> inDF = sparkSession.createDataFrame(javaRddRow, schema);
	Dataset<Row> outDF = RDDConverterUtilsExt.stringDataFrameToVectorDataFrame(sparkSession, inDF);

	List<String> expectedResults = new ArrayList<>();
	expectedResults.add("[[1.2,3.4]]");
	expectedResults.add("[null]");

	List<Row> outputList = outDF.collectAsList();
	for (Row row : outputList) {
		assertTrue("Expected results don't contain: " + row, expectedResults.contains(row.toString()));
	}
}
 

@Test
public void testOutputDataFrameFromMatrixDML() {
	System.out.println("MLContextTest - output DataFrame from matrix DML");

	String s = "M = matrix('1 2 3 4', rows=2, cols=2);";
	Script script = dml(s).out("M");
	Dataset<Row> df = ml.execute(script).getMatrix("M").toDF();
	Dataset<Row> sortedDF = df.sort(RDDConverterUtils.DF_ID_COLUMN);
	List<Row> list = sortedDF.collectAsList();
	Row row1 = list.get(0);
	Assert.assertEquals(1.0, row1.getDouble(0), 0.0);
	Assert.assertEquals(1.0, row1.getDouble(1), 0.0);
	Assert.assertEquals(2.0, row1.getDouble(2), 0.0);

	Row row2 = list.get(1);
	Assert.assertEquals(2.0, row2.getDouble(0), 0.0);
	Assert.assertEquals(3.0, row2.getDouble(1), 0.0);
	Assert.assertEquals(4.0, row2.getDouble(2), 0.0);
}
 

@Test
public void testOutputDataFrameDMLDoublesNoIDColumn() {
	System.out.println("MLContextTest - output DataFrame DML, doubles no ID column");

	String s = "M = matrix('1 2 3 4', rows=2, cols=2);";
	Script script = dml(s).out("M");
	MLResults results = ml.execute(script);
	Dataset<Row> dataFrame = results.getDataFrameDoubleNoIDColumn("M");
	List<Row> list = dataFrame.collectAsList();

	Row row1 = list.get(0);
	Assert.assertEquals(1.0, row1.getDouble(0), 0.0);
	Assert.assertEquals(2.0, row1.getDouble(1), 0.0);

	Row row2 = list.get(1);
	Assert.assertEquals(3.0, row2.getDouble(0), 0.0);
	Assert.assertEquals(4.0, row2.getDouble(1), 0.0);
}
 

@Test(expected = IllegalArgumentException.class)
public void testSnapshotSelectionByInvalidTimestamp() throws IOException {
  long timestamp = System.currentTimeMillis();

  String tableLocation = temp.newFolder("iceberg-table").toString();
  HadoopTables tables = new HadoopTables(CONF);
  PartitionSpec spec = PartitionSpec.unpartitioned();
  tables.create(SCHEMA, spec, tableLocation);

  Dataset<Row> df = spark.read()
      .format("iceberg")
      .option("as-of-timestamp", timestamp)
      .load(tableLocation);

  df.collectAsList();
}
 

@Test
public void testStringDataFrameToVectorDataFrameNull() {
	List<String> list = new ArrayList<>();
	list.add("[1.2, 3.4]");
	list.add(null);
	JavaRDD<String> javaRddString = sc.parallelize(list);
	JavaRDD<Row> javaRddRow = javaRddString.map(new StringToRow());
	SparkSession sparkSession = SparkSession.builder().sparkContext(sc.sc()).getOrCreate();
	List<StructField> fields = new ArrayList<>();
	fields.add(DataTypes.createStructField("C1", DataTypes.StringType, true));
	StructType schema = DataTypes.createStructType(fields);
	Dataset<Row> inDF = sparkSession.createDataFrame(javaRddRow, schema);
	Dataset<Row> outDF = RDDConverterUtilsExt.stringDataFrameToVectorDataFrame(sparkSession, inDF);

	List<String> expectedResults = new ArrayList<>();
	expectedResults.add("[[1.2,3.4]]");
	expectedResults.add("[null]");

	List<Row> outputList = outDF.collectAsList();
	for (Row row : outputList) {
		assertTrue("Expected results don't contain: " + row, expectedResults.contains(row.toString()));
	}
}
 

public Map<String, List<Row>> collectAll() {
    Map<String, List<Row>> rowMap = new HashMap<>();
    DataFrameLoader.loadDataFrame(collectionName, javaSparkContext);
    Dataset<Row> dataRows = null;
    if (collectionName.contains("pipelines")) {
        dataRows = getDataRowsForPipelines();
    }
    else {
        dataRows = sparkSession.sql(query);
    }
    List<Row> rowList = dataRows.collectAsList();
        rowList.forEach(row -> {
            String item = (String) ((GenericRowWithSchema) row.getAs("collectorItemId")).get(0);
            boolean matchingCollectorItem = !Objects.equals(
                    Optional.ofNullable(collectorItemIds).orElseGet(Collections::emptyList).stream()
                    .filter(Predicate.isEqual(item)).findFirst().orElse(""), "");
            Date timeWindowDt = row.getAs("timeWindow");
            long daysAgo = HygieiaExecutiveUtil.getDaysAgo(timeWindowDt);
            if (matchingCollectorItem && (daysAgo < 90)) {
                String collectorItemId = (String) ((GenericRowWithSchema) row.getAs("collectorItemId")).get(0);
                List<Row> existingRowList = rowMap.get(collectorItemId);
                if (existingRowList == null) {
                    List<Row> newRow = new ArrayList<>();
                    newRow.add(row);
                    rowMap.put(collectorItemId, newRow);
                } else {
                    existingRowList.add(row);
                    rowMap.put(collectorItemId, existingRowList);
                }
            }
        });
    return rowMap;
}
 

@Test
public void testOutputDataFrameVectorsWithIDColumnFromMatrixDML() {
	System.out.println("MLContextTest - output DataFrame of vectors with ID column from matrix DML");

	String s = "M = matrix('1 2 3 4', rows=1, cols=4);";
	Script script = dml(s).out("M");
	Dataset<Row> df = ml.execute(script).getMatrix("M").toDFVectorWithIDColumn();
	List<Row> list = df.collectAsList();

	Row row = list.get(0);
	Assert.assertEquals(1.0, row.getDouble(0), 0.0);
	Assert.assertArrayEquals(new double[] { 1.0, 2.0, 3.0, 4.0 }, ((Vector) row.get(1)).toArray(), 0.0);
}
 

public static void main(String[] args) {
	System.setProperty("hadoop.home.dir", "C:\\softwares\\Winutils");
	SparkSession sparkSession = SparkSession.builder()
			.master("local")
			.appName("CSV Read Example")
			.config("spark.sql.warehouse.dir", "file:////C:/Users/sgulati/spark-warehouse")
			.getOrCreate();
	
	Dataset<Row> csv = sparkSession.read().format("com.databricks.spark.csv").option("header","true")
			.load("C:\\Users\\sgulati\\Documents\\my_docs\\book\\testdata\\emp.csv");
	
	csv.createOrReplaceTempView("test");
	Dataset<Row> sql = sparkSession.sql("select * from test");
	sql.collectAsList();
}
 

public static void main(String[] args) {
  SparkSession spark = SparkSession
    .builder()
    .appName("JavaWord2VecExample")
    .getOrCreate();

  // $example on$
  // Input data: Each row is a bag of words from a sentence or document.
  List<Row> data = Arrays.asList(
    RowFactory.create(Arrays.asList("Hi I heard about Spark".split(" "))),
    RowFactory.create(Arrays.asList("I wish Java could use case classes".split(" "))),
    RowFactory.create(Arrays.asList("Logistic regression models are neat".split(" ")))
  );
  StructType schema = new StructType(new StructField[]{
    new StructField("text", new ArrayType(DataTypes.StringType, true), false, Metadata.empty())
  });
  Dataset<Row> documentDF = spark.createDataFrame(data, schema);

  // Learn a mapping from words to Vectors.
  Word2Vec word2Vec = new Word2Vec()
    .setInputCol("text")
    .setOutputCol("result")
    .setVectorSize(3)
    .setMinCount(0);

  Word2VecModel model = word2Vec.fit(documentDF);
  Dataset<Row> result = model.transform(documentDF);

  for (Row row : result.collectAsList()) {
    List<String> text = row.getList(0);
    Vector vector = (Vector) row.get(1);
    System.out.println("Text: " + text + " => \nVector: " + vector + "\n");
  }
  // $example off$

  spark.stop();
}
 

@Test
public void testOutputDataFrameVectorsWithIDColumnFromMatrixDML() {
	System.out.println("MLContextTest - output DataFrame of vectors with ID column from matrix DML");

	String s = "M = matrix('1 2 3 4', rows=1, cols=4);";
	Script script = dml(s).out("M");
	Dataset<Row> df = ml.execute(script).getMatrix("M").toDFVectorWithIDColumn();
	List<Row> list = df.collectAsList();

	Row row = list.get(0);
	Assert.assertEquals(1.0, row.getDouble(0), 0.0);
	Assert.assertArrayEquals(new double[] { 1.0, 2.0, 3.0, 4.0 }, ((Vector) row.get(1)).toArray(), 0.0);
}
 

@Test
public void testRowLevelRules() throws Exception {
  Config config = ConfigUtils.configFromResource("/dq/dq-dataset-good.conf").getConfig("steps.checkrows");

  SparkSession sparkSession = Contexts.getSparkSession();

  List<Row> dataList = Lists.newArrayList(
      new RowWithSchema(SCHEMA, "Apple", "One Infinite Loop", 151),
      (Row)new RowWithSchema(SCHEMA, "Microsoft", "One Microsoft Way", -1)
  );
  Dataset<Row> mydata = sparkSession.createDataFrame(dataList, SCHEMA);

  Map<String, Dataset<Row>> dependencies = Maps.newHashMap();
  dependencies.put("mydata", mydata);

  DataQualityDeriver dq = new DataQualityDeriver();
  assertNoValidationFailures(dq, config.getConfig("deriver"));
  dq.configure(config.getConfig("deriver"));
  Dataset<Row> dqResults = dq.derive(dependencies);
  List<Row> dqRows = dqResults.collectAsList();

  assertEquals("Should be two rows", 2, dqRows.size());
  for (Row row : dqRows) {
    scala.collection.immutable.Map<String, Boolean> scalaResults = row.getAs("results");
    Map<String, Boolean> ruleResults = fromScalaMap(scalaResults);
    assertEquals("Rule results map should have three entries", 3, ruleResults.size());
    assertTrue("Checkfields should pass", ruleResults.get("r1"));
    assertTrue("Regex should pass", ruleResults.get("r2"));
    assertFalse("Range should fail", ruleResults.get("r3"));
  }
}
 

private void writeAndValidateWithLocations(Table table, File location, File expectedDataDir) throws IOException {
  Schema tableSchema = table.schema(); // use the table schema because ids are reassigned

  table.updateProperties().set(TableProperties.DEFAULT_FILE_FORMAT, format).commit();

  List<Record> expected = RandomData.generateList(tableSchema, 100, 0L);
  Dataset<Row> df = createDataset(expected, tableSchema);
  DataFrameWriter<?> writer = df.write().format("iceberg").mode("append");

  writer.save(location.toString());

  table.refresh();

  Dataset<Row> result = spark.read()
      .format("iceberg")
      .load(location.toString());

  List<Row> actual = result.collectAsList();

  Assert.assertEquals("Result size should match expected", expected.size(), actual.size());
  for (int i = 0; i < expected.size(); i += 1) {
    assertEqualsSafe(tableSchema.asStruct(), expected.get(i), actual.get(i));
  }

  table.currentSnapshot().addedFiles().forEach(dataFile ->
      Assert.assertTrue(
          String.format(
              "File should have the parent directory %s, but has: %s.",
              expectedDataDir.getAbsolutePath(),
              dataFile.path()),
          URI.create(dataFile.path().toString()).getPath().startsWith(expectedDataDir.getAbsolutePath())));
}
 

public static void printDatasetContent(Dataset<WindowedValue> dataset) {
  // cannot use dataset.show because dataset schema is binary so it will print binary
  // code.
  List<WindowedValue> windowedValues = dataset.collectAsList();
  for (WindowedValue windowedValue : windowedValues) {
    LOG.debug("**** dataset content {} ****", windowedValue.toString());
  }
}
 

@Test
public void testStringDataFrameToVectorDataFrame() {
	List<String> list = new ArrayList<>();
	list.add("((1.2, 4.3, 3.4))");
	list.add("(1.2, 3.4, 2.2)");
	list.add("[[1.2, 34.3, 1.2, 1.25]]");
	list.add("[1.2, 3.4]");
	JavaRDD<String> javaRddString = sc.parallelize(list);
	JavaRDD<Row> javaRddRow = javaRddString.map(new StringToRow());
	SparkSession sparkSession = SparkSession.builder().sparkContext(sc.sc()).getOrCreate();
	List<StructField> fields = new ArrayList<>();
	fields.add(DataTypes.createStructField("C1", DataTypes.StringType, true));
	StructType schema = DataTypes.createStructType(fields);
	Dataset<Row> inDF = sparkSession.createDataFrame(javaRddRow, schema);
	Dataset<Row> outDF = RDDConverterUtilsExt.stringDataFrameToVectorDataFrame(sparkSession, inDF);

	List<String> expectedResults = new ArrayList<>();
	expectedResults.add("[[1.2,4.3,3.4]]");
	expectedResults.add("[[1.2,3.4,2.2]]");
	expectedResults.add("[[1.2,34.3,1.2,1.25]]");
	expectedResults.add("[[1.2,3.4]]");

	List<Row> outputList = outDF.collectAsList();
	for (Row row : outputList) {
		assertTrue("Expected results don't contain: " + row, expectedResults.contains(row.toString()));
	}
}
 

/**
 * Returns a bundle containing all resources in the dataset. This should be used
 * with caution for large datasets, since the returned bundle will include all data.
 *
 * @param dataset a dataset of FHIR resoruces
 * @param resourceTypeUrl the FHIR resource type
 * @return a bundle containing those resources.
 */
public static Bundle toBundle(Dataset<Row> dataset,
    String resourceTypeUrl) {

  List<Row> resources = (List<Row>) dataset.collectAsList();

  SparkRowConverter converter = SparkRowConverter.forResource(CONTEXT,resourceTypeUrl);

  Bundle bundle = new Bundle();

  for (Row row : resources) {

    IBaseResource resource = converter.rowToResource(row);

    bundle.addEntry().setResource((Resource) resource);
  }

  return bundle;
}
 

private static List<Row> read(String table, String expr, String select0, String... selectN) {
  Dataset<Row> dataset = spark.read().format("iceberg").load(table).filter(expr)
      .select(select0, selectN);
  return dataset.collectAsList();
}
 

public static void main(String[] args) {
  SparkSession spark = SparkSession
    .builder()
    .appName("JavaEstimatorTransformerParamExample")
    .getOrCreate();

  // $example on$
  // Prepare training data.
  List<Row> dataTraining = Arrays.asList(
      RowFactory.create(1.0, Vectors.dense(0.0, 1.1, 0.1)),
      RowFactory.create(0.0, Vectors.dense(2.0, 1.0, -1.0)),
      RowFactory.create(0.0, Vectors.dense(2.0, 1.3, 1.0)),
      RowFactory.create(1.0, Vectors.dense(0.0, 1.2, -0.5))
  );
  StructType schema = new StructType(new StructField[]{
      new StructField("label", DataTypes.DoubleType, false, Metadata.empty()),
      new StructField("features", new VectorUDT(), false, Metadata.empty())
  });
  Dataset<Row> training = spark.createDataFrame(dataTraining, schema);

  // Create a LogisticRegression instance. This instance is an Estimator.
  LogisticRegression lr = new LogisticRegression();
  // Print out the parameters, documentation, and any default values.
  System.out.println("LogisticRegression parameters:\n" + lr.explainParams() + "\n");

  // We may set parameters using setter methods.
  lr.setMaxIter(10).setRegParam(0.01);

  // Learn a LogisticRegression model. This uses the parameters stored in lr.
  LogisticRegressionModel model1 = lr.fit(training);
  // Since model1 is a Model (i.e., a Transformer produced by an Estimator),
  // we can view the parameters it used during fit().
  // This prints the parameter (name: value) pairs, where names are unique IDs for this
  // LogisticRegression instance.
  System.out.println("Model 1 was fit using parameters: " + model1.parent().extractParamMap());

  // We may alternatively specify parameters using a ParamMap.
  ParamMap paramMap = new ParamMap()
    .put(lr.maxIter().w(20))  // Specify 1 Param.
    .put(lr.maxIter(), 30)  // This overwrites the original maxIter.
    .put(lr.regParam().w(0.1), lr.threshold().w(0.55));  // Specify multiple Params.

  // One can also combine ParamMaps.
  ParamMap paramMap2 = new ParamMap()
    .put(lr.probabilityCol().w("myProbability"));  // Change output column name
  ParamMap paramMapCombined = paramMap.$plus$plus(paramMap2);

  // Now learn a new model using the paramMapCombined parameters.
  // paramMapCombined overrides all parameters set earlier via lr.set* methods.
  LogisticRegressionModel model2 = lr.fit(training, paramMapCombined);
  System.out.println("Model 2 was fit using parameters: " + model2.parent().extractParamMap());

  // Prepare test documents.
  List<Row> dataTest = Arrays.asList(
      RowFactory.create(1.0, Vectors.dense(-1.0, 1.5, 1.3)),
      RowFactory.create(0.0, Vectors.dense(3.0, 2.0, -0.1)),
      RowFactory.create(1.0, Vectors.dense(0.0, 2.2, -1.5))
  );
  Dataset<Row> test = spark.createDataFrame(dataTest, schema);

  // Make predictions on test documents using the Transformer.transform() method.
  // LogisticRegression.transform will only use the 'features' column.
  // Note that model2.transform() outputs a 'myProbability' column instead of the usual
  // 'probability' column since we renamed the lr.probabilityCol parameter previously.
  Dataset<Row> results = model2.transform(test);
  Dataset<Row> rows = results.select("features", "label", "myProbability", "prediction");
  for (Row r: rows.collectAsList()) {
    System.out.println("(" + r.get(0) + ", " + r.get(1) + ") -> prob=" + r.get(2)
      + ", prediction=" + r.get(3));
  }
  // $example off$

  spark.stop();
}
 

@Override
protected void writeAndValidate(Schema schema) throws IOException {
  System.out.println("Starting ORC test with " + schema);
  final int ROW_COUNT = 100;
  final long SEED = 1;
  File parent = temp.newFolder("orc");
  File location = new File(parent, "test");
  File dataFolder = new File(location, "data");
  dataFolder.mkdirs();

  File orcFile = new File(dataFolder,
      FileFormat.ORC.addExtension(UUID.randomUUID().toString()));

  HadoopTables tables = new HadoopTables(CONF);
  Table table = tables.create(schema, PartitionSpec.unpartitioned(),
      location.toString());

  // Important: use the table's schema for the rest of the test
  // When tables are created, the column ids are reassigned.
  Schema tableSchema = table.schema();

  Metrics metrics;
  SparkOrcWriter writer = new SparkOrcWriter(ORC.write(localOutput(orcFile))
      .schema(tableSchema)
      .build());
  try {
    writer.addAll(RandomData.generateSpark(tableSchema, ROW_COUNT, SEED));
  } finally {
    writer.close();
    // close writes the last batch, so metrics are not correct until after close is called
    metrics = writer.metrics();
  }

  DataFile file = DataFiles.builder(PartitionSpec.unpartitioned())
      .withFileSizeInBytes(orcFile.length())
      .withPath(orcFile.toString())
      .withMetrics(metrics)
      .build();

  table.newAppend().appendFile(file).commit();

  Dataset<Row> df = spark.read()
      .format("iceberg")
      .load(location.toString());

  List<Row> rows = df.collectAsList();
  Assert.assertEquals("Wrong number of rows", ROW_COUNT, rows.size());
  Iterator<InternalRow> expected = RandomData.generateSpark(tableSchema,
      ROW_COUNT, SEED);
  for(int i=0; i < ROW_COUNT; ++i) {
    TestHelpers.assertEquals("row " + i, schema.asStruct(), expected.next(),
        rows.get(i));
  }
}
 

private static List<Row> read(String table, String expr, String select0, String... selectN) {
  Dataset<Row> dataset = spark.read().format("iceberg").load(table).filter(expr)
      .select(select0, selectN);
  return dataset.collectAsList();
}