web-dev-qa-db-ja.com

Spark SQLでユーザー定義の集計関数を定義して使用するには？

Spark SQLでUDFを書く方法を知っています。

_def belowThreshold(power: Int): Boolean = {
        return power < -40
      }

sqlContext.udf.register("belowThreshold", belowThreshold _)
_

集約関数を定義するのと似たようなことはできますか？これはどのように行われますか？

コンテキストでは、次のSQLクエリを実行します。

_val aggDF = sqlContext.sql("""SELECT span, belowThreshold(opticalReceivePower), timestamp
                                    FROM ifDF
                                    WHERE opticalReceivePower IS NOT null
                                    GROUP BY span, timestamp
                                    ORDER BY span""")
_

次のようなものが返されるはずです

Row(span1, false, T0)

opticalReceivePowerとspanで定義されたグループ内のtimestampの値がしきい値を下回っているかどうかを集約関数に教えてほしい。上で貼り付けたUDFとは異なる方法でUDAFを作成する必要がありますか？

scalaapache-sparkapache-spark-sqlaggregate-functionsuser-defined-functions
36
Rory Byrne

サポートされている方法

Spark> = 2.3

ベクトル化されたudf（Pythonのみ）：

from pyspark.sql.functions import pandas_udf
from pyspark.sql.functions import PandasUDFType

from pyspark.sql.types import *
import pandas as pd

df = sc.parallelize([
    ("a", 0), ("a", 1), ("b", 30), ("b", -50)
]).toDF(["group", "power"])

def below_threshold(threshold, group="group", power="power"):
    @pandas_udf("struct<group: string, below_threshold: boolean>", PandasUDFType.GROUPED_MAP)
    def below_threshold_(df):
        df = pd.DataFrame(
           df.groupby(group).apply(lambda x: (x[power] < threshold).any()))
        df.reset_index(inplace=True, drop=False)
        return df

    return below_threshold_

使用例：

df.groupBy("group").apply(below_threshold(-40)).show()

## +-----+---------------+
## |group|below_threshold|
## +-----+---------------+
## |    b|           true|
## |    a|          false|
## +-----+---------------+

PySparkでGroupedDataにUDFを適用する（機能するpython例）

Spark> = 2.0（オプションで1.6ですが、APIが若干異なります）：

型付きAggregatorsでDatasetsを使用することができます。

import org.Apache.spark.sql.expressions.Aggregator
import org.Apache.spark.sql.{Encoder, Encoders}

class BelowThreshold[I](f: I => Boolean)  extends Aggregator[I, Boolean, Boolean]
    with Serializable {
  def zero = false
  def reduce(acc: Boolean, x: I) = acc | f(x)
  def merge(acc1: Boolean, acc2: Boolean) = acc1 | acc2
  def finish(acc: Boolean) = acc

  def bufferEncoder: Encoder[Boolean] = Encoders.scalaBoolean
  def outputEncoder: Encoder[Boolean] = Encoders.scalaBoolean
}

val belowThreshold = new BelowThreshold[(String, Int)](_._2 < - 40).toColumn
df.as[(String, Int)].groupByKey(_._1).agg(belowThreshold)

Spark> = 1.5：

Spark 1.5では、次のようにUDAFを作成できますが、これはおそらくやり過ぎです。

import org.Apache.spark.sql.expressions._
import org.Apache.spark.sql.types._
import org.Apache.spark.sql.Row

object belowThreshold extends UserDefinedAggregateFunction {
    // Schema you get as an input
    def inputSchema = new StructType().add("power", IntegerType)
    // Schema of the row which is used for aggregation
    def bufferSchema = new StructType().add("ind", BooleanType)
    // Returned type
    def dataType = BooleanType
    // Self-explaining 
    def deterministic = true
    // zero value
    def initialize(buffer: MutableAggregationBuffer) = buffer.update(0, false)
    // Similar to seqOp in aggregate
    def update(buffer: MutableAggregationBuffer, input: Row) = {
        if (!input.isNullAt(0))
          buffer.update(0, buffer.getBoolean(0) | input.getInt(0) < -40)
    }
    // Similar to combOp in aggregate
    def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = {
      buffer1.update(0, buffer1.getBoolean(0) | buffer2.getBoolean(0))    
    }
    // Called on exit to get return value
    def evaluate(buffer: Row) = buffer.getBoolean(0)
}

使用例：

df
  .groupBy($"group")
  .agg(belowThreshold($"power").alias("belowThreshold"))
  .show

// +-----+--------------+
// |group|belowThreshold|
// +-----+--------------+
// |    a|         false|
// |    b|          true|
// +-----+--------------+

Spark 1.4の回避策：

私はあなたの要件を正しく理解しているかどうかはわかりませんが、ここでは単純な古い集約で十分であると言えます。

val df = sc.parallelize(Seq(
    ("a", 0), ("a", 1), ("b", 30), ("b", -50))).toDF("group", "power")

df
  .withColumn("belowThreshold", ($"power".lt(-40)).cast(IntegerType))
  .groupBy($"group")
  .agg(sum($"belowThreshold").notEqual(0).alias("belowThreshold"))
  .show

// +-----+--------------+
// |group|belowThreshold|
// +-----+--------------+
// |    a|         false|
// |    b|          true|
// +-----+--------------+

スパーク<= 1.4：

私の知る限り、現時点（Spark 1.4.1）では、Hive以外のUDAFのサポートはありません。 Spark 1.5（ SPARK-3947 を参照）で可能になるはずです。

サポートされていない/内部メソッド

内部的にSparkは ImperativeAggregates および DeclarativeAggregates を含む多くのクラスを使用します。

内部での使用を目的としており、予告なしに変更される可能性があるため、本番コードで使用したいものではないかもしれませんが、完全を期すためにBelowThreshold with DeclarativeAggregateをこのように実装できます（Spark 2.2-SNAPSHOT）でテスト済み）：

import org.Apache.spark.sql.catalyst.expressions.aggregate.DeclarativeAggregate
import org.Apache.spark.sql.catalyst.expressions._
import org.Apache.spark.sql.types._

case class BelowThreshold(child: Expression, threshold: Expression) 
    extends  DeclarativeAggregate  {
  override def children: Seq[Expression] = Seq(child, threshold)

  override def nullable: Boolean = false
  override def dataType: DataType = BooleanType

  private lazy val belowThreshold = AttributeReference(
    "belowThreshold", BooleanType, nullable = false
  )()

  // Used to derive schema
  override lazy val aggBufferAttributes = belowThreshold :: Nil

  override lazy val initialValues = Seq(
    Literal(false)
  )

  override lazy val updateExpressions = Seq(Or(
    belowThreshold,
    If(IsNull(child), Literal(false), LessThan(child, threshold))
  ))

  override lazy val mergeExpressions = Seq(
    Or(belowThreshold.left, belowThreshold.right)
  )

  override lazy val evaluateExpression = belowThreshold
  override def defaultResult: Option[Literal] = Option(Literal(false))
}

withAggregateFunction と同等のものでさらにラップする必要があります。

74
zero323