TaraReferenceExtractionUtils { public static Dataset<Row> runReferenceExtraction(SparkSession spark, Dataset<Row> documentMetadataDF, PipeExecutionEnvironment environment) throws IOException { logger.info("Running reference extraction for input document metadata."); String pipeCommandStr = environment.pipeCommand(); JavaRDD<Row> piped = documentMetadataDF .toJSON() .javaRDD() .pipe(pipeCommandStr) .map(RowFactory::create); StructType schema = (StructType) SchemaConverters.toSqlType(DocumentToProject.SCHEMA$).dataType(); Dataset<Row> resultAsNullableDF = spark.createDataFrame(piped, PIPE_RESULT_SCHEMA) .withColumn("json_struct", from_json(col("value"), schema.asNullable())) .select(expr("json_struct.*")); return spark.createDataFrame(resultAsNullableDF.javaRDD(), schema); }  static Dataset<Row> buildDocumentMetadata(Dataset<Row> documentTextDF,
                                                     Dataset<Row> extractedDocumentMetadataMergedWithOriginalDF); static Dataset<Row> runReferenceExtraction(SparkSession spark,
                                                      Dataset<Row> documentMetadataDF,
                                                      PipeExecutionEnvironment environment);  }
@Test public void shouldRunReferenceExtraction() throws IOException { Dataset<Row> documentMetadataDF = spark.createDataFrame( Collections.singletonList( createDocumentMetadata("id-1", "text-1") ), TaraReferenceExtractionJob.DOCUMENT_METADATA_SCHEMA); PipeExecutionEnvironment pipeExecutionEnvironment = () -> { Path scriptWithInputCheck = createTestScriptWithInputCheck(); spark.sparkContext().addFile(scriptWithInputCheck.toString()); return String.format("bash %s/%s", SparkFiles.getRootDirectory(), scriptWithInputCheck.getFileName().toString()); }; Dataset<Row> resultDF = runReferenceExtraction(spark, documentMetadataDF, pipeExecutionEnvironment); assertEquals(SchemaConverters.toSqlType(DocumentToProject.SCHEMA$).dataType().asNullable(), resultDF.schema().asNullable()); List<Row> results = resultDF.collectAsList(); assertEquals(1, results.size()); Row row = results.get(0); assertForDocumentToProject(row, "docId-1", "projId-1", 1.0f); }