forked from D-Net/dnet-hadoop
645 lines
23 KiB
Scala
645 lines
23 KiB
Scala
package eu.dnetlib.pace.model
|
|
|
|
import com.jayway.jsonpath.{Configuration, JsonPath, Option}
|
|
import eu.dnetlib.pace.config.{DedupConfig, Type}
|
|
import eu.dnetlib.pace.tree.support.TreeProcessor
|
|
import eu.dnetlib.pace.util.MapDocumentUtil.truncateValue
|
|
import eu.dnetlib.pace.util.{BlockProcessor, MapDocumentUtil, SparkReporter}
|
|
import org.apache.spark.SparkContext
|
|
import org.apache.spark.rdd.RDD.rddToPairRDDFunctions
|
|
import org.apache.spark.sql.catalyst.encoders.{ExpressionEncoder, RowEncoder}
|
|
import org.apache.spark.sql.{Column, Dataset, Encoder, Encoders, Row, functions}
|
|
import org.apache.spark.sql.catalyst.expressions.{GenericRowWithSchema, Literal}
|
|
import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, UserDefinedAggregateFunction, UserDefinedFunction, Window}
|
|
import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, Metadata, StructField, StructType}
|
|
|
|
import java.util
|
|
import java.util.function.Predicate
|
|
import java.util.regex.Pattern
|
|
import scala.collection.JavaConverters._
|
|
import scala.collection.mutable
|
|
import org.apache.spark.sql.functions.{col, lit, udf}
|
|
|
|
import java.util.Collections
|
|
import java.util.stream.Collectors
|
|
|
|
case class SparkDedupConfig(conf: DedupConfig, numPartitions: Int) extends Serializable {
|
|
|
|
private val URL_REGEX: Pattern = Pattern.compile("^\\s*(http|https|ftp)\\://.*")
|
|
|
|
private val CONCAT_REGEX: Pattern = Pattern.compile("\\|\\|\\|")
|
|
|
|
private val urlFilter = (s: String) => URL_REGEX.matcher(s).matches
|
|
|
|
val modelExtractor: (Dataset[String] => Dataset[Row]) = df => {
|
|
|
|
df.withColumn("mapDocument", rowFromJsonUDF.apply(df.col(df.columns(0))))
|
|
.withColumn("identifier", new Column("mapDocument.identifier"))
|
|
//.repartition(new Column("identifier"))
|
|
.dropDuplicates("identifier")
|
|
.select("mapDocument.*")
|
|
|
|
df.map(r => rowFromJson(r))(RowEncoder(rowDataType))
|
|
.dropDuplicates("identifier")
|
|
}
|
|
|
|
val generateClusters: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
val df_with_keys = conf
|
|
.clusterings()
|
|
.asScala
|
|
.foldLeft(df_with_filters)((res, cd) => {
|
|
res.withColumn(
|
|
cd.getName + "_clustered",
|
|
functions.explode_outer(
|
|
clusterValuesUDF(cd).apply(
|
|
functions.array(
|
|
cd.getFields.asScala
|
|
.map(f => res.col(if (conf.blacklists.containsKey(f)) f.concat("_filtered") else f)): _*
|
|
)
|
|
)
|
|
)
|
|
)
|
|
})
|
|
|
|
// filter blacklisted values// filter blacklisted values
|
|
// create one column per cluster prefix// create one column per cluster prefix
|
|
|
|
// GROUPING sets approach// GROUPING sets approach
|
|
val tempTable = this.getClass.getSimpleName + "__generateClusters";
|
|
|
|
df_with_keys.createOrReplaceTempView(this.getClass.getSimpleName + "__generateClusters")
|
|
|
|
val keys = conf.clusterings().asScala.map(_.getName + "_clustered").mkString(",")
|
|
val fields = rowDataType.fieldNames.mkString(",")
|
|
|
|
// Using SQL because GROUPING SETS are not available through Scala/Java DSL
|
|
df_with_keys.sqlContext.sql(
|
|
("SELECT coalesce(" + keys + ") as key, sort_array(collect_sort_slice(" + fields + ")) as block FROM " + tempTable + " WHERE coalesce(" + keys + ") IS NOT NULL GROUP BY GROUPING SETS (" + keys + ") HAVING size(block) > 1")
|
|
)
|
|
|
|
}
|
|
|
|
val generateClustersWithDFAPI: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
System.out.println(conf.getWf.getEntityType + "::" +conf.getWf.getSubEntityType)
|
|
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
var relBlocks: Dataset[Row] = null
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
for (cd <- conf.clusterings()) {
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
val tmp: Dataset[Row] = df_with_filters.withColumn("key", functions.explode(clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))))
|
|
|
|
|
|
|
|
|
|
/*.select((Seq(rowDataType.fieldNames: _*) ++ Seq("key")).map(col): _*)
|
|
.groupByKey(r => r.getAs[String]("key"))(Encoders.STRING)
|
|
.agg(collectSortSliceAggregator.toColumn)
|
|
.toDF("key", "block")
|
|
.select(col("block.block").as("block"))*/
|
|
|
|
System.out.println(cd.getName)
|
|
|
|
val ds = tmp.groupBy("key")
|
|
// .agg(functions.sort_array(collectSortSliceUDAF(rowDataType.fieldNames.map(col): _*)).as("block"))
|
|
.agg(functions.collect_set(functions.struct(rowDataType.fieldNames.map(col): _*)).as("block"))
|
|
//.filter(functions.size(new Column("block")).geq(new Literal(2, DataTypes.IntegerType)))
|
|
|
|
|
|
//df_with_filters.printSchema()
|
|
//ds.printSchema()
|
|
|
|
if (relBlocks == null) relBlocks = ds
|
|
else relBlocks = relBlocks.union(ds)
|
|
}
|
|
|
|
// System.out.println()
|
|
|
|
relBlocks
|
|
}
|
|
|
|
val generateClustersWithWindows: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
var relBlocks: Dataset[Row] = null
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
for (cd <- conf.clusterings()) {
|
|
System.out.println(conf.getWf.getEntityType + "::" + conf.getWf.getSubEntityType+ ": " + cd.getName + " " + cd.toString)
|
|
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
// Add 'key' column with the value generated by the given clustering definition
|
|
val ds: Dataset[Row] = df_with_filters.withColumn("key", functions.explode(clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))))
|
|
// Add position column having the position of the row within the set of rows having the same key value ordered by the sorting value
|
|
.withColumn("position", functions.row_number().over(Window.partitionBy("key").orderBy(col(conf.getWf.getOrderField))))
|
|
// filter out rows with position exceeding the maxqueuesize parameter
|
|
.filter(col("position").leq(conf.getWf.getQueueMaxSize))
|
|
.groupBy("key")
|
|
.agg(functions.collect_set(functions.struct(rowDataType.fieldNames.map(col): _*)).as("block"))
|
|
.filter(functions.size(new Column("block")).geq(new Literal(2, DataTypes.IntegerType)))
|
|
|
|
if (relBlocks == null) relBlocks = ds
|
|
else relBlocks = relBlocks.union(ds)
|
|
}
|
|
|
|
relBlocks
|
|
}
|
|
|
|
val generateClustersWithDFAPIMerged: (Dataset[Row] => Dataset[Row]) = df => {
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
val keys = conf.clusterings().foldLeft(null : Column)((res, cd) => {
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
if (res != null)
|
|
functions.array_union(res, clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*)))
|
|
else
|
|
clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))
|
|
})
|
|
|
|
val ds: Dataset[Row] = df_with_filters.withColumn("key", functions.explode(keys))
|
|
.select((Seq(rowDataType.fieldNames: _*) ++ Seq("key")).map(col): _*)
|
|
.groupByKey(r => r.getAs[String]("key"))(Encoders.STRING)
|
|
.agg(collectSortSliceAggregator.toColumn)
|
|
.toDF("key", "block")
|
|
.select(col("block.block").as("block"))
|
|
|
|
/*.groupBy("key")
|
|
.agg(collectSortSliceUDAF(rowDataType.fieldNames.map(col): _*).as("block"))*/
|
|
.filter(functions.size(new Column("block")).geq(new Literal(2, DataTypes.IntegerType)))
|
|
|
|
ds
|
|
}
|
|
|
|
val generateClustersWithRDDReduction: (Dataset[Row] => Dataset[Row]) = df => {
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
var relBlocks: Dataset[Row] = null
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
for (cd <- conf.clusterings()) {
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
val ds: Dataset[Row] = df.sparkSession.createDataFrame(df_with_filters.withColumn("key", functions.explode(clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))))
|
|
.select(col("key"), functions.array(functions.struct(rowDataType.fieldNames.map(col): _*).as("value")))
|
|
.rdd.keyBy(_.getString(0))
|
|
.reduceByKey((a, b) => {
|
|
val b1 = a.getSeq[Row](1)
|
|
val b2 = b.getSeq[Row](1)
|
|
|
|
if (b1.size + b2.size > conf.getWf.getQueueMaxSize)
|
|
Row(a.get(0), b1.union(b2).sortBy(_.getString(orderingFieldPosition)).slice(0, conf.getWf.getQueueMaxSize))
|
|
else
|
|
Row(a.get(0), b1.union(b2))
|
|
})
|
|
.map(_._2)
|
|
.filter(k => k.getSeq(1).size > 1),
|
|
new StructType().add(StructField("key", DataTypes.StringType)).add(StructField("block", ArrayType(rowDataType)))
|
|
)
|
|
|
|
if (relBlocks == null) relBlocks = ds
|
|
else relBlocks = relBlocks.union(ds)
|
|
}
|
|
|
|
relBlocks
|
|
}
|
|
|
|
val printAnalytics: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
var relBlocks: Dataset[Row] = null
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
for (cd <- conf.clusterings()) {
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
// Add 'key' column with the value generated by the given clustering definition
|
|
val ds: Dataset[Row] = df_with_filters.withColumn("key", functions.explode(clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))))
|
|
// Add position column having the position of the row within the set of rows having the same key value ordered by the sorting value
|
|
.withColumn("position", functions.row_number().over(Window.partitionBy("key").orderBy(conf.getWf.getOrderField)))
|
|
// filter out rows with position exceeding the maxqueuesize parameter
|
|
.filter(col("position").lt(conf.getWf.getQueueMaxSize))
|
|
|
|
|
|
// inner join to compute all combination of rows to compare
|
|
// note the condition on position to obtain 'windowing': given a row this is compared at most with the next
|
|
// SlidingWindowSize rows following the sort order
|
|
val dsWithMatch = ds.as("l").join(ds.as("r"),
|
|
col("l.key").equalTo(col("r.key")),
|
|
"inner"
|
|
)
|
|
.filter((col("l.position").lt(col("r.position")))
|
|
&& (col("r.position").lt(col("l.position").plus(lit(conf.getWf.getSlidingWindowSize)))))
|
|
// Add match column with the result of comparison
|
|
// dsWithMatch.show(false)
|
|
|
|
if (relBlocks == null)
|
|
relBlocks = dsWithMatch
|
|
else
|
|
relBlocks = relBlocks.union(dsWithMatch)
|
|
}
|
|
|
|
System.out.println(conf.getWf.getEntityType + "::" + conf.getWf.getSubEntityType)
|
|
System.out.println("Total number of comparations: " + relBlocks.count())
|
|
|
|
|
|
df
|
|
}
|
|
|
|
|
|
val generateAndProcessClustersWithJoins: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
val df_with_filters = conf.getPace.getModel.asScala.foldLeft(df)((res, fdef) => {
|
|
if (conf.blacklists.containsKey(fdef.getName)) {
|
|
res.withColumn(
|
|
fdef.getName + "_filtered",
|
|
filterColumnUDF(fdef).apply(new Column(fdef.getName))
|
|
)
|
|
} else {
|
|
res
|
|
}
|
|
})
|
|
|
|
var relBlocks: Dataset[Row] = null
|
|
|
|
import scala.collection.JavaConversions._
|
|
|
|
for (cd <- conf.clusterings()) {
|
|
val columns: util.List[Column] = new util.ArrayList[Column](cd.getFields().size)
|
|
|
|
for (fName <- cd.getFields()) {
|
|
if (conf.blacklists.containsKey(fName))
|
|
columns.add(new Column(fName + "_filtered"))
|
|
else
|
|
columns.add(new Column(fName))
|
|
}
|
|
|
|
// Add 'key' column with the value generated by the given clustering definition
|
|
val ds: Dataset[Row] = df_with_filters.withColumn("key", functions.explode(clusterValuesUDF(cd).apply(functions.array(columns.asScala: _*))))
|
|
// Add position column having the position of the row within the set of rows having the same key value ordered by the sorting value
|
|
.withColumn("position", functions.row_number().over(Window.partitionBy("key").orderBy(conf.getWf.getOrderField)))
|
|
// filter out rows with position exceeding the maxqueuesize parameter
|
|
.filter(col("position").lt(conf.getWf.getQueueMaxSize))
|
|
|
|
|
|
// inner join to compute all combination of rows to compare
|
|
// note the condition on position to obtain 'windowing': given a row this is compared at most with the next
|
|
// SlidingWindowSize rows following the sort order
|
|
val dsWithMatch = ds.as("l").join(ds.as("r"),
|
|
col("l.key").equalTo(col("r.key")),
|
|
"inner"
|
|
)
|
|
.filter((col("l.position").lt(col("r.position")))
|
|
&& (col("r.position").lt(col("l.position").plus(lit(conf.getWf.getSlidingWindowSize)))))
|
|
// Add match column with the result of comparison
|
|
.withColumn("match", udf[Boolean, Row, Row]((a, b) => {
|
|
val treeProcessor = new TreeProcessor(conf)
|
|
|
|
treeProcessor.compare(a, b)
|
|
}).apply(functions.struct(rowDataType.fieldNames.map(s => col("l.".concat(s))): _*), functions.struct(rowDataType.fieldNames.map(s => col("r.".concat(s))): _*)))
|
|
.filter(col("match").equalTo(true))
|
|
.select(col("l.identifier").as("from"), col("r.identifier").as("to"))
|
|
|
|
// dsWithMatch.show(false)
|
|
|
|
if (relBlocks == null)
|
|
relBlocks = dsWithMatch
|
|
else
|
|
relBlocks = relBlocks.union(dsWithMatch)
|
|
}
|
|
|
|
val res = relBlocks
|
|
//.select(col("l.identifier").as("from"), col("r.identifier").as("to"))
|
|
//.repartition()
|
|
.distinct()
|
|
|
|
// res.show(false)
|
|
res.select(functions.struct("from", "to"))
|
|
}
|
|
|
|
val processClusters: (Dataset[Row] => Dataset[Row]) = df => {
|
|
|
|
val entity = conf.getWf.getEntityType
|
|
|
|
df.filter(functions.size(new Column("block")).geq(new Literal(2, DataTypes.IntegerType)))
|
|
.withColumn("relations", processBlock(df.sqlContext.sparkContext).apply(new Column("block")))
|
|
.select(functions.explode(new Column("relations")).as("relation"))
|
|
//.repartition(new Column("relation"))
|
|
.dropDuplicates("relation")
|
|
}
|
|
|
|
val rowDataType: StructType = {
|
|
// val unordered = conf.getPace.getModel.asScala.foldLeft(
|
|
// new StructType()
|
|
// )((resType, fdef) => {
|
|
// resType.add(fdef.getType match {
|
|
// case Type.List | Type.JSON =>
|
|
// StructField(fdef.getName, DataTypes.createArrayType(DataTypes.StringType), true, Metadata.empty)
|
|
// case Type.DoubleArray =>
|
|
// StructField(fdef.getName, DataTypes.createArrayType(DataTypes.DoubleType), true, Metadata.empty)
|
|
// case _ =>
|
|
// StructField(fdef.getName, DataTypes.StringType, true, Metadata.empty)
|
|
// })
|
|
// })
|
|
//
|
|
// conf.getPace.getModel.asScala.filterNot(_.getName.equals(conf.getWf.getOrderField)).foldLeft(
|
|
// new StructType()
|
|
// .add(unordered(conf.getWf.getOrderField))
|
|
// .add(StructField("identifier", DataTypes.StringType, false, Metadata.empty))
|
|
// )((resType, fdef) => resType.add(unordered(fdef.getName)))
|
|
|
|
val identifier = new FieldDef()
|
|
identifier.setName("identifier")
|
|
identifier.setType(Type.String)
|
|
|
|
(conf.getPace.getModel.asScala ++ Seq(identifier)).sortBy(_.getName)
|
|
.foldLeft(
|
|
new StructType()
|
|
)((resType, fdef) => {
|
|
resType.add(fdef.getType match {
|
|
case Type.List | Type.JSON =>
|
|
StructField(fdef.getName, DataTypes.createArrayType(DataTypes.StringType), true, Metadata.empty)
|
|
case Type.DoubleArray =>
|
|
StructField(fdef.getName, DataTypes.createArrayType(DataTypes.DoubleType), true, Metadata.empty)
|
|
case _ =>
|
|
StructField(fdef.getName, DataTypes.StringType, true, Metadata.empty)
|
|
})
|
|
})
|
|
|
|
|
|
|
|
}
|
|
|
|
val identityFieldPosition: Int = rowDataType.fieldIndex("identifier")
|
|
|
|
val orderingFieldPosition: Int = rowDataType.fieldIndex(conf.getWf.getOrderField)
|
|
|
|
def rowFromJson(json: String) : Row = {
|
|
val documentContext =
|
|
JsonPath.using(Configuration.defaultConfiguration.addOptions(Option.SUPPRESS_EXCEPTIONS)).parse(json)
|
|
val values = new Array[Any](rowDataType.size)
|
|
|
|
values(identityFieldPosition) = MapDocumentUtil.getJPathString(conf.getWf.getIdPath, documentContext)
|
|
|
|
rowDataType.fieldNames.zipWithIndex.foldLeft(values) {
|
|
case ((res, (fname, index))) => {
|
|
val fdef = conf.getPace.getModelMap.get(fname)
|
|
|
|
if (fdef != null) {
|
|
res(index) = fdef.getType match {
|
|
case Type.String | Type.Int =>
|
|
MapDocumentUtil.truncateValue(
|
|
MapDocumentUtil.getJPathString(fdef.getPath, documentContext),
|
|
fdef.getLength
|
|
)
|
|
|
|
case Type.URL =>
|
|
var uv = MapDocumentUtil.getJPathString(fdef.getPath, documentContext)
|
|
if (!urlFilter(uv)) uv = ""
|
|
uv
|
|
|
|
case Type.List | Type.JSON =>
|
|
MapDocumentUtil.truncateList(
|
|
MapDocumentUtil.getJPathList(fdef.getPath, documentContext, fdef.getType),
|
|
fdef.getSize
|
|
).toArray
|
|
|
|
case Type.StringConcat =>
|
|
val jpaths = CONCAT_REGEX.split(fdef.getPath)
|
|
|
|
truncateValue(
|
|
jpaths
|
|
.map(jpath => MapDocumentUtil.getJPathString(jpath, documentContext))
|
|
.mkString(" "),
|
|
fdef.getLength
|
|
)
|
|
|
|
case Type.DoubleArray =>
|
|
MapDocumentUtil.getJPathArray(fdef.getPath, json)
|
|
}
|
|
}
|
|
|
|
res
|
|
}
|
|
}
|
|
|
|
new GenericRowWithSchema(values, rowDataType)
|
|
}
|
|
|
|
val rowFromJsonUDF = udf(rowFromJson(_), rowDataType)
|
|
|
|
def filterColumnUDF(fdef: FieldDef): UserDefinedFunction = {
|
|
|
|
val blacklist: Predicate[String] = conf.blacklists().get(fdef.getName)
|
|
|
|
if (blacklist == null) {
|
|
throw new IllegalArgumentException("Column: " + fdef.getName + " does not have any filter")
|
|
} else {
|
|
fdef.getType match {
|
|
case Type.List | Type.JSON =>
|
|
udf[Array[String], Array[String]](values => {
|
|
values.filter((v: String) => !blacklist.test(v))
|
|
})
|
|
|
|
case _ =>
|
|
udf[String, String](v => {
|
|
if (blacklist.test(v)) ""
|
|
else v
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
def clusterValuesUDF(cd: ClusteringDef) = {
|
|
udf[mutable.WrappedArray[String], mutable.WrappedArray[Object]](values => {
|
|
values.flatMap(f => cd.clusteringFunction().apply(conf, Seq(f.toString).asJava).asScala).map(cd.getName.concat(_))
|
|
})
|
|
}
|
|
|
|
def processBlock(implicit sc: SparkContext) = {
|
|
val accumulators = SparkReporter.constructAccumulator(conf, sc)
|
|
|
|
udf[Array[Tuple2[String, String]], mutable.WrappedArray[Row]](block => {
|
|
val reporter = new SparkReporter(accumulators)
|
|
|
|
val mapDocuments = block.asJava.stream
|
|
.sorted(new RowDataOrderingComparator(orderingFieldPosition))
|
|
.limit(conf.getWf.getQueueMaxSize)
|
|
.collect(Collectors.toList[Row]())
|
|
|
|
new BlockProcessor(conf, identityFieldPosition, orderingFieldPosition).processSortedRows(mapDocuments, reporter)
|
|
|
|
reporter.getRelations.asScala.toArray
|
|
}).asNondeterministic()
|
|
}
|
|
|
|
val collectSortSliceAggregator : Aggregator[Row,Seq[Row], Row] = new Aggregator[Row, Seq[Row], Row] () {
|
|
override def zero: Seq[Row] = Seq[Row]()
|
|
|
|
|
|
override def reduce(buffer: Seq[Row], input: Row): Seq[Row] = {
|
|
merge(buffer, Seq(input))
|
|
}
|
|
|
|
override def merge(buffer: Seq[Row], toMerge: Seq[Row]): Seq[Row] = {
|
|
val newBlock = buffer ++ toMerge
|
|
|
|
if (newBlock.size > conf.getWf.getQueueMaxSize)
|
|
newBlock.sortBy(_.getString(orderingFieldPosition)).slice(0, conf.getWf.getQueueMaxSize)
|
|
else
|
|
newBlock
|
|
}
|
|
|
|
override def finish(reduction: Seq[Row]): Row = {
|
|
Row(reduction.toArray)
|
|
}
|
|
|
|
override def bufferEncoder: Encoder[Seq[Row]] = Encoders.kryo[Seq[Row]]
|
|
|
|
override def outputEncoder: Encoder[Row] = RowEncoder.apply(new StructType().add("block", DataTypes.createArrayType(rowDataType), nullable = true))
|
|
}
|
|
|
|
val collectSortSliceUDAF : UserDefinedAggregateFunction = new UserDefinedAggregateFunction {
|
|
override def inputSchema: StructType = rowDataType
|
|
|
|
override def bufferSchema: StructType = {
|
|
new StructType().add("block", DataTypes.createArrayType(rowDataType), nullable = true)
|
|
}
|
|
|
|
override def dataType: DataType = DataTypes.createArrayType(rowDataType)
|
|
|
|
override def deterministic: Boolean = true
|
|
|
|
override def initialize(buffer: MutableAggregationBuffer): Unit = {
|
|
buffer(0) = Seq[Row]()
|
|
}
|
|
|
|
override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
|
|
val newBlock = buffer.getSeq[Row](0) ++ Seq(input)
|
|
|
|
if (newBlock.size > conf.getWf.getQueueMaxSize)
|
|
buffer(0) = newBlock.sortBy(_.getString(orderingFieldPosition)).slice(0, conf.getWf.getQueueMaxSize)
|
|
else
|
|
buffer(0) = newBlock
|
|
}
|
|
|
|
override def merge(buffer: MutableAggregationBuffer, row: Row): Unit = {
|
|
val newBlock = buffer.getSeq[Row](0) ++ row.getSeq[Row](0)
|
|
|
|
if (newBlock.size > conf.getWf.getQueueMaxSize)
|
|
buffer(0) = newBlock.sortBy(_.getString(orderingFieldPosition)).slice(0, conf.getWf.getQueueMaxSize)
|
|
else
|
|
buffer(0) = newBlock
|
|
}
|
|
|
|
override def evaluate(buffer: Row): Any = {
|
|
buffer.getSeq[Row](0)
|
|
}
|
|
}
|
|
}
|