Apache Spark SQL UDAF над окном, показывающим нечетное поведение с дублирующимся вводом

Я обнаружил, что в Apache Spark SQL (версия 2.2.0), когда пользовательская агрегированная функция (UDAF), которая используется по спецификации окна, снабжена несколькими строками одинакового ввода, UDAF (как бы) не вызывает правильно evaluate метод.

Я смог воспроизвести это поведение как на Java, так и на Scala, локально и на кластере. В приведенном ниже коде показан пример, в котором строки помечены как ложные, если они находятся в пределах 1 секунды от предыдущей строки.

class ExampleUDAF(val timeLimit: Long) extends UserDefinedAggregateFunction {
  def deterministic: Boolean = true
  def inputSchema: StructType = StructType(Array(StructField("unix_time", LongType)))
  def dataType: DataType = BooleanType

  def bufferSchema = StructType(Array(
    StructField("previousKeepTime", LongType),
    StructField("keepRow", BooleanType)
  ))

  def initialize(buffer: MutableAggregationBuffer) = {
    buffer(0) = 0L
    buffer(1) = false
  }

  def update(buffer: MutableAggregationBuffer, input: Row) = {    
    if (buffer(0) == 0L) {
      buffer(0) = input.getLong(0)
      buffer(1) = true
    } else {
      val timeDiff = input.getLong(0) - buffer.getLong(0)

      if (timeDiff < timeLimit) {
        buffer(1) = false
      } else {
        buffer(0) = input.getLong(0)
        buffer(1) = true
      }
    }
  }

  def merge(buffer1: MutableAggregationBuffer, buffer2: Row) = {} // Not implemented
  def evaluate(buffer: Row): Boolean = buffer.getBoolean(1)
 }

val timeLimit = 1000 // 1 second
val udaf = new ExampleUDAF(timeLimit)

val window = Window
  .orderBy(column("unix_time"))
  .partitionBy(column("category"))

val df = spark.createDataFrame(Arrays.asList(
    Row(1510000001000L, "a", true), 
    Row(1510000001000L, "a", false), 
    Row(1510000001000L, "a", false),
    Row(1510000001000L, "a", false),
    Row(1510000700000L, "a", true),
    Row(1510000700000L, "a", false)
  ), new StructType().add("unix_time", LongType).add("category", StringType).add("expected_result", BooleanType))

df.withColumn("actual_result", udaf(column("unix_time")).over(window)).show

Ниже приведен результат выполнения приведенного выше кода. Ожидается, что первая строка будет иметь значение actual_result true, так как предварительных данных нет. Когда вход unix_time изменяется на 1 миллисекунду между каждой записью, UDAF работает так, как ожидалось.

Добавление операторов печати в методы UDAF показало, что evaluate вызывается только один раз, в конце, и этот буфер был правильно обновлен до true в методе update, но это не то, что возвращается после завершения UDAF.

+-------------+--------+---------------+-------------+
|    unix_time|category|expected_result|actual_result|
+-------------+--------+---------------+-------------+
|1510000001000|       a|           true|        false|  // Should true as first element
|1510000001000|       a|          false|        false|
|1510000001000|       a|          false|        false|
|1510000001000|       a|          false|        false|
|1510000700000|       a|           true|        false|  // Should be true as more than 1000 milliseconds between self and previous
|1510000700000|       a|          false|        false|
+-------------+--------+---------------+-------------+

Я правильно понимаю поведение Spark UDAF при использовании над спецификациями окна? Если нет, может ли кто-нибудь предложить какое-либо понимание в этой области. Если мое понимание поведения UDAF над окнами верное, может ли это быть ошибкой в Spark? Спасибо.

ОТВЕТЫ

Ответ 1

Одна из проблем с вашим UDAF заключается в том, что он не указывает, на каких строках вы хотите запустить свое окно с помощью rowsBetween(). Если спецификация rowsBetween(), для каждой строки функция окна будет принимать все (см. Обновление ниже) строки до и после текущей, включая текущую (в данной категории). Таким образом, actual_result для всех строк будет в основном учитывать только две последние строки в вашем примере DataFrame с unix_time=1510000700000 который эффективно возвращает false для всех строк.

С объявлением window следующим образом:

Window.partitionBy(col("category")).orderBy(col("unix_time")).rowsBetween(-1L, 0L)

Вы всегда смотрите только на предыдущую строку и текущую строку. С предыдущей строкой, принятой первой. Это создает правильный вывод. Но поскольку упорядочение строк с одним и тем же unix_time не является уникальным, невозможно предсказать, какая строка будет иметь значение true среди строк с одинаковым unix_time.

Результат может выглядеть так:

+-------------+--------+---------------+-------------+
|    unix_time|category|expected_result|actual_result|
+-------------+--------+---------------+-------------+
|1510000001000|       a|          false|         true|
|1510000001000|       a|          false|        false|
|1510000001000|       a|          false|        false|
|1510000001000|       a|           true|        false|
|1510000700000|       a|           true|         true|
|1510000700000|       a|          false|        false|
+-------------+--------+---------------+-------------+

Обновить

После дальнейшего исследования кажется, что при orderBy столбца orderBy он принимает все элементы перед текущей строкой + текущей строкой. Не все элементы раздела, как я уже говорил. Кроме того, если столбец OrderBy содержит окно с двойными значениями для каждой повторяющейся строки, все дублированные значения будут содержать. Вы можете ясно видеть это:

val wA = Window.partitionBy(col("category")).orderBy(col("unix_time"))
val wB = Window.partitionBy(col("category"))
val wC = Window.partitionBy(col("category")).orderBy(col("unix_time")).rowsBetween(-1L, 0L)

df.withColumn("countRows", count(col("unix_time")).over(wA)).show()
df.withColumn("countRows", count(col("unix_time")).over(wB)).show()
df.withColumn("countRows", count(col("unix_time")).over(wC)).show()

который будет подсчитывать количество элементов в каждом окне.

  • Окно wA будет содержать 4 элемента в каждой строке 1510000001000 и 6 элементов для каждого 1510000700000.
  • Для wB когда нет orderBy все строки включены в окно для каждого раздела, поэтому все окна будут иметь 6 элементов.
  • Последний wC задает выбор строк, поэтому не оставляет двусмысленности, какие строки выбраны для какого окна. Существует только 1 элемент для первой строки и 2 элемента в окнах всех последующих строк. Что дает правильный результат.

Сегодня я узнал что-то новое :)