This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new feed6c18a30a [SPARK-58963][ML][SQL] Add internal ML-specific vector
posexplode function
feed6c18a30a is described below
commit feed6c18a30a9465c8fa579facacef9f8ead45a9
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Aug 25 16:41:54 2026 +0800
[SPARK-58963][ML][SQL] Add internal ML-specific vector posexplode function
### What changes were proposed in this pull request?
This PR adds an internal ML-specific `vector_posexplode` generator
expression for MLlib vector
columns.
The Catalyst expression is registered as an internal expression and
operates on the SQL struct
representation produced by `unwrap_udt`:
```scala
struct<type:tinyint,size:int,indices:array<int>,values:array<double>>
```
The ML-side helper is `private[ml]` and wraps vector UDT columns by calling
`unwrap_udt` before
invoking the internal expression, so it can be used by Spark ML internals
with both
`org.apache.spark.ml.linalg.VectorUDT` and
`org.apache.spark.mllib.linalg.VectorUDT` columns.
The generator emits `(index, value)` rows with two modes:
* `sparse` emits nonzero vector entries.
* `dense` emits every vector entry, including zeros.
For each non-null vector, it first emits a marker row with index `-1 -
vector.size` and value
`Double.NaN`. Null vectors emit no rows.
### Why are the changes needed?
Spark SQL generator functions such as `posexplode` only accept arrays and
maps. Spark ML internals
sometimes need to explode vector columns into index-value rows while
preserving sparse vectors
without densifying through `vector_to_array`.
This adds a vector-aware internal generator that works from the shared
vector SQL representation
and keeps the API surface internal to ML.
### Does this PR introduce _any_ user-facing change?
No. The Scala helper is `private[ml]`, and the Catalyst expression is
registered with
`registerInternalExpression`. This is intended only for Spark ML internals.
### How was this patch tested?
Added coverage in `FunctionsSuite` for:
* `spark.ml` and `spark.mllib` vector UDT columns
* dense vectors
* sparse vectors with explicit zero values
* null vectors
* empty sparse and dense vectors
* `dense` and `sparse` modes
* output schema
Also ran:
```bash
build/sbt "mllib/Test/compile"
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58194 from zhengruifeng/vector-posexplode-dev5.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../main/scala/org/apache/spark/ml/functions.scala | 15 ++
.../scala/org/apache/spark/ml/FunctionsSuite.scala | 80 +++++-
.../sql/catalyst/analysis/FunctionRegistry.scala | 2 +
.../catalyst/expressions/ml/VectorGenerators.scala | 291 +++++++++++++++++++++
4 files changed, 387 insertions(+), 1 deletion(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/functions.scala
b/mllib/src/main/scala/org/apache/spark/ml/functions.scala
index 07db59e53ba7..b98d0523ab51 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/functions.scala
+++ b/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)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
index 4608d9bd683d..c942a9a8c4f4 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/FunctionsSuite.scala
+++ b/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),
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index 4facf16cefbe..ca4225b645fb 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++
b/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(
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala
new file mode 100644
index 000000000000..5d5e8f7829bc
--- /dev/null
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/ml/VectorGenerators.scala
@@ -0,0 +1,291 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions.ml
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
+import
org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch,
TypeCheckFailure}
+import org.apache.spark.sql.catalyst.expressions.{Expression, Generator,
Literal}
+import org.apache.spark.sql.catalyst.expressions.Cast._
+import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
+import org.apache.spark.sql.catalyst.util.ArrayData
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+
+/**
+ * Explodes the SQL struct representation of an MLlib vector into index-value
pairs. This
+ * expression is dedicated only for Spark ML and should be used together with
`unwrap_udt`.
+ * The mode controls whether it emits all entries or nonzero entries. It
always emits a marker
+ * row before each vector for ML computations that need a per-vector row. The
marker index is
+ * `-1 - vector.size`.
+ *
+ * Sparse vector examples:
+ * {{{
+ * // v = {type: 0, size: 4, indices: [1, 3], values: [2.0, 4.0]}
+ * vector_posexplode(v)
+ * index value
+ * -5 NaN
+ * 1 2.0
+ * 3 4.0
+ *
+ * vector_posexplode(v, mode = "dense")
+ * index value
+ * -5 NaN
+ * 0 0.0
+ * 1 2.0
+ * 2 0.0
+ * 3 4.0
+ * }}}
+ *
+ * Dense vector examples:
+ * {{{
+ * // v = {type: 1, size: null, indices: null, values: [1.0, 0.0, 3.0]}
+ * vector_posexplode(v)
+ * index value
+ * -4 NaN
+ * 0 1.0
+ * 2 3.0
+ *
+ * vector_posexplode(v, mode = "dense")
+ * index value
+ * -4 NaN
+ * 0 1.0
+ * 1 0.0
+ * 2 3.0
+ * }}}
+ */
+case class VectorPosExplode(child: Expression, mode: Expression)
+ extends Generator with CodegenFallback {
+
+ def this(child: Expression) = this(child, Literal("sparse"))
+
+ override def children: Seq[Expression] = Seq(child, mode)
+
+ @transient private lazy val vectorMode: VectorPosExplode.VectorMode.Value =
+ VectorPosExplode.toMode(mode.eval().asInstanceOf[UTF8String].toString)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ if (!VectorPosExplode.isVectorType(child.dataType)) {
+ return DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(0),
+ "requiredType" ->
+ toSQLType(s"STRUCT with SQL type
${VectorPosExplode.vectorSqlType.sql}"),
+ "inputSql" -> toSQLExpr(child),
+ "inputType" -> toSQLType(child.dataType)))
+ }
+ if (!mode.foldable || !mode.dataType.isInstanceOf[StringType]) {
+ return DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_INPUT_TYPE",
+ messageParameters = Map(
+ "paramIndex" -> ordinalNumber(1),
+ "requiredType" -> toSQLType("foldable STRING"),
+ "inputSql" -> toSQLExpr(mode),
+ "inputType" -> toSQLType(mode.dataType)))
+ }
+ val modeValue = mode.eval()
+ if (modeValue == null) {
+ return TypeCheckFailure("The second argument of vector_posexplode cannot
be null.")
+ }
+ VectorPosExplode.toModeOption(modeValue.asInstanceOf[UTF8String].toString)
match {
+ case Some(_) =>
+ case None =>
+ return TypeCheckFailure(
+ "The second argument of vector_posexplode must be one of: dense,
sparse.")
+ }
+ TypeCheckResult.TypeCheckSuccess
+ }
+
+ override def elementSchema: StructType = VectorPosExplode.elementSchema
+
+ override def eval(input: InternalRow): IterableOnce[InternalRow] = {
+ val vector = child.eval(input).asInstanceOf[InternalRow]
+ if (vector == null) {
+ Iterator.empty
+ } else {
+ val values = vector.getArray(3)
+ val (size, rows) = vector.getByte(0) match {
+ case VectorPosExplode.SparseVectorType =>
+ val indices = vector.getArray(2)
+ if (indices == null || values == null || vector.isNullAt(1)) {
+ return Iterator.empty
+ }
+ val size = vector.getInt(1)
+ (size, VectorPosExplode.explodeSparse(vectorMode, size, indices,
values))
+ case VectorPosExplode.DenseVectorType =>
+ if (values == null) {
+ return Iterator.empty
+ }
+ (values.numElements(), VectorPosExplode.explodeDense(vectorMode,
values))
+ case vectorType =>
+ throw new IllegalArgumentException(s"Unknown vector type
$vectorType.")
+ }
+ Iterator.single(VectorPosExplode.markerRow(size)) ++ rows
+ }
+ }
+
+ override protected def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): VectorPosExplode = {
+ copy(child = newChildren(0), mode = newChildren(1))
+ }
+}
+
+object VectorPosExplode {
+ object VectorMode extends Enumeration {
+ val Dense, Sparse = Value
+ }
+
+ private val SparseVectorType: Byte = 0
+ private val DenseVectorType: Byte = 1
+
+ private val vectorSqlType = StructType(Array(
+ StructField("type", ByteType, nullable = false),
+ StructField("size", IntegerType, nullable = true),
+ StructField("indices", ArrayType(IntegerType, containsNull = false),
nullable = true),
+ StructField("values", ArrayType(DoubleType, containsNull = false),
nullable = true)))
+
+ private val elementSchema = new StructType()
+ .add("index", IntegerType, nullable = false)
+ .add("value", DoubleType, nullable = false)
+
+ private def isVectorType(dataType: DataType): Boolean = dataType match {
+ case struct: StructType => struct == vectorSqlType
+ case _ => false
+ }
+
+ private def toModeOption(mode: String): Option[VectorMode.Value] = mode
match {
+ case "dense" => Some(VectorMode.Dense)
+ case "sparse" => Some(VectorMode.Sparse)
+ case _ => None
+ }
+
+ private def toMode(mode: String): VectorMode.Value = toModeOption(mode).get
+
+ private def markerRow(size: Int): InternalRow = InternalRow(-1 - size,
Double.NaN)
+
+ private def explodeSparse(
+ mode: VectorMode.Value,
+ size: Int,
+ indices: ArrayData,
+ values: ArrayData): Iterator[InternalRow] = mode match {
+ case VectorMode.Dense =>
+ explodeSparseAsDense(size, indices, values)
+ case VectorMode.Sparse =>
+ explodeSparseNonzero(indices, values)
+ }
+
+ private def explodeSparseAsDense(
+ vectorSize: Int,
+ indices: ArrayData,
+ values: ArrayData): Iterator[InternalRow] = {
+ val numActives = values.numElements()
+ // Mirrors SparseVector.iterator without depending on MLlib from Catalyst.
+ new Iterator[InternalRow] {
+ private var index = 0
+ private var activeIndex = 0
+ private var nextActiveIndex = if (numActives > 0) indices.getInt(0) else
-1
+
+ override def hasNext: Boolean = index < vectorSize
+
+ override def next(): InternalRow = {
+ if (!hasNext) {
+ throw new NoSuchElementException("next on empty iterator")
+ }
+ val value = if (index == nextActiveIndex) {
+ val activeValue = values.getDouble(activeIndex)
+ activeIndex += 1
+ nextActiveIndex = if (activeIndex < numActives)
indices.getInt(activeIndex) else -1
+ activeValue
+ } else {
+ 0.0
+ }
+ val row = InternalRow(index, value)
+ index += 1
+ row
+ }
+ }
+ }
+
+ private def explodeSparseNonzero(
+ indices: ArrayData,
+ values: ArrayData): Iterator[InternalRow] = {
+ val numElements = values.numElements()
+ new Iterator[InternalRow] {
+ private var index = 0
+ private var nextRow: InternalRow = _
+
+ override def hasNext: Boolean = {
+ while (nextRow == null && index < numElements) {
+ val value = values.getDouble(index)
+ if (value != 0.0) {
+ nextRow = InternalRow(indices.getInt(index), value)
+ }
+ index += 1
+ }
+ nextRow != null
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) {
+ throw new NoSuchElementException("next on empty iterator")
+ }
+ val row = nextRow
+ nextRow = null
+ row
+ }
+ }
+ }
+
+ private def explodeDense(mode: VectorMode.Value, values: ArrayData):
Iterator[InternalRow] = {
+ mode match {
+ case VectorMode.Dense =>
+ explodeDense(values, skipZero = false)
+ case VectorMode.Sparse =>
+ explodeDense(values, skipZero = true)
+ }
+ }
+
+ private def explodeDense(values: ArrayData, skipZero: Boolean):
Iterator[InternalRow] = {
+ val numElements = values.numElements()
+ new Iterator[InternalRow] {
+ private var index = 0
+ private var nextRow: InternalRow = _
+
+ override def hasNext: Boolean = {
+ while (nextRow == null && index < numElements) {
+ val value = values.getDouble(index)
+ if (!skipZero || value != 0.0) {
+ nextRow = InternalRow(index, value)
+ }
+ index += 1
+ }
+ nextRow != null
+ }
+
+ override def next(): InternalRow = {
+ if (!hasNext) {
+ throw new NoSuchElementException("next on empty iterator")
+ }
+ val row = nextRow
+ nextRow = null
+ row
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
