Skip to content

[SPARK-47563][SQL] Add map normalization on creation #45721

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
stevomitric wants to merge 7 commits into from
Closed

[SPARK-47563][SQL] Add map normalization on creation #45721

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
95aad6e
Added map normalization on creation
stevomitric Mar 26, 2024
9231547
Added a check if a key needs to be normalized
stevomitric Mar 26, 2024
9649c5d
new test with struct of arrays
stevomitric Mar 26, 2024
5f008da
Added test for successful map normalization
stevomitric Mar 26, 2024
cd4bb3d
Changed normalize function to return a lambda function based on data …
stevomitric Mar 26, 2024
15065a1
reverted needNormalize to private scope
stevomitric Mar 26, 2024
e9322d3
refactored normalizer function
stevomitric Mar 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...yst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/NormalizeFloatingNumbers.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ object NormalizeFloatingNumbers extends Rule[LogicalPlan] {
case _ => needNormalize(expr.dataType)
}

private def needNormalize(dt: DataType): Boolean = dt match {
def needNormalize(dt: DataType): Boolean = dt match {
case FloatType | DoubleType => true
case StructType(fields) => fields.exists(f => needNormalize(f.dataType))
case ArrayType(et, _) => needNormalize(et)
Expand Down
26 changes: 23 additions & 3 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ArrayBasedMapBuilder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@ import scala.collection.mutable

import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
import org.apache.spark.sql.catalyst.optimizer.NormalizeFloatingNumbers
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
Expand Down Expand Up @@ -52,18 +54,36 @@ class ArrayBasedMapBuilder(keyType: DataType, valueType: DataType) extends Seria

private val mapKeyDedupPolicy = SQLConf.get.getConf(SQLConf.MAP_KEY_DEDUP_POLICY)

private lazy val keyNeedNormalize = NormalizeFloatingNumbers.needNormalize(keyType)

def normalize(value: Any, dataType: DataType): Any = dataType match {

@cloud-fan cloud-fan Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we should return a lambda function to do normalization based on the data type, instead of matching the data type per row.

case FloatType => NormalizeFloatingNumbers.FLOAT_NORMALIZER(value)
case DoubleType => NormalizeFloatingNumbers.DOUBLE_NORMALIZER(value)
case ArrayType(dt, _) =>

@cloud-fan cloud-fan Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no need to handle complex types, as we use TreeMap for complex types which should handle floating points well.

new GenericArrayData(value.asInstanceOf[GenericArrayData].array.map { element =>

@stefankandic stefankandic Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if we have an array of 1 million strings we will go through each value even though we know we don't need to normalize strings

what about doing the same as in NormalizeFloatingNumbers and first check if we need to perform normalization

@stevomitric stevomitric Mar 26, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Applied NormalizeFloatingNumbers.needNormalize here

normalize(element, dt)
})
case StructType(sf) =>
new GenericInternalRow(
value.asInstanceOf[GenericInternalRow].values.zipWithIndex.map { element =>
normalize(element._1, sf(element._2).dataType)

@stefankandic stefankandic Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you could also check if you need to do normalization here right?

this way we would avoid normalization of all fields of a struct if only one actually needs it

@stevomitric stevomitric Mar 26, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As noted by @cloud-fan below, complex types have been dropped.

})
case _ => value
}

def put(key: Any, value: Any): Unit = {
if (key == null) {
throw QueryExecutionErrors.nullAsMapKeyNotAllowedError()
}

val index = keyToIndex.getOrDefault(key, -1)
val keyNormalized = if (keyNeedNormalize) normalize(key, keyType) else key
val index = keyToIndex.getOrDefault(keyNormalized, -1)
if (index == -1) {
if (size >= ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH) {
throw QueryExecutionErrors.exceedMapSizeLimitError(size)
}
keyToIndex.put(key, values.length)
keys.append(key)
keyToIndex.put(keyNormalized, values.length)
keys.append(keyNormalized)
values.append(value)
} else {
if (mapKeyDedupPolicy == SQLConf.MapKeyDedupPolicy.EXCEPTION.toString) {
Expand Down
56 changes: 55 additions & 1 deletion ...atalyst/src/test/scala/org/apache/spark/sql/catalyst/util/ArrayBasedMapBuilderSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{UnsafeArrayData, UnsafeRow}
import org.apache.spark.sql.catalyst.plans.SQLHelper
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{ArrayType, BinaryType, IntegerType, StructType}
import org.apache.spark.sql.types.{ArrayType, BinaryType, DoubleType, IntegerType, StructType}
import org.apache.spark.unsafe.Platform

class ArrayBasedMapBuilderSuite extends SparkFunSuite with SQLHelper {
Expand Down Expand Up @@ -60,6 +60,60 @@ class ArrayBasedMapBuilderSuite extends SparkFunSuite with SQLHelper {
)
}

test("apply key normalization when creating") {

@stefankandic stefankandic Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add another test for successful normalization

val builderDouble = new ArrayBasedMapBuilder(DoubleType, IntegerType)
builderDouble.put(-0.0, 1)
checkError(
exception = intercept[SparkRuntimeException](builderDouble.put(0.0, 2)),
errorClass = "DUPLICATED_MAP_KEY",
parameters = Map(
"key" -> "0.0",
"mapKeyDedupPolicy" -> "\"spark.sql.mapKeyDedupPolicy\"")
)

val builderArray = new ArrayBasedMapBuilder(ArrayType(DoubleType), IntegerType)
builderArray.put(new GenericArrayData(Seq(-0.0)), 1)
checkError(
exception = intercept[SparkRuntimeException](
builderArray.put(new GenericArrayData(Seq(0.0)), 1)),
errorClass = "DUPLICATED_MAP_KEY",
parameters = Map(
"key" -> "[0.0]",
"mapKeyDedupPolicy" -> "\"spark.sql.mapKeyDedupPolicy\"")
)

val builderStruct = new ArrayBasedMapBuilder(new StructType().add("i", "double"), IntegerType)

@stefankandic stefankandic Mar 26, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe add a case when array is inside of a struct

builderStruct.put(InternalRow(-0.0), 1)
// By default duplicated map key fails the query.
checkError(
exception = intercept[SparkRuntimeException](builderStruct.put(InternalRow(0.0), 3)),
errorClass = "DUPLICATED_MAP_KEY",
parameters = Map(
"key" -> "[0.0]",
"mapKeyDedupPolicy" -> "\"spark.sql.mapKeyDedupPolicy\"")
)

val builderStructWithArray = new ArrayBasedMapBuilder(
new StructType().add("array", ArrayType(DoubleType) ), IntegerType)
builderStructWithArray.put(InternalRow(new GenericArrayData(Seq(-0.0))), 1)
checkError(
exception = intercept[SparkRuntimeException](
builderStructWithArray.put(InternalRow(new GenericArrayData(Seq(0.0))), 1)),
errorClass = "DUPLICATED_MAP_KEY",
parameters = Map(
"key" -> "[[0.0]]",
"mapKeyDedupPolicy" -> "\"spark.sql.mapKeyDedupPolicy\"")
)
}

test("successful map normalization on build") {
val builder = new ArrayBasedMapBuilder(DoubleType, IntegerType)
builder.put(-0.0, 1)
val map = builder.build()
assert(map.numElements() == 1)
assert(ArrayBasedMapData.toScalaMap(map) == Map(0.0 -> 1))
}

test("remove duplicated keys with last wins policy") {
withSQLConf(SQLConf.MAP_KEY_DEDUP_POLICY.key -> SQLConf.MapKeyDedupPolicy.LAST_WIN.toString) {
val builder = new ArrayBasedMapBuilder(IntegerType, IntegerType)
Expand Down