Skip to content

Commit

Permalink
h
Browse files Browse the repository at this point in the history
  • Loading branch information
@GideonPotok
GideonPotok committed Jun 5, 2024
1 parent 52d625b commit 4f50358
Show file tree
Hide file tree
Showing 4 changed files with 42 additions and 208 deletions.
1 change: 1 addition & 0 deletions core/src/main/scala/org/apache/spark/util/collection/OpenHashMap.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ class OpenHashMap[K : ClassTag, @specialized(Long, Int, Double) V: ClassTag](
def this() = this(64)

protected var _keySet = new OpenHashSet[K](initialCapacity, 0.7)

/*
specialCase match {
case -1 => None
Expand Down
36 changes: 28 additions & 8 deletions core/src/main/scala/org/apache/spark/util/collection/OpenHashSet.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ import scala.reflect._
import com.google.common.hash.Hashing

import org.apache.spark.annotation.Private
import org.apache.spark.sql.catalyst.util.CollationFactory
import org.apache.spark.unsafe.types.UTF8String
// import org.apache.spark.sql.catalyst.util.CollationFactory
// import org.apache.spark.unsafe.types.UTF8String

Expand All @@ -46,7 +48,15 @@ import org.apache.spark.annotation.Private
class OpenHashSet[@specialized(Long, Int, Double, Float) T: ClassTag](
initialCapacity: Int,
loadFactor: Double,
specialPassedInHasher: Option[Any => Int] = None
var specialPassedInHasher: Option[Object => Int] = Some(o => {
val i = CollationFactory.fetchCollation(1)
.hashFunction.applyAsLong(o.asInstanceOf[UTF8String])
.toInt
// scalastyle:off println
println(s"Hashing: $o -> $i")
// scalastyle:on println
i
})
)
extends Serializable {

Expand All @@ -72,12 +82,8 @@ class OpenHashSet[@specialized(Long, Int, Double, Float) T: ClassTag](
case ClassTag.Double => new DoubleHasher().asInstanceOf[Hasher[T]]
case ClassTag.Float => new FloatHasher().asInstanceOf[Hasher[T]]
case _ =>
// scalastyle:off println
println("hasher: " + classTag[T].toString())
println("hasher: " + classTag[T].runtimeClass.toString())
// scalastyle:on println
specialPassedInHasher.map(f =>
new CustomHasher(f).asInstanceOf[Hasher[T]]).getOrElse(
new CustomHasher(f.asInstanceOf[Any => Int]).asInstanceOf[Hasher[T]]).getOrElse(
new Hasher[T])
}

Expand Down Expand Up @@ -129,8 +135,22 @@ class OpenHashSet[@specialized(Long, Int, Double, Float) T: ClassTag](
* See: https://issues.apache.org/jira/browse/SPARK-45599
*/
@annotation.nowarn("cat=other-non-cooperative-equals")
private def keyExistsAtPos(k: T, pos: Int) =
_data(pos) equals k
private def keyExistsAtPos(k: T, pos: Int) = {
val isItEqual = classTag[T] match {
case ClassTag.Long => _data(pos) equals k
case ClassTag.Int => _data(pos) equals k
case ClassTag.Double => _data(pos) equals k
case ClassTag.Float => _data(pos) equals k
case _ => _data(pos).asInstanceOf[UTF8String].semanticEquals(k.asInstanceOf[UTF8String], 1)
}
// scalastyle:off println
println("keyExistsAtPos: " + k + " == " + _data(pos) +
" ==" + (k == _data(pos)) + "... equals() " +
(k equals _data(pos)) + "... semanticEquals() " +
isItEqual)
// scalastyle:on println
isItEqual
}

/**
* Add an element to the set. This one differs from add in that it doesn't trigger rehashing.
Expand Down
24 changes: 4 additions & 20 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Mode.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,9 @@ import org.apache.spark.sql.catalyst.analysis.{ExpressionBuilder, UnresolvedWith
import org.apache.spark.sql.catalyst.expressions.{Ascending, Descending, Expression, ExpressionDescription, ImplicitCastInputTypes, SortOrder}
import org.apache.spark.sql.catalyst.trees.UnaryLike
import org.apache.spark.sql.catalyst.types.PhysicalDataType
import org.apache.spark.sql.catalyst.util.{CollationFactory, GenericArrayData, UnsafeRowUtils}
import org.apache.spark.sql.catalyst.util.{GenericArrayData}
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType, StringType}
import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.sql.types.{AbstractDataType, AnyDataType, ArrayType, BooleanType, DataType}
import org.apache.spark.util.collection.OpenHashMap

case class Mode(
Expand Down Expand Up @@ -75,30 +74,15 @@ case class Mode(
if (buffer.isEmpty) {
return null
}
val collationAwareBuffer = child.dataType match {
case c: StringType if
!CollationFactory.fetchCollation(c.collationId).supportsBinaryEquality =>
val collationId = c.collationId
val modeMap = buffer.toSeq.groupMapReduce {
case (k, _) => CollationFactory.getCollationKey(k.asInstanceOf[UTF8String], collationId)
}(x => x)((x, y) => (x._1, x._2 + y._2)).values
modeMap
case complexType if !UnsafeRowUtils.isBinaryStable(child.dataType) =>
val modeMap = buffer.toSeq.groupMapReduce {
case (k, _) => UnsafeRowUtils.getBinaryStableKey(k, child.dataType)
}(x => x)((x, y) => (x._1, x._2 + y._2)).values
modeMap
case _ => buffer
}
reverseOpt.map { reverse =>
val defaultKeyOrdering = if (reverse) {
PhysicalDataType.ordering(child.dataType).asInstanceOf[Ordering[AnyRef]].reverse
} else {
PhysicalDataType.ordering(child.dataType).asInstanceOf[Ordering[AnyRef]]
}
val ordering = Ordering.Tuple2(Ordering.Long, defaultKeyOrdering)
collationAwareBuffer.maxBy { case (key, count) => (count, key) }(ordering)
}.getOrElse(collationAwareBuffer.maxBy(_._2))._1
buffer.maxBy { case (key, count) => (count, key) }(ordering)
}.getOrElse(buffer.maxBy(_._2))._1
}

override def withNewMutableAggBufferOffset(newMutableAggBufferOffset: Int): Mode =
Expand Down
189 changes: 9 additions & 180 deletions sql/core/src/test/scala/org/apache/spark/sql/CollationSQLExpressionsSuite.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,13 @@ import java.text.SimpleDateFormat
import scala.collection.immutable.Seq

import org.apache.spark.{SparkConf, SparkException, SparkIllegalArgumentException, SparkRuntimeException}
import org.apache.spark.sql.catalyst.expressions.Literal
import org.apache.spark.sql.catalyst.expressions.aggregate.Mode
// import org.apache.spark.sql.catalyst.expressions.Literal
// import org.apache.spark.sql.catalyst.expressions.aggregate.Mode
import org.apache.spark.sql.internal.{SqlApiConf, SQLConf}
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.collection.OpenHashMap
// import org.apache.spark.unsafe.types.UTF8String
// import org.apache.spark.util.collection.OpenHashMap

// scalastyle:off nonascii
class CollationSQLExpressionsSuite
Expand Down Expand Up @@ -1650,12 +1650,13 @@ class CollationSQLExpressionsSuite
test("Support mode for string expression with collation - Advanced Test") {
case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R)
val testCases = Seq(
ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("utf8_binary_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a")
// ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("utf8_binary_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b")
// ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
// ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a")
)
testCases.foreach(t => {
println(s" TEST CASE: ${t}")
val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) =>
(0L to numRepeats).map(_ => s"('$elt')").mkString(",")
}.mkString(",")
Expand All @@ -1672,178 +1673,6 @@ class CollationSQLExpressionsSuite
})
}

test("Support Mode.eval(buffer)") {
case class UTF8StringModeTestCase[R](
collationId: String,
bufferValues: Map[UTF8String, Long],
result: R)

val bufferValuesUTF8String = Map(
UTF8String.fromString("a") -> 5L,
UTF8String.fromString("b") -> 4L,
UTF8String.fromString("B") -> 3L,
UTF8String.fromString("d") -> 2L,
UTF8String.fromString("e") -> 1L)

val testCasesUTF8String = Seq(
UTF8StringModeTestCase("utf8_binary", bufferValuesUTF8String, "a"),
UTF8StringModeTestCase("utf8_binary_lcase", bufferValuesUTF8String, "b"),
UTF8StringModeTestCase("unicode_ci", bufferValuesUTF8String, "b"),
UTF8StringModeTestCase("unicode", bufferValuesUTF8String, "a"))

testCasesUTF8String.foreach(t => {
val buffer = new OpenHashMap[AnyRef, Long](5)
val myMode = Mode(child = Literal.create("some_column_name", StringType(t.collationId)))
t.bufferValues.foreach { case (k, v) => buffer.update(k, v) }
assert(myMode.eval(buffer).toString.toLowerCase() == t.result.toLowerCase())
})
}

test("Support mode for string expression with collated strings in struct") {
case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R)
val testCases = Seq(
ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("utf8_binary_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b")
)
testCases.foreach(t => {
val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) =>
(0L to numRepeats).map(_ => s"named_struct('f1'," +
s" collate('$elt', '${t.collationId}'), 'f2', 1)").mkString(",")
}.mkString(",")

val tableName = s"t_${t.collationId}_mode_struct"
withTable(tableName) {
sql(s"CREATE TABLE ${tableName}(i STRUCT<f1: STRING COLLATE " +
t.collationId + ", f2: INT>) USING parquet")
sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd)
val query = s"SELECT lower(mode(i).f1) FROM ${tableName}"
if(t.collationId == "") { // || t.collationId == "unicode_ci") {
// Cannot resolve "mode(i)" due to data type mismatch:
// Input to function mode was a complex type with strings collated on non-binary
// collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13;
val params = Seq(("sqlExpr", "\"mode(i)\""),
("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
" fields, which are currently not supported by 'mode'."),
("hint", "")).toMap
checkError(
exception = intercept[AnalysisException] {
sql(query)
},
errorClass = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT",
parameters = params,
queryContext = Array(
ExpectedContext(objectType = "",
objectName = "",
startIndex = 13,
stopIndex = 19,
fragment = "mode(i)")
)
)
} else {
checkAnswer(sql(query), Row(t.result))
}
}
})
}

test("Support mode for string expression with collated strings in recursively nested struct") {
case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R)
val testCases = Seq(
ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("utf8_binary_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b")
)
testCases.foreach(t => {
val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) =>
(0L to numRepeats).map(_ => s"named_struct('f1', " +
s"named_struct('f2', collate('$elt', '${t.collationId}')), 'f3', 1)").mkString(",")
}.mkString(",")

val tableName = s"t_${t.collationId}_mode_nested_struct"
withTable(tableName) {
sql(s"CREATE TABLE ${tableName}(i STRUCT<f1: STRUCT<f2: STRING COLLATE " +
t.collationId + ">, f3: INT>) USING parquet")
sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd)
val query = s"SELECT lower(mode(i).f1.f2) FROM ${tableName}"
if(t.collationId == "utf8_binary_lcase" || t.collationId == "unicode_ci") {
// Cannot resolve "mode(i)" due to data type mismatch:
// Input to function mode was a complex type with strings collated on non-binary
// collations, which is not yet supported.. SQLSTATE: 42K09; line 1 pos 13;
val params = Seq(("sqlExpr", "\"mode(i)\""),
("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
" fields, which are currently not supported by 'mode'."),
("hint", "")).toMap
checkError(
exception = intercept[AnalysisException] {
sql(query)
},
errorClass = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT",
parameters = params,
queryContext = Array(
ExpectedContext(objectType = "",
objectName = "",
startIndex = 13,
stopIndex = 19,
fragment = "mode(i)")
)
)
} else {
checkAnswer(sql(query), Row(t.result))
}
}
})
}

test("Support mode for string expression with collated strings in array complex type") {
case class ModeTestCase[R](collationId: String, bufferValues: Map[String, Long], result: R)
val testCases = Seq(
ModeTestCase("utf8_binary", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("utf8_binary_lcase", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b"),
ModeTestCase("unicode", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "a"),
ModeTestCase("unicode_ci", Map("a" -> 3L, "b" -> 2L, "B" -> 2L), "b")
)
testCases.foreach(t => {
val valuesToAdd = t.bufferValues.map { case (elt, numRepeats) =>
(0L to numRepeats).map(_ => s"array(named_struct('s1', named_struct('a2', " +
s"array(collate('$elt', '${t.collationId}'))), 'f3', 1))").mkString(",")
}.mkString(",")

val tableName = s"t_${t.collationId}_mode_nested_struct"
withTable(tableName) {
sql(s"CREATE TABLE ${tableName}(" +
s"i ARRAY<STRUCT<s1: STRUCT<a2: ARRAY<STRING COLLATE ${t.collationId}>>, f3: INT>>)" +
s" USING parquet")
sql(s"INSERT INTO ${tableName} VALUES " + valuesToAdd)
val query = s"SELECT lower(element_at(element_at(mode(i), 1).s1.a2, 1)) FROM ${tableName}"
if(t.collationId == "utf8_binary_lcase" || t.collationId == "unicode_ci") {
val params = Seq(("sqlExpr", "\"mode(i)\""),
("msg", "The input to the function 'mode' was a complex type with non-binary collated" +
" fields, which are currently not supported by 'mode'."),
("hint", "")).toMap
checkError(
exception = intercept[AnalysisException] {
sql(query)
},
errorClass = "DATATYPE_MISMATCH.TYPE_CHECK_FAILURE_WITH_HINT",
parameters = params,
queryContext = Array(
ExpectedContext(objectType = "",
objectName = "",
startIndex = 35,
stopIndex = 41,
fragment = "mode(i)")
)
)
} else {
checkAnswer(sql(query), Row(t.result))
}
}
})
}

test("CurrentTimeZone expression with collation") {
// Supported collations
testSuppCollations.foreach(collationName => {
Expand Down

0 comments on commit 4f50358

Please sign in to comment.