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 07db59e53ba7d..b98d0523ab510 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 4608d9bd683dc..c942a9a8c4f4a 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") + } + 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 4facf16cefbe9..ca4225b645fb1 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 0000000000000..5d5e8f7829bc6 --- /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 + } + } + } +}