[SPARK-58963][ML][SQL] Add internal ML-specific vector posexplode function

mllib/src/main/scala/org/apache/spark/ml/functions.scala

@@ -44,6 +44,21 @@ object functions {
    */
   def array_to_vector(v: Column): Column = Column.internalFn("array_to_vector", v)
 
+  /**
+   * Creates a new row for each index-value pair in the given vector column. This expression is
+   * dedicated only for Spark ML. It always emits a marker row with index `-1 - vector.size` and
+   * value `Double.NaN` before each non-null vector.
+   * @param v: the column of MLlib sparse/dense vectors
+   * @param mode: `dense` emits all elements, and `sparse` emits nonzero elements
+   * @return the index and value columns of the vector elements
+   * @since 4.4.0
+   */
+  private[ml] def vector_posexplode(
+      v: Column,
+      mode: String = "sparse"): Column = {
+    Column.internalFn("vector_posexplode", sf.unwrap_udt(v), sf.lit(mode))
+  }
+
   private[ml] def array_binary_search(a: Column, v: Column): Column =
     Column.internalFn("array_binary_search", a, v)
 

mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala

@@ -22,7 +22,7 @@ import org.apache.spark.ml.functions._
 import org.apache.spark.ml.linalg.{Matrices, MatrixUDT, Vector, Vectors, VectorUDT}
 import org.apache.spark.ml.util.MLTest
 import org.apache.spark.mllib.linalg.{Matrices => OldMatrices, MatrixUDT => OldMatrixUDT,
-  Vectors => OldVectors, VectorUDT => OldVectorUDT}
+  Vector => OldVector, Vectors => OldVectors, VectorUDT => OldVectorUDT}
 import org.apache.spark.sql.{AnalysisException, DataFrame, Row}
 import org.apache.spark.sql.functions.{col, unwrap_udt, wrap_udt}
 import org.apache.spark.sql.types.{StructField, StructType, UserDefinedType}
@@ -63,6 +63,13 @@ class FunctionsSuite extends MLTest {
     assert(converted.collect().map(_.get(0)).toSeq === Seq(expected, null))
   }
 
+  private def normalizeNaN(rows: Seq[(Int, Int, Double)]): Seq[(Int, Int, String)] = {
+    rows.map {
+      case (id, index, value) if value.isNaN => (id, index, "NaN")
+      case (id, index, value) => (id, index, value.toString)
+    }
+  }
+
   test("test vector_to_array") {
     val df = Seq(
       (Vectors.dense(1.0, 2.0, 3.0), OldVectors.dense(10.0, 20.0, 30.0)),
@@ -137,6 +144,77 @@ class FunctionsSuite extends MLTest {
     assert(resultVec3 === Vectors.dense(Array(1.0, 2.0)))
   }
 
+  test("test vector_posexplode with vector UDT") {
+    val df = Seq(
+      (0, Vectors.dense(1.0, 0.0, 3.0), OldVectors.dense(10.0, 0.0, 30.0)),
+      (1, Vectors.sparse(4, Seq((1, 2.0), (2, 0.0), (3, 4.0))),
+        OldVectors.sparse(4, Seq((0, 20.0), (1, 0.0), (2, 30.0)))),
+      (2, null.asInstanceOf[Vector], null.asInstanceOf[OldVector]),
+      (3, Vectors.sparse(10, Array.emptyIntArray, Array.emptyDoubleArray),
+        OldVectors.sparse(10, Array.emptyIntArray, Array.emptyDoubleArray)),
+      (4, Vectors.dense(Array.emptyDoubleArray),
+        OldVectors.dense(Array.emptyDoubleArray))
+    ).toDF("id", "vec", "oldVec")
+
+    val result = df.select($"id", vector_posexplode($"vec"))
+      .as[(Int, Int, Double)]
+      .collect()
+      .toSeq
+    assert(normalizeNaN(result) === Seq(
+      (0, -4, "NaN"),
+      (0, 0, "1.0"),
+      (0, 2, "3.0"),
+      (1, -5, "NaN"),
+      (1, 1, "2.0"),
+      (1, 3, "4.0"),
+      (3, -11, "NaN"),
+      (4, -1, "NaN")))
+
+    val oldResult = df.select($"id", vector_posexplode($"oldVec"))
+      .as[(Int, Int, Double)]
+      .collect()
+      .toSeq
+    assert(normalizeNaN(oldResult) === Seq(
+      (0, -4, "NaN"),
+      (0, 0, "10.0"),
+      (0, 2, "30.0"),
+      (1, -5, "NaN"),
+      (1, 0, "20.0"),
+      (1, 2, "30.0"),
+      (3, -11, "NaN"),
+      (4, -1, "NaN")))
+
+    val denseResult = df
+      .where($"id" === 1)
+      .select($"id", vector_posexplode($"vec", mode = "dense"))
+      .as[(Int, Int, Double)]
+      .collect()
+      .toSeq
+    assert(normalizeNaN(denseResult) === Seq(
+      (1, -5, "NaN"),
+      (1, 0, "0.0"),
+      (1, 1, "2.0"),
+      (1, 2, "0.0"),
+      (1, 3, "4.0")))
+
+    val sparseResult = df.select($"id", vector_posexplode($"vec", mode = "sparse"))
+      .as[(Int, Int, Double)]
+      .collect()
+      .toSeq
+    assert(normalizeNaN(sparseResult) === Seq(
+      (0, -4, "NaN"),
+      (0, 0, "1.0"),
+      (0, 2, "3.0"),
+      (1, -5, "NaN"),
+      (1, 1, "2.0"),
+      (1, 3, "4.0"),
+      (3, -11, "NaN"),
+      (4, -1, "NaN")))
+
+    val schema = df.select(vector_posexplode($"vec")).schema
+    assert(schema.simpleString === "struct<index:int,value:double>")
+  }
+
   test("test get_vector") {
     val df = Seq(
       (Vectors.dense(1.0, 2.0, 3.0), 0),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -30,6 +30,7 @@ import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.FunctionIdentifier
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
+import org.apache.spark.sql.catalyst.expressions.ml._
 import org.apache.spark.sql.catalyst.expressions.st._
 import org.apache.spark.sql.catalyst.expressions.variant._
 import org.apache.spark.sql.catalyst.expressions.xml._
@@ -1137,6 +1138,7 @@ object FunctionRegistry {
   registerInternalExpression[NullIndex]("null_index")
   registerInternalExpression[CastTimestampNTZToLong]("timestamp_ntz_to_long")
   registerInternalExpression[ArrayBinarySearch]("array_binary_search")
+  registerInternalExpression[VectorPosExplode]("vector_posexplode")
 
   private def makeExprInfoForVirtualOperator(name: String, usage: String): ExpressionInfo = {
     new ExpressionInfo(