用java,scala,python三种语言实现spark rdd的treeAggregate算子

Java实现：

import org.apache.spark.api.java.JavaRDD;

public class TreeAggregateJava {

    public static <T> T treeAggregate(JavaRDD<T> rdd, T zeroValue, TreeAggregateFunction<T> func) {
        int depth = (int) Math.ceil(Math.log(rdd.getNumPartitions()) / Math.log(2));
        return rdd.treeAggregate(zeroValue, func, depth);
    }

    public interface TreeAggregateFunction<T> extends org.apache.spark.api.java.function.Function2<T, T, T> {
        T apply(T t1, T t2);

        @Override
        default T call(T t1, T t2) throws Exception {
            return apply(t1, t2);
        }
    }

}

Scala实现：

import org.apache.spark.rdd.RDD

object TreeAggregateScala {

  def treeAggregate[T](rdd: RDD[T], zeroValue: T)(func: (T, T) => T): T = {
    val depth = math.ceil(math.log(rdd.getNumPartitions) / math.log(2)).toInt
    rdd.treeAggregate(zeroValue)(func, func, depth)
  }

}

Python实现：

from pyspark import RDD

def tree_aggregate(rdd: RDD, zero_value, func):
    depth = int(math.ceil(math.log(rdd.getNumPartitions()) / math.log(2)))
    return rdd.treeAggregate(zero_value, func, func, depth)