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
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 all 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
14 changes: 11 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,7 @@ import scala.collection.mutable

import org.apache.spark.SparkException
import org.apache.spark.sql.catalyst.InternalRow
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 +53,25 @@ class ArrayBasedMapBuilder(keyType: DataType, valueType: DataType) extends Seria

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

private lazy val keyNormalizer: Any => Any = keyType match {
case FloatType => NormalizeFloatingNumbers.FLOAT_NORMALIZER
case DoubleType => NormalizeFloatingNumbers.DOUBLE_NORMALIZER
case _ => identity
}

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

val index = keyToIndex.getOrDefault(key, -1)
val keyNormalized = keyNormalizer(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
22 changes: 21 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,26 @@ class ArrayBasedMapBuilderSuite extends SparkFunSuite with SQLHelper {
)
}

test("apply key normalization when creating") {
Copy link
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\"")
)
}

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