Функция Scala возвращает разные результаты при применении в UDF

Question

У меня есть функция, которая возвращает точку вставки в «правом поиска»:

import scala.collection.Searching._

def getPit(samples: Seq[Int], obs: Int) : Int = { 
    val pitValue = samples.map(-_).reverse.search(-obs) match {
        case Found(i)          => sampless.length - i
        case InsertionPoint(i) => sampless.length - i
        } // looking for the insertion in a "search right" manner
    pitValue
}

Так что для этих примеров я получаю следующие результаты:

println(getPit(Seq(1, 1, 2, 2), 0)) # 0
println(getPit(Seq(2, 2, 2, 5), 0)) # 0
println(getPit(Seq(0, 1, 7, 10), 7)) # 3
println(getPit(Seq(0, 4, 5, 6), 7)) # 4
println(getPit(Seq(0, 1, 2, 2), 2)) # 4
println(getPit(Seq(1, 1, 1, 1), 1)) # 4
println(getPit(Seq(0, 1, 2, 2), 4)) # 4 

Чтоименно то, что я хочу.

Теперь, применяя эту функцию в UDF следующим образом:

val pitValuesUDF = udf { (samples: Seq[Int], obs: Int) =>
  getPit(samples, obs)
}

val   df =  Seq( ("a",0, Seq(1, 1, 2,2)),
    ("b", 0, Seq(2, 2, 2, 5)),
    ("b", 7, Seq(0, 1, 7, 10)),
    ("b", 7, Seq(0, 4, 5, 6)),
    ("b", 2, Seq(0, 1, 2, 2)),
    ("a", 1, Seq(1, 1, 1, 1)),
    ("b", 4, Seq(0, 1, 2, 2))
 ).toDF("sku_id", "sale_qty", "samples")

df.withColumn("pit", pitValuesUDF(col("samples"), col("sale_qty"))).show

Я получил разные результаты:

+------+--------+-------------+---+
|sku_id|sale_qty|      samples|pit|
+------+--------+-------------+---+
|     a|       0| [1, 1, 2, 2]|  0|
|     b|       0| [2, 2, 2, 5]|  0|
|     b|       7|[0, 1, 7, 10]|  3|
|     b|       7| [0, 4, 5, 6]|  4|
|     b|       2| [0, 1, 2, 2]|  3|
|     a|       1| [1, 1, 1, 1]|  3|
|     b|       4| [0, 1, 2, 2]|  4|
+------+--------+-------------+---+

У вас есть идея, почему результаты различаются в двух случаях? И как я могу заставить UDF возвращать значение в первом случае?

Answer 1

Хорошо, я получил ответ. Глядя на реализацию поиска scala здесь . Мы видим, что есть случай:

 final def search[B >: A](elem: B)(implicit ord: Ordering[B]): SearchResult =
      coll match {
        case _: IndexedSeqLike[A, Repr] => binarySearch(elem, 0, coll.length)(ord)
        case _ => linearSearch(coll.view, elem, 0)(ord)
      }

И угадайте, что binarySearch и linearSearch не будут возвращать один и тот же индекс , когда в коллекции есть дубликаты. Это объясняет, почему могут быть разные результаты для разных типов коллекций.

Чтобы быть более конкретным в первом случае, тип Seq не является IndexedSeqLike, так что это линейная функция, которая выполняется (и это то, что я хочу). Принимая во внимание, что в udf, проверяя samples.getClass, я заметил, что это на самом деле WrappedArray, который является дочерним по отношению к IndexedSeqLike, так что выполняет binarySearch (а это не то, что я хочу).

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