[WIP][ML] Update vector posexplode semantics

Author: zhengruifeng

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

@@ -46,19 +46,17 @@ object functions {
 
   /**
    * Creates a new row for each index-value pair in the given vector column. This expression is
-   * dedicated only for Spark ML.
+   * 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, `active` emits active elements, and `nonzero` emits
-   *              active elements whose values are not zero
-   * @param withMarker: whether to emit `(-1, Double.NaN)` before each vector's values
+   * @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 = "active",
-      withMarker: Boolean = false): Column = {
-    Column.internalFn("vector_posexplode", v, sf.lit(mode), sf.lit(withMarker))
+      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 =

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

@@ -22,8 +22,8 @@ import org.apache.spark.ml.functions._
 import org.apache.spark.ml.linalg.{Vector, Vectors}
 import org.apache.spark.ml.util.MLTest
 import org.apache.spark.mllib.linalg.{Vector => OldVector, Vectors => OldVectors}
-import org.apache.spark.sql.{AnalysisException, Row}
-import org.apache.spark.sql.functions.col
+import org.apache.spark.sql.{AnalysisException, Column, Row}
+import org.apache.spark.sql.functions.{col, lit}
 import org.apache.spark.sql.types._
 
 class FunctionsSuite extends MLTest {
@@ -125,61 +125,53 @@ class FunctionsSuite extends MLTest {
       .as[(Int, Int, Double)]
       .collect()
       .toSeq
-    assert(result === Seq(
-      (0, 0, 1.0),
-      (0, 1, 0.0),
-      (0, 2, 3.0),
-      (1, 1, 2.0),
-      (1, 2, 0.0),
-      (1, 3, 4.0)))
+    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")))
 
     val oldResult = df.select($"id", vector_posexplode($"oldVec"))
       .as[(Int, Int, Double)]
       .collect()
       .toSeq
-    assert(oldResult === Seq(
-      (0, 0, 10.0),
-      (0, 1, 0.0),
-      (0, 2, 30.0),
-      (1, 0, 20.0),
-      (1, 1, 0.0),
-      (1, 2, 30.0)))
+    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")))
 
     val denseResult = df
       .where($"id" === 1)
       .select($"id", vector_posexplode($"vec", mode = "dense"))
       .as[(Int, Int, Double)]
       .collect()
       .toSeq
-    assert(denseResult === Seq(
-      (1, 0, 0.0),
-      (1, 1, 2.0),
-      (1, 2, 0.0),
-      (1, 3, 4.0)))
+    assert(normalizeNaN(denseResult) === Seq(
+      (1, -5, "NaN"),
+      (1, 0, "0.0"),
+      (1, 1, "2.0"),
+      (1, 2, "0.0"),
+      (1, 3, "4.0")))
 
-    val nonzeroResult = df.select($"id", vector_posexplode($"vec", mode = "nonzero"))
+    val sparseResult = df.select($"id", vector_posexplode($"vec", mode = "sparse"))
       .as[(Int, Int, Double)]
       .collect()
       .toSeq
-    assert(nonzeroResult === Seq(
-      (0, 0, 1.0),
-      (0, 2, 3.0),
-      (1, 1, 2.0),
-      (1, 3, 4.0)))
-
-    val markerResult = df.select($"id",
-        vector_posexplode($"vec", mode = "nonzero", withMarker = true))
-      .as[(Int, Int, Double)]
-      .collect()
-      .toSeq
-    assert(normalizeNaN(markerResult) === Seq(
-      (0, -1, "NaN"),
+    assert(normalizeNaN(sparseResult) === Seq(
+      (0, -4, "NaN"),
       (0, 0, "1.0"),
       (0, 2, "3.0"),
-      (1, -1, "NaN"),
+      (1, -5, "NaN"),
       (1, 1, "2.0"),
       (1, 3, "4.0"),
-      (3, -1, "NaN")))
+      (3, -11, "NaN")))
 
     val schema = df.select(vector_posexplode($"vec")).schema
     assert(schema.simpleString === "struct<index:int,value:double>")
@@ -197,20 +189,12 @@ class FunctionsSuite extends MLTest {
       StructType(Seq(
         StructField("id", IntegerType, nullable = false),
         StructField("vec", vectorSqlType, nullable = true))))
-    val structResult = structDf.select($"id", vector_posexplode($"vec"))
-      .as[(Int, Int, Double)]
-      .collect()
-      .toSeq
-    assert(structResult === Seq((0, 0, 5.0), (0, 2, 0.0)))
-    val structMarkerResult = structDf.select($"id",
-        vector_posexplode($"vec", withMarker = true))
+    val result = structDf
+      .select($"id", Column.internalFn("vector_posexplode", $"vec", lit("sparse")))
       .as[(Int, Int, Double)]
       .collect()
       .toSeq
-    assert(normalizeNaN(structMarkerResult) === Seq(
-      (0, -1, "NaN"),
-      (0, 0, "5.0"),
-      (0, 2, "0.0")))
+    assert(normalizeNaN(result) === Seq((0, -4, "NaN"), (0, 0, "5.0")))
   }
 
   test("test get_vector") {