Spark：求出分组内的TopN

制作测试数据源：

spark scala实现代码：

import org.apache.spark.SparkConf
import org.apache.spark.sql.SparkSession

object GroupTopN1 {
  System.setProperty("hadoop.home.dir", "D:\Java_Study\hadoop-common-2.2.0-bin-master")

  case class Rating(userId: String, rating: Long)

  def main(args: Array[String]) {
    val sparkConf = new SparkConf().setAppName("ALS with ML Pipeline")
    val spark = SparkSession
      .builder()
      .config(sparkConf)
      .master("local")
      .config("spark.sql.warehouse.dir", "/")
      .getOrCreate()

    import spark.implicits._
    import spark.sql

    val lines = spark.read.textFile("C:\Users\Administrator\Desktop\group.txt")
    val classScores = lines.map(line => Rating(line.split(" ")(0).toString, line.split(" ")(1).toLong))

    classScores.createOrReplaceTempView("tb_test")

    var df = sql(
      s"""|select
          | userId,
          | rating,
          | row_number()over(partition by userId order by rating desc) rn
          |from tb_test
          |having(rn<=3)
          |""".stripMargin)
    df.show()

    spark.stop()
  }
}

打印结果：

+------+------+---+
|userId|rating| rn|
+------+------+---+
|    c1|    85|  1|
|    c1|    66|  2|
|    c1|    65|  3|
|    c4|    65|  1|
|    c3|    95|  1|
|    c3|    88|  2|
|    c3|    54|  3|
|    c2|    91|  1|
|    c2|    77|  2|
|    c2|    66|  3|
+------+------+---+

spark java代码实现：

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.MapFunction;
import org.apache.spark.sql.*;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import scala.Function1;

import javax.management.RuntimeErrorException;
import java.util.List;
import java.util.ArrayList;

public class Test {
    public static void main(String[] args) {
        System.out.println("Hello");
        SparkConf sparkConf = new SparkConf().setAppName("ALS with ML Pipeline");
        SparkSession spark = SparkSession
                .builder()
                .config(sparkConf)
                .master("local")
                .config("spark.sql.warehouse.dir", "/")
                .getOrCreate();


        // Create an RDD
        JavaRDD<String> peopleRDD = spark.sparkContext()
                .textFile("C:\Users\Administrator\Desktop\group.txt", 1)
                .toJavaRDD();

        // The schema is encoded in a string
        String schemaString = "userId rating";

        // Generate the schema based on the string of schema
        List<StructField> fields = new ArrayList<>();
        StructField field1 = DataTypes.createStructField("userId", DataTypes.StringType, true);
        StructField field2 = DataTypes.createStructField("rating", DataTypes.LongType, true);
        fields.add(field1);
        fields.add(field2);
        StructType schema = DataTypes.createStructType(fields);

        // Convert records of the RDD (people) to Rows
        JavaRDD<Row> rowRDD = peopleRDD.map((Function<String, Row>) record -> {
            String[] attributes = record.split(" ");
            if(attributes.length!=2){
                throw new Exception();
            }
            return RowFactory.create(attributes[0],Long.valueOf( attributes[1].trim()));
        });

        // Apply the schema to the RDD
        Dataset<Row> peopleDataFrame = spark.createDataFrame(rowRDD, schema);

        peopleDataFrame.createOrReplaceTempView("tb_test");

        Dataset<Row> items = spark.sql("select userId,rating,row_number()over(partition by userId order by rating desc) rn " +
                "from tb_test " +
                "having(rn<=3)");
        items.show();

        spark.stop();
    }
}

输出结果同上边输出结果。

Java 中使用combineByKey实现TopN:

import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import scala.Tuple2;

import java.util.*;

public class SparkJava {
    public static void main(String[] args) {
        SparkSession spark = SparkSession.builder().master("local[*]").appName("Spark").getOrCreate();
        final JavaSparkContext ctx = JavaSparkContext.fromSparkContext(spark.sparkContext());

        List<String> data = Arrays.asList("a,110,a1", "b,122,b1", "c,123,c1", "a,210,a2", "b,212,b2", "a,310,a3", "b,312,b3", "a,410,a4", "b,412,b4");
        JavaRDD<String> javaRDD = ctx.parallelize(data);

        JavaPairRDD<String, Integer> javaPairRDD = javaRDD.mapToPair(new PairFunction<String, String, Integer>() {
            public Tuple2<String, Integer> call(String key) throws Exception {
                return new Tuple2<String, Integer>(key.split(",")[0], Integer.valueOf(key.split(",")[1]));
            }
        });

        final int topN = 3;
        JavaPairRDD<String, List<Integer>> combineByKeyRDD2 = javaPairRDD.combineByKey(new Function<Integer, List<Integer>>() {
            public List<Integer> call(Integer v1) throws Exception {
                List<Integer> items = new ArrayList<Integer>();
                items.add(v1);
                return items;
            }
        }, new Function2<List<Integer>, Integer, List<Integer>>() {
            public List<Integer> call(List<Integer> v1, Integer v2) throws Exception {
                if (v1.size() > topN) {
                    Integer item = Collections.min(v1);
                    v1.remove(item);
                    v1.add(v2);
                }
                return v1;
            }
        }, new Function2<List<Integer>, List<Integer>, List<Integer>>() {
            public List<Integer> call(List<Integer> v1, List<Integer> v2) throws Exception {
                v1.addAll(v2);
                while (v1.size() > topN) {
                    Integer item = Collections.min(v1);
                    v1.remove(item);
                }

                return v1;
            }
        });

        // 由K:String,V:List<Integer> 转化为 K:String,V:Integer
        // old:[(a,[210, 310, 410]), (b,[122, 212, 312]), (c,[123])]
        // new:[(a,210), (a,310), (a,410), (b,122), (b,212), (b,312), (c,123)]
        JavaRDD<Tuple2<String, Integer>> javaTupleRDD = combineByKeyRDD2.flatMap(new FlatMapFunction<Tuple2<String, List<Integer>>, Tuple2<String, Integer>>() {
            public Iterator<Tuple2<String, Integer>> call(Tuple2<String, List<Integer>> stringListTuple2) throws Exception {
                List<Tuple2<String, Integer>> items=new ArrayList<Tuple2<String, Integer>>();
                for(Integer v:stringListTuple2._2){
                    items.add(new Tuple2<String, Integer>(stringListTuple2._1,v));
                }
                return items.iterator();
            }
        });

        JavaRDD<Row> rowRDD = javaTupleRDD.map(new Function<Tuple2<String, Integer>, Row>() {
            public Row call(Tuple2<String, Integer> kv) throws Exception {
                String key = kv._1;
                Integer num = kv._2;

                return RowFactory.create(key, num);
            }
        });

        ArrayList<StructField> fields = new ArrayList<StructField>();
        StructField field = null;
        field = DataTypes.createStructField("key", DataTypes.StringType, true);
        fields.add(field);
        field = DataTypes.createStructField("TopN_values", DataTypes.IntegerType, true);
        fields.add(field);

        StructType schema = DataTypes.createStructType(fields);

        Dataset<Row> df = spark.createDataFrame(rowRDD, schema);
        df.printSchema();
        df.show();

        spark.stop();
    }
}

输出：

root
 |-- key: string (nullable = true)
 |-- TopN_values: integer (nullable = true)

+---+-----------+
|key|TopN_values|
+---+-----------+
|  a|        210|
|  a|        310|
|  a|        410|
|  b|        122|
|  b|        212|
|  b|        312|
|  c|        123|
+---+-----------+

Spark使用combineByKeyWithClassTag函数实现TopN

combineByKeyWithClassTag函数，借助HashSet的排序，此例是取组内最大的N个元素一下是代码：

createcombiner就简单的将首个元素装进HashSet然后返回就可以了；
mergevalue插入元素之后，如果元素的个数大于N就删除最小的元素；
mergeCombiner在合并之后，如果总的个数大于N，就从一次删除最小的元素，知道Hashset内只有N 个元素。

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

object Main {
  val N = 3

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .master("local[*]")
      .appName("Spark")
      .getOrCreate()
    val sc = spark.sparkContext
    var SampleDataset = List(
      ("apple.com", 3L),
      ("apple.com", 4L),
      ("apple.com", 1L),
      ("apple.com", 9L),
      ("google.com", 4L),
      ("google.com", 1L),
      ("google.com", 2L),
      ("google.com", 3L),
      ("google.com", 11L),
      ("google.com", 32L),
      ("slashdot.org", 11L),
      ("slashdot.org", 12L),
      ("slashdot.org", 13L),
      ("slashdot.org", 14L),
      ("slashdot.org", 15L),
      ("slashdot.org", 16L),
      ("slashdot.org", 17L),
      ("slashdot.org", 18L),
      ("microsoft.com", 5L),
      ("microsoft.com", 2L),
      ("microsoft.com", 6L),
      ("microsoft.com", 9L),
      ("google.com", 4L))
    val urdd: RDD[(String, Long)] = sc.parallelize(SampleDataset).map((t) => (t._1, t._2))
    var topNs = urdd.combineByKeyWithClassTag(
      //createCombiner
      (firstInt: Long) => {
        var uset = new mutable.TreeSet[Long]()
        uset += firstInt
      },
      // mergeValue
      (uset: mutable.TreeSet[Long], value: Long) => {
        uset += value
        while (uset.size > N) {
          uset.remove(uset.min)
        }
        uset
      },
      //mergeCombiners
      (uset1: mutable.TreeSet[Long], uset2: mutable.TreeSet[Long]) => {
        var resultSet = uset1 ++ uset2
        while (resultSet.size > N) {
          resultSet.remove(resultSet.min)
        }
        resultSet
      }
    )
    import spark.implicits._
    topNs.flatMap(rdd => {
      var uset = new mutable.HashSet[String]()
      for (i <- rdd._2.toList) {
        uset += rdd._1 + "/" + i.toString
      }
      uset
    }).map(rdd => {
      (rdd.split("/")(0), rdd.split("/")(1))
    }).toDF("key", "TopN_values").show()
  }
}

参考《https://blog.csdn.net/gpwner/article/details/78455234》

输出结果：

+-------------+-----------+
|          key|TopN_values|
+-------------+-----------+
|   google.com|          4|
|   google.com|         11|
|   google.com|         32|
|microsoft.com|          9|
|microsoft.com|          6|
|microsoft.com|          5|
|    apple.com|          4|
|    apple.com|          9|
|    apple.com|          3|
| slashdot.org|         16|
| slashdot.org|         17|
| slashdot.org|         18|
+-------------+-----------+