You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "yhuang-db (via GitHub)" <gi...@apache.org> on 2025/07/15 23:17:27 UTC
[PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
yhuang-db opened a new pull request, #51505:
URL: https://github.com/apache/spark/pull/51505
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430764055
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -317,8 +322,51 @@ object ApproxTopK {
def getSketchStateDataType(itemDataType: DataType): StructType =
StructType(
StructField("sketch", BinaryType, nullable = false) ::
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("itemDataTypeDDL", StringType, nullable = false) :: Nil)
+
+ def dataTypeToDDL(dataType: DataType): String = dataType match {
+ case _: StringType =>
+ // Hide collation information in DDL format
Review Comment:
It fails with the following error in [CollationExpressionWalkerSuite.scala:674](https://github.com/apache/spark/blob/master/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala#L674), where it asserts the utf8BinaryResult should match with utf8LcaseResult.
> [info] - SPARK-48280: Expression Walker for codeGen generation *** FAILED *** (38 seconds, 60 milliseconds)
> [info] ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string collate utf8_binary not null}") did not equal ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string collate utf8_lcase not null}") (CollationExpressionWalkerSuite.scala:674)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on PR #51505:
URL: https://github.com/apache/spark/pull/51505#issuecomment-3417737072
Thanks, merging to master
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2462104886
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
+ def checkMixedTypeError(mixedTypeSeq: Seq[(DataType, String, Seq[Any])]): Unit = {
+ for (i <- 0 until mixedTypeSeq.size - 1) {
+ for (j <- i + 1 until mixedTypeSeq.size) {
+ val (type1, _, seq1) = mixedTypeSeq(i)
+ val (type2, _, seq2) = mixedTypeSeq(j)
+ setupMixedTypeAccumulation(seq1, seq2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(type2))
+ )
+ }
+ }
+ }
+
+ checkMixedTypeError(mixedNumberTypes)
+ checkMixedTypeError(mixedDateTimeTypes)
+ }
+
+ // enumerate all combinations of number and datetime types
+ gridTest("SPARK-52798: number vs datetime - fail on UNION")(
+ for {
+ (type1, typeName1, seq1) <- mixedNumberTypes
+ (type2, typeName2, seq2) <- mixedDateTimeTypes
+ } yield ((type1, typeName1, seq1), (type2, typeName2, seq2))) {
+ case ((_, type1, seq1), (_, type2, seq2)) =>
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type2 + ", itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs string - fail at combine")(mixedNumberTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs boolean - fail at UNION")(mixedNumberTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs string - fail at combine")(mixedDateTimeTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs boolean - fail at UNION")(mixedDateTimeTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ test("SPARK-52798: string vs boolean - fail at combine") {
Review Comment:
I verified the failure manually.
```
$ SPARK_ANSI_SQL_MODE=false build/sbt "sql/testOnly *.ApproxTopKSuite -- -z SPARK-52798"
...
[info] - SPARK-52798: string vs boolean - fail at combine *** FAILED *** (10 milliseconds)
[info] org.apache.spark.sql.catalyst.ExtendedAnalysisException: [INCOMPATIBLE_COLUMN_TYPE] UNION can only be performed on tables with compatible column types. The first column of the second table is "STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>" type which is not compatible with "STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, itemDataType: STRING, itemDataTypeDDL: STRING NOT NULL>" at the same column of the first table.. SQLSTATE: 42825; line 1 pos 0;
...
[info] *** 1 TEST FAILED ***
[error] Failed tests:
[error] org.apache.spark.sql.ApproxTopKSuite
[error] (sql / Test / testOnly) sbt.TestsFailedException: Tests unsuccessful
[error] Total time: 34 s, completed Oct 24, 2025, 2:54:17 PM
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2427250444
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(2)
+ val typeCode = inputState.getBinary(3)
+ val actualItemDataType = ApproxTopK.bytesToDataType(typeCode)
+ buffer.setItemDataType(actualItemDataType)
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ if (!combineSizeSpecified) {
+ buffer.setMaxItemsTracked(inputMaxItemsTracked)
+ }
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
+ : CombineInternal[Any] = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (buffer.getMaxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer is a placeholder sketch, set it to the input sketch's max items tracked
+ buffer.setMaxItemsTracked(input.getMaxItemsTracked)
+ }
+ if (buffer.getMaxItemsTracked != input.getMaxItemsTracked) {
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ buffer.getMaxItemsTracked,
+ input.getMaxItemsTracked
+ )
+ }
+ }
+ // check item data type
+ if (buffer.getItemDataType != null && input.getItemDataType != null &&
+ buffer.getItemDataType != input.getItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ buffer.getItemDataType,
+ input.getItemDataType
+ )
+ } else if (buffer.getItemDataType == null) {
+ // If buffer is a placeholder sketch, set it to the input sketch's item data type
+ buffer.setItemDataType(input.getItemDataType)
+ }
+ buffer.getSketch.merge(input.getSketch)
+ buffer
+ }
+
+ override def eval(buffer: CombineInternal[Any]): Any = {
+ val sketchBytes =
+ buffer.getSketch.toByteArray(ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTracked = buffer.getMaxItemsTracked
+ val typeCode = ApproxTopK.dataTypeToBytes(buffer.getItemDataType)
+ InternalRow.apply(sketchBytes, null, maxItemsTracked, typeCode)
+ }
+
+ override def serialize(buffer: CombineInternal[Any]): Array[Byte] = {
Review Comment:
Thanks for the comment. Now use java.nio.ByteBuffer for ser/deSer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430785760
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
Review Comment:
QQ: are there test cases for combine with different input types. For example, int and long, date and string
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430780027
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +491,283 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+/**
+ * In internal class used as the aggregation buffer for ApproxTopKCombine.
+ *
+ * @param sketch the ItemsSketch instance
+ * @param itemDataType the data type of items in the sketch
+ * @param maxItemsTracked the maximum number of items tracked in the sketch
+ */
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def updateMaxItemsTracked(combineSizeSpecified: Boolean, newMaxItemsTracked: Int): Unit = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (this.maxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer's maxItemsTracked VOID_MAX_ITEMS_TRACKED, it means the buffer is a placeholder
+ // sketch that has not beed updated by any input sketch yet.
+ // So we can set it to the input sketch's max items tracked.
+ this.maxItemsTracked = newMaxItemsTracked
+ } else {
+ if (this.maxItemsTracked != newMaxItemsTracked) {
+ // If buffer's maxItemsTracked is not VOID_MAX_ITEMS_TRACKED, it means the buffer has been
+ // updated by some input sketch. So if buffer and input sketch have different
+ // maxItemsTracked values, it means at least two of the input sketches have different
+ // maxItemsTracked values. In this case, we should throw an error.
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ this.maxItemsTracked, newMaxItemsTracked)
+ }
+ }
+ }
+ }
+
+ def updateItemDataType(inputItemDataType: DataType): Unit = {
+ // When the buffer's dataType hasn't been set, set it to the input sketch's item data type
+ // When input sketch's item data type is null, buffer's item data type will remain null
+ if (this.itemDataType == null) {
+ this.itemDataType = inputItemDataType
+ } else {
+ // When the buffer's dataType has been set, throw an error
+ // if the input sketch's item data type is not null the two data types don't match
+ if (inputItemDataType != null && this.itemDataType != inputItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ this.itemDataType, inputItemDataType)
+ }
+ }
+ }
+
+ /**
+ * Serialize the CombineInternal instance to a byte array.
+ * Serialization format:
+ * maxItemsTracked (4 bytes int) +
+ * itemDataTypeDDL length n in byte (4 bytes int) +
+ * itemDataTypeDDL (n bytes) +
+ * sketchBytes
+ */
+ def serialize(): Array[Byte] = {
+ val sketchBytes = sketch.toByteArray(
+ ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]])
+ val itemDataTypeDDL = ApproxTopK.dataTypeToDDL(itemDataType)
+ val ddlBytes: Array[Byte] = itemDataTypeDDL.getBytes(StandardCharsets.UTF_8)
+ val byteArray = new Array[Byte](sketchBytes.length + 4 + 4 + ddlBytes.length)
+
+ val byteBuffer = ByteBuffer.wrap(byteArray)
+ byteBuffer.putInt(maxItemsTracked)
+ byteBuffer.putInt(ddlBytes.length)
+ byteBuffer.put(ddlBytes)
+ byteBuffer.put(sketchBytes)
+ byteArray
+ }
+}
+
+object CombineInternal {
+ /**
+ * Deserialize a byte array to a CombineInternal instance.
+ * Serialization format:
+ * maxItemsTracked (4 bytes int) +
+ * itemDataTypeDDL length n in byte (4 bytes int) +
+ * itemDataTypeDDL (n bytes) +
+ * sketchBytes
+ */
+ def deserialize(buffer: Array[Byte]): CombineInternal[Any] = {
+ val byteBuffer = ByteBuffer.wrap(buffer)
+ // read maxItemsTracked
+ val maxItemsTracked = byteBuffer.getInt
+ // read itemDataTypeDDL
+ val ddlLength = byteBuffer.getInt
+ val ddlBytes = new Array[Byte](ddlLength)
+ byteBuffer.get(ddlBytes)
+ val itemDataTypeDDL = new String(ddlBytes, StandardCharsets.UTF_8)
+ val itemDataType = ApproxTopK.DDLToDataType(itemDataTypeDDL)
+ // read sketchBytes
+ val sketchBytes = new Array[Byte](buffer.length - 4 - 4 - ddlLength)
+ byteBuffer.get(sketchBytes)
+ val sketch = ItemsSketch.getInstance(
+ Memory.wrap(sketchBytes), ApproxTopK.genSketchSerDe(itemDataType))
+ new CombineInternal[Any](sketch, itemDataType, maxItemsTracked)
+ }
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is specified, it will be set for the combined sketch. If maxItemsTracked is not specified, the input sketches must have the same maxItemsTracked value, otherwise an error will be thrown. The output sketch will use the same value from the input sketches.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ // If maxItemsTracked is not specified, set it to VOID_MAX_ITEMS_TRACKED.
+ // This indicates that there is no explicit maxItemsTracked input from the function call.
+ // Hence, function needs to check the input sketches' maxItemsTracked values during merge.
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ // The item data type extracted from the third field of the state struct.
+ // It is named "unchecked" because it may be inaccurate when input sketches have different
+ // item data types. For example, if one sketch has int type null and another has string type
+ // null, the union of the two sketches will have bigint type null.
+ // The accurate item data type will be tracked in the aggregation buffer during update/merge.
+ // It is okay to use uncheckedItemDataType to create the output data type of this function,
+ // because if the input sketches have different item data types, an error will be thrown
+ // during update/merge. Otherwise, the uncheckedItemDataType is accurate.
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](2).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ /**
+ * If maxItemsTracked is specified in function call, use it for the output sketch.
+ * Otherwise, create a placeholder sketch with VOID_MAX_ITEMS_TRACKED. The actual value will be
+ * decided during the first update.
+ */
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ // If maxItemsTracked is not specified, create a sketch with the maximum allowed size.
+ // No need to worry about memory waste, as the sketch always grows from a small init size.
+ // The actual maxItemsTracked will be checked during the updates.
+ val maxMapSize = ApproxTopK.calMaxMapSize(ApproxTopK.MAX_ITEMS_TRACKED_LIMIT)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ /**
+ * Update the aggregation buffer with an input sketch. The input has the same schema as the
+ * ApproxTopKAccumulate output, i.e., sketchBytes + maxItemsTracked + null + DDL.
+ */
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(1)
+ val inputItemDataTypeDDL = inputState.getUTF8String(3).toString
+ val inputItemDataType = ApproxTopK.DDLToDataType(inputItemDataTypeDDL)
+ // update maxItemsTracked (throw error if not match)
+ buffer.updateMaxItemsTracked(combineSizeSpecified, inputMaxItemsTracked)
+ // update itemDataType (throw error if not match)
+ buffer.updateItemDataType(inputItemDataType)
+ // update sketch
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
Review Comment:
```suggestion
override def merge(
buffer: CombineInternal[Any],
input: CombineInternal[Any]): CombineInternal[Any] = {```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang closed pull request #51505: [SPARK-52798] [SQL] Add function approx_top_k_combine
URL: https://github.com/apache/spark/pull/51505
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "cboumalh (via GitHub)" <gi...@apache.org>.
cboumalh commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2229725214
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(2)
+ val typeCode = inputState.getBinary(3)
+ val actualItemDataType = ApproxTopK.bytesToDataType(typeCode)
+ buffer.setItemDataType(actualItemDataType)
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ if (!combineSizeSpecified) {
+ buffer.setMaxItemsTracked(inputMaxItemsTracked)
+ }
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
+ : CombineInternal[Any] = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (buffer.getMaxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer is a placeholder sketch, set it to the input sketch's max items tracked
+ buffer.setMaxItemsTracked(input.getMaxItemsTracked)
+ }
+ if (buffer.getMaxItemsTracked != input.getMaxItemsTracked) {
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ buffer.getMaxItemsTracked,
+ input.getMaxItemsTracked
+ )
+ }
+ }
+ // check item data type
+ if (buffer.getItemDataType != null && input.getItemDataType != null &&
+ buffer.getItemDataType != input.getItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ buffer.getItemDataType,
+ input.getItemDataType
+ )
+ } else if (buffer.getItemDataType == null) {
+ // If buffer is a placeholder sketch, set it to the input sketch's item data type
+ buffer.setItemDataType(input.getItemDataType)
+ }
+ buffer.getSketch.merge(input.getSketch)
+ buffer
+ }
+
+ override def eval(buffer: CombineInternal[Any]): Any = {
+ val sketchBytes =
+ buffer.getSketch.toByteArray(ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTracked = buffer.getMaxItemsTracked
+ val typeCode = ApproxTopK.dataTypeToBytes(buffer.getItemDataType)
+ InternalRow.apply(sketchBytes, null, maxItemsTracked, typeCode)
+ }
+
+ override def serialize(buffer: CombineInternal[Any]): Array[Byte] = {
Review Comment:
Serialization and deserialization here are a bit concerning. The current implementation has brittle byte-level manipulations with no robust error checking or flexibility. Hardcoded assumptions about byte sizes and type conversions create significant risks for future compatibility and data integrity. I think it may be worth looking into using a ByteBuffer and adding versioning to the sketch buffer, or any other approaches!
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -177,6 +177,8 @@ object ApproxTopK {
val DEFAULT_K: Int = 5
val DEFAULT_MAX_ITEMS_TRACKED: Int = 10000
private val MAX_ITEMS_TRACKED_LIMIT: Int = 1000000
+ val VOID_MAX_ITEMS_TRACKED = -1
Review Comment:
Can we document what exactly this VOID_MAX_ITEMS_TRACKED and SKETCH_SIZE_PLACEHOLDER mean?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(2)
+ val typeCode = inputState.getBinary(3)
+ val actualItemDataType = ApproxTopK.bytesToDataType(typeCode)
+ buffer.setItemDataType(actualItemDataType)
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ if (!combineSizeSpecified) {
+ buffer.setMaxItemsTracked(inputMaxItemsTracked)
+ }
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
+ : CombineInternal[Any] = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (buffer.getMaxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer is a placeholder sketch, set it to the input sketch's max items tracked
+ buffer.setMaxItemsTracked(input.getMaxItemsTracked)
+ }
+ if (buffer.getMaxItemsTracked != input.getMaxItemsTracked) {
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ buffer.getMaxItemsTracked,
+ input.getMaxItemsTracked
+ )
+ }
+ }
+ // check item data type
+ if (buffer.getItemDataType != null && input.getItemDataType != null &&
+ buffer.getItemDataType != input.getItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ buffer.getItemDataType,
+ input.getItemDataType
+ )
+ } else if (buffer.getItemDataType == null) {
+ // If buffer is a placeholder sketch, set it to the input sketch's item data type
+ buffer.setItemDataType(input.getItemDataType)
+ }
+ buffer.getSketch.merge(input.getSketch)
+ buffer
+ }
+
+ override def eval(buffer: CombineInternal[Any]): Any = {
+ val sketchBytes =
+ buffer.getSketch.toByteArray(ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTracked = buffer.getMaxItemsTracked
+ val typeCode = ApproxTopK.dataTypeToBytes(buffer.getItemDataType)
+ InternalRow.apply(sketchBytes, null, maxItemsTracked, typeCode)
+ }
+
+ override def serialize(buffer: CombineInternal[Any]): Array[Byte] = {
+ val sketchBytes = buffer.getSketch.toByteArray(
+ ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTrackedByte = buffer.getMaxItemsTracked.toByte
Review Comment:
I believe this is casting an `Int` as a `Byte`. Could this cause any overflow?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430822628
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -317,8 +322,51 @@ object ApproxTopK {
def getSketchStateDataType(itemDataType: DataType): StructType =
StructType(
StructField("sketch", BinaryType, nullable = false) ::
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("itemDataTypeDDL", StringType, nullable = false) :: Nil)
+
+ def dataTypeToDDL(dataType: DataType): String = dataType match {
+ case _: StringType =>
+ // Hide collation information in DDL format
Review Comment:
There is a short skip list https://github.com/apache/spark/blob/master/sql/core/src/test/scala/org/apache/spark/sql/collation/CollationExpressionWalkerSuite.scala#L375. Or maybe we can add accumulate and combine to the list.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2426974888
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(2)
+ val typeCode = inputState.getBinary(3)
+ val actualItemDataType = ApproxTopK.bytesToDataType(typeCode)
+ buffer.setItemDataType(actualItemDataType)
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ if (!combineSizeSpecified) {
+ buffer.setMaxItemsTracked(inputMaxItemsTracked)
+ }
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
+ : CombineInternal[Any] = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (buffer.getMaxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer is a placeholder sketch, set it to the input sketch's max items tracked
+ buffer.setMaxItemsTracked(input.getMaxItemsTracked)
+ }
+ if (buffer.getMaxItemsTracked != input.getMaxItemsTracked) {
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ buffer.getMaxItemsTracked,
+ input.getMaxItemsTracked
+ )
+ }
+ }
+ // check item data type
+ if (buffer.getItemDataType != null && input.getItemDataType != null &&
+ buffer.getItemDataType != input.getItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ buffer.getItemDataType,
+ input.getItemDataType
+ )
+ } else if (buffer.getItemDataType == null) {
+ // If buffer is a placeholder sketch, set it to the input sketch's item data type
+ buffer.setItemDataType(input.getItemDataType)
+ }
+ buffer.getSketch.merge(input.getSketch)
+ buffer
+ }
+
+ override def eval(buffer: CombineInternal[Any]): Any = {
+ val sketchBytes =
+ buffer.getSketch.toByteArray(ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTracked = buffer.getMaxItemsTracked
+ val typeCode = ApproxTopK.dataTypeToBytes(buffer.getItemDataType)
+ InternalRow.apply(sketchBytes, null, maxItemsTracked, typeCode)
+ }
+
+ override def serialize(buffer: CombineInternal[Any]): Array[Byte] = {
+ val sketchBytes = buffer.getSketch.toByteArray(
+ ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTrackedByte = buffer.getMaxItemsTracked.toByte
Review Comment:
Thanks for noticing. Fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2417592705
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
Review Comment:
let's add comment about the syntax when maxItemsTracked is not provided
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2417562983
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
+ }
+ }
+
+ override def update(buffer: CombineInternal[Any], input: InternalRow): CombineInternal[Any] = {
+ val inputState = state.eval(input).asInstanceOf[InternalRow]
+ val inputSketchBytes = inputState.getBinary(0)
+ val inputMaxItemsTracked = inputState.getInt(2)
+ val typeCode = inputState.getBinary(3)
+ val actualItemDataType = ApproxTopK.bytesToDataType(typeCode)
+ buffer.setItemDataType(actualItemDataType)
+ val inputSketch = ItemsSketch.getInstance(
+ Memory.wrap(inputSketchBytes), ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ buffer.getSketch.merge(inputSketch)
+ if (!combineSizeSpecified) {
+ buffer.setMaxItemsTracked(inputMaxItemsTracked)
+ }
+ buffer
+ }
+
+ override def merge(buffer: CombineInternal[Any], input: CombineInternal[Any])
+ : CombineInternal[Any] = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (buffer.getMaxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer is a placeholder sketch, set it to the input sketch's max items tracked
+ buffer.setMaxItemsTracked(input.getMaxItemsTracked)
+ }
+ if (buffer.getMaxItemsTracked != input.getMaxItemsTracked) {
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ buffer.getMaxItemsTracked,
+ input.getMaxItemsTracked
+ )
+ }
+ }
+ // check item data type
+ if (buffer.getItemDataType != null && input.getItemDataType != null &&
+ buffer.getItemDataType != input.getItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ buffer.getItemDataType,
+ input.getItemDataType
+ )
+ } else if (buffer.getItemDataType == null) {
+ // If buffer is a placeholder sketch, set it to the input sketch's item data type
+ buffer.setItemDataType(input.getItemDataType)
+ }
+ buffer.getSketch.merge(input.getSketch)
+ buffer
+ }
+
+ override def eval(buffer: CombineInternal[Any]): Any = {
+ val sketchBytes =
+ buffer.getSketch.toByteArray(ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTracked = buffer.getMaxItemsTracked
+ val typeCode = ApproxTopK.dataTypeToBytes(buffer.getItemDataType)
+ InternalRow.apply(sketchBytes, null, maxItemsTracked, typeCode)
+ }
+
+ override def serialize(buffer: CombineInternal[Any]): Array[Byte] = {
+ val sketchBytes = buffer.getSketch.toByteArray(
+ ApproxTopK.genSketchSerDe(buffer.getItemDataType))
+ val maxItemsTrackedByte = buffer.getMaxItemsTracked.toByte
Review Comment:
Nice catch. This will cause overflow.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2462102987
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
+ def checkMixedTypeError(mixedTypeSeq: Seq[(DataType, String, Seq[Any])]): Unit = {
+ for (i <- 0 until mixedTypeSeq.size - 1) {
+ for (j <- i + 1 until mixedTypeSeq.size) {
+ val (type1, _, seq1) = mixedTypeSeq(i)
+ val (type2, _, seq2) = mixedTypeSeq(j)
+ setupMixedTypeAccumulation(seq1, seq2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(type2))
+ )
+ }
+ }
+ }
+
+ checkMixedTypeError(mixedNumberTypes)
+ checkMixedTypeError(mixedDateTimeTypes)
+ }
+
+ // enumerate all combinations of number and datetime types
+ gridTest("SPARK-52798: number vs datetime - fail on UNION")(
+ for {
+ (type1, typeName1, seq1) <- mixedNumberTypes
+ (type2, typeName2, seq2) <- mixedDateTimeTypes
+ } yield ((type1, typeName1, seq1), (type2, typeName2, seq2))) {
+ case ((_, type1, seq1), (_, type2, seq2)) =>
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type2 + ", itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs string - fail at combine")(mixedNumberTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs boolean - fail at UNION")(mixedNumberTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs string - fail at combine")(mixedDateTimeTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs boolean - fail at UNION")(mixedDateTimeTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ test("SPARK-52798: string vs boolean - fail at combine") {
Review Comment:
Hi, @yhuang-db and @gengliangwang . This test case seems to break non-ANSI CI for a week.
- https://github.com/apache/spark/actions/workflows/build_non_ansi.yml
- https://github.com/apache/spark/actions/runs/18623628999/job/53098232441
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2462104886
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
+ def checkMixedTypeError(mixedTypeSeq: Seq[(DataType, String, Seq[Any])]): Unit = {
+ for (i <- 0 until mixedTypeSeq.size - 1) {
+ for (j <- i + 1 until mixedTypeSeq.size) {
+ val (type1, _, seq1) = mixedTypeSeq(i)
+ val (type2, _, seq2) = mixedTypeSeq(j)
+ setupMixedTypeAccumulation(seq1, seq2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(type2))
+ )
+ }
+ }
+ }
+
+ checkMixedTypeError(mixedNumberTypes)
+ checkMixedTypeError(mixedDateTimeTypes)
+ }
+
+ // enumerate all combinations of number and datetime types
+ gridTest("SPARK-52798: number vs datetime - fail on UNION")(
+ for {
+ (type1, typeName1, seq1) <- mixedNumberTypes
+ (type2, typeName2, seq2) <- mixedDateTimeTypes
+ } yield ((type1, typeName1, seq1), (type2, typeName2, seq2))) {
+ case ((_, type1, seq1), (_, type2, seq2)) =>
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type2 + ", itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs string - fail at combine")(mixedNumberTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs boolean - fail at UNION")(mixedNumberTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs string - fail at combine")(mixedDateTimeTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs boolean - fail at UNION")(mixedDateTimeTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ test("SPARK-52798: string vs boolean - fail at combine") {
Review Comment:
```
$ SPARK_ANSI_SQL_MODE=false build/sbt "sql/testOnly *.ApproxTopKSuite -- -z SPARK-52798"
...
[info] - SPARK-52798: string vs boolean - fail at combine *** FAILED *** (10 milliseconds)
[info] org.apache.spark.sql.catalyst.ExtendedAnalysisException: [INCOMPATIBLE_COLUMN_TYPE] UNION can only be performed on tables with compatible column types. The first column of the second table is "STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>" type which is not compatible with "STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, itemDataType: STRING, itemDataTypeDDL: STRING NOT NULL>" at the same column of the first table.. SQLSTATE: 42825; line 1 pos 0;
...
[info] *** 1 TEST FAILED ***
[error] Failed tests:
[error] org.apache.spark.sql.ApproxTopKSuite
[error] (sql / Test / testOnly) sbt.TestsFailedException: Tests unsuccessful
[error] Total time: 34 s, completed Oct 24, 2025, 2:54:17 PM
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2427269883
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -177,6 +177,8 @@ object ApproxTopK {
val DEFAULT_K: Int = 5
val DEFAULT_MAX_ITEMS_TRACKED: Int = 10000
private val MAX_ITEMS_TRACKED_LIMIT: Int = 1000000
+ val VOID_MAX_ITEMS_TRACKED = -1
Review Comment:
Add comments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "cboumalh (via GitHub)" <gi...@apache.org>.
cboumalh commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430661588
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +491,283 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+/**
+ * In internal class used as the aggregation buffer for ApproxTopKCombine.
+ *
+ * @param sketch the ItemsSketch instance
+ * @param itemDataType the data type of items in the sketch
+ * @param maxItemsTracked the maximum number of items tracked in the sketch
+ */
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def updateMaxItemsTracked(combineSizeSpecified: Boolean, newMaxItemsTracked: Int): Unit = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (this.maxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer's maxItemsTracked VOID_MAX_ITEMS_TRACKED, it means the buffer is a placeholder
+ // sketch that has not beed updated by any input sketch yet.
+ // So we can set it to the input sketch's max items tracked.
+ this.maxItemsTracked = newMaxItemsTracked
+ } else {
+ if (this.maxItemsTracked != newMaxItemsTracked) {
+ // If buffer's maxItemsTracked is not VOID_MAX_ITEMS_TRACKED, it means the buffer has been
+ // updated by some input sketch. So if buffer and input sketch have different
+ // maxItemsTracked values, it means at least two of the input sketches have different
+ // maxItemsTracked values. In this case, we should throw an error.
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ this.maxItemsTracked, newMaxItemsTracked)
+ }
+ }
+ }
+ }
+
+ def updateItemDataType(inputItemDataType: DataType): Unit = {
+ // When the buffer's dataType hasn't been set, set it to the input sketch's item data type
+ // When input sketch's item data type is null, buffer's item data type will remain null
+ if (this.itemDataType == null) {
+ this.itemDataType = inputItemDataType
+ } else {
+ // When the buffer's dataType has been set, throw an error
+ // if the input sketch's item data type is not null the two data types don't match
+ if (inputItemDataType != null && this.itemDataType != inputItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ this.itemDataType, inputItemDataType)
+ }
+ }
+ }
+
+ /**
+ * Serialize the CombineInternal instance to a byte array.
+ * Serialization format:
+ * maxItemsTracked (4 bytes int) +
+ * itemDataTypeDDL length n in byte (4 bytes int) +
+ * itemDataTypeDDL (n bytes) +
+ * sketchBytes
+ */
+ def serialize(): Array[Byte] = {
+ val sketchBytes = sketch.toByteArray(
+ ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]])
+ val itemDataTypeDDL = ApproxTopK.dataTypeToDDL(itemDataType)
+ val ddlBytes: Array[Byte] = itemDataTypeDDL.getBytes(StandardCharsets.UTF_8)
+ val byteArray = new Array[Byte](sketchBytes.length + 4 + 4 + ddlBytes.length)
Review Comment:
super nit, we could consider `Integer.BYTES` instead of 4 here for clarity about what the extra bytes represent. Not necessary to make the change, your call.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430785760
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
Review Comment:
QQ: are there test cases for **valid** combine with different input types. For example, int and long, date and string
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2462123267
##########
sql/core/src/test/scala/org/apache/spark/sql/ApproxTopKSuite.scala:
##########
@@ -328,4 +331,326 @@ class ApproxTopKSuite extends QueryTest with SharedSparkSession {
parameters = Map("maxItemsTracked" -> "5", "k" -> "10")
)
}
+
+ /////////////////////////////////
+ // approx_top_k_combine
+ /////////////////////////////////
+
+ def setupMixedSizeAccumulations(size1: Int, size2: Int): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, $size1) as acc " +
+ "FROM VALUES (0), (0), (0), (1), (1), (2), (2), (3) AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, $size2) as acc " +
+ "FROM VALUES (1), (1), (2), (2), (3), (3), (4), (4) AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ def setupMixedTypeAccumulation(seq1: Seq[Any], seq2: Seq[Any]): Unit = {
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq1.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+
+ sql(s"SELECT approx_top_k_accumulate(expr, 10) as acc " +
+ s"FROM VALUES ${seq2.mkString(", ")} AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+
+ sql("SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2")
+ .createOrReplaceTempView("unioned")
+ }
+
+ val mixedNumberTypes: Seq[(DataType, String, Seq[Any])] = Seq(
+ (IntegerType, "INT",
+ Seq(0, 0, 0, 1, 1, 2, 2, 3)),
+ (ByteType, "TINYINT",
+ Seq("cast(0 AS BYTE)", "cast(0 AS BYTE)", "cast(1 AS BYTE)")),
+ (ShortType, "SMALLINT",
+ Seq("cast(0 AS SHORT)", "cast(0 AS SHORT)", "cast(1 AS SHORT)")),
+ (LongType, "BIGINT",
+ Seq("cast(0 AS LONG)", "cast(0 AS LONG)", "cast(1 AS LONG)")),
+ (FloatType, "FLOAT",
+ Seq("cast(0 AS FLOAT)", "cast(0 AS FLOAT)", "cast(1 AS FLOAT)")),
+ (DoubleType, "DOUBLE",
+ Seq("cast(0 AS DOUBLE)", "cast(0 AS DOUBLE)", "cast(1 AS DOUBLE)")),
+ (DecimalType(4, 2), "DECIMAL(4,2)",
+ Seq("cast(0 AS DECIMAL(4, 2))", "cast(0 AS DECIMAL(4, 2))", "cast(1 AS DECIMAL(4, 2))")),
+ (DecimalType(10, 2), "DECIMAL(10,2)",
+ Seq("cast(0 AS DECIMAL(10, 2))", "cast(0 AS DECIMAL(10, 2))", "cast(1 AS DECIMAL(10, 2))")),
+ (DecimalType(20, 3), "DECIMAL(20,3)",
+ Seq("cast(0 AS DECIMAL(20, 3))", "cast(0 AS DECIMAL(20, 3))", "cast(1 AS DECIMAL(20, 3))"))
+ )
+
+ val mixedDateTimeTypes: Seq[(DataType, String, Seq[String])] = Seq(
+ (DateType, "DATE",
+ Seq("DATE'2025-01-01'", "DATE'2025-01-01'", "DATE'2025-01-02'")),
+ (TimestampType, "TIMESTAMP",
+ Seq("TIMESTAMP'2025-01-01 00:00:00'", "TIMESTAMP'2025-01-01 00:00:00'")),
+ (TimestampNTZType, "TIMESTAMP_NTZ",
+ Seq("TIMESTAMP_NTZ'2025-01-01 00:00:00'", "TIMESTAMP_NTZ'2025-01-01 00:00:00'")
+ )
+ )
+
+ // positive tests for approx_top_k_combine on every types
+ gridTest("SPARK-52798: same type, same size, specified combine size - success")(itemsWithTopK) {
+ case (input, expected) =>
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation1")
+ sql(s"SELECT approx_top_k_accumulate(expr) AS acc FROM VALUES $input AS tab(expr);")
+ .createOrReplaceTempView("accumulation2")
+ sql("SELECT approx_top_k_combine(acc, 30) as com " +
+ "FROM (SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2);")
+ .createOrReplaceTempView("combined")
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ // expected should be doubled because we combine two identical sketches
+ val expectedDoubled = expected.map {
+ case Row(value: Any, count: Int) => Row(value, count * 2)
+ }
+ checkAnswer(est, Row(expectedDoubled))
+ }
+
+ test("SPARK-52798: same type, same size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, same size, unspecified combine size - success") {
+ setupMixedSizeAccumulations(10, 10)
+
+ sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+ .createOrReplaceTempView("combined")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combined;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, specified combine size - success") {
+ setupMixedSizeAccumulations(10, 20)
+
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned")
+ .createOrReplaceTempView("combination")
+
+ val est = sql("SELECT approx_top_k_estimate(com) FROM combination;")
+ checkAnswer(est, Row(Seq(Row(2, 4), Row(1, 4), Row(0, 3), Row(3, 3), Row(4, 2))))
+ }
+
+ test("SPARK-52798: same type, different size, unspecified combine size - fail") {
+ setupMixedSizeAccumulations(10, 20)
+
+ val comb = sql("SELECT approx_top_k_combine(acc) as com FROM unioned")
+
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ comb.collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_SIZE_NOT_MATCH",
+ parameters = Map("size1" -> "10", "size2" -> "20")
+ )
+ }
+
+ gridTest("SPARK-52798: invalid combine size - fail")(Seq((10, 10), (10, 20))) {
+ case (size1, size2) =>
+ setupMixedSizeAccumulations(size1, size2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 0) as com FROM unioned").collect()
+ },
+ condition = "APPROX_TOP_K_NON_POSITIVE_ARG",
+ parameters = Map("argName" -> "`maxItemsTracked`", "argValue" -> "0")
+ )
+ }
+
+ test("SPARK-52798: among different number or datetime types - fail at combine") {
+ def checkMixedTypeError(mixedTypeSeq: Seq[(DataType, String, Seq[Any])]): Unit = {
+ for (i <- 0 until mixedTypeSeq.size - 1) {
+ for (j <- i + 1 until mixedTypeSeq.size) {
+ val (type1, _, seq1) = mixedTypeSeq(i)
+ val (type2, _, seq2) = mixedTypeSeq(j)
+ setupMixedTypeAccumulation(seq1, seq2)
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(type2))
+ )
+ }
+ }
+ }
+
+ checkMixedTypeError(mixedNumberTypes)
+ checkMixedTypeError(mixedDateTimeTypes)
+ }
+
+ // enumerate all combinations of number and datetime types
+ gridTest("SPARK-52798: number vs datetime - fail on UNION")(
+ for {
+ (type1, typeName1, seq1) <- mixedNumberTypes
+ (type2, typeName2, seq2) <- mixedDateTimeTypes
+ } yield ((type1, typeName1, seq1), (type2, typeName2, seq2))) {
+ case ((_, type1, seq1), (_, type2, seq2)) =>
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type2 + ", itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs string - fail at combine")(mixedNumberTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: number vs boolean - fail at UNION")(mixedNumberTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs string - fail at combine")(mixedDateTimeTypes) {
+ case (type1, _, seq1) =>
+ setupMixedTypeAccumulation(seq1, Seq("'a'", "'b'", "'c'", "'c'", "'c'", "'c'", "'d'", "'d'"))
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT approx_top_k_combine(acc, 30) as com FROM unioned;").collect()
+ },
+ condition = "APPROX_TOP_K_SKETCH_TYPE_NOT_MATCH",
+ parameters = Map("type1" -> toSQLType(type1), "type2" -> toSQLType(StringType))
+ )
+ }
+
+ gridTest("SPARK-52798: datetime vs boolean - fail at UNION")(mixedDateTimeTypes) {
+ case (_, type1, seq1) =>
+ val seq2 = Seq("(true)", "(true)", "(false)", "(false)")
+ checkError(
+ exception = intercept[ExtendedAnalysisException] {
+ setupMixedTypeAccumulation(seq1, seq2)
+ },
+ condition = "INCOMPATIBLE_COLUMN_TYPE",
+ parameters = Map(
+ "tableOrdinalNumber" -> "second",
+ "columnOrdinalNumber" -> "first",
+ "dataType2" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: " + type1 + ", itemDataTypeDDL: STRING NOT NULL>\""),
+ "operator" -> "UNION",
+ "hint" -> "",
+ "dataType1" -> ("\"STRUCT<sketch: BINARY NOT NULL, maxItemsTracked: INT NOT NULL, " +
+ "itemDataType: BOOLEAN, itemDataTypeDDL: STRING NOT NULL>\"")
+ ),
+ queryContext = Array(
+ ExpectedContext(
+ "SELECT acc from accumulation1 UNION ALL SELECT acc FROM accumulation2", 0, 68))
+ )
+ }
+
+ test("SPARK-52798: string vs boolean - fail at combine") {
Review Comment:
Here is a follow-up to recover non-ANSI CI.
- https://github.com/apache/spark/pull/52726
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2417560164
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
Review Comment:
Shall we simply set it as `MAX_ITEMS_TRACKED_LIMIT: Int = 1000000` here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430754056
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -317,8 +322,51 @@ object ApproxTopK {
def getSketchStateDataType(itemDataType: DataType): StructType =
StructType(
StructField("sketch", BinaryType, nullable = false) ::
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("itemDataTypeDDL", StringType, nullable = false) :: Nil)
+
+ def dataTypeToDDL(dataType: DataType): String = dataType match {
+ case _: StringType =>
+ // Hide collation information in DDL format
Review Comment:
so when with collation, the fromDDL/toDDL won't work?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430820005
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -317,8 +322,51 @@ object ApproxTopK {
def getSketchStateDataType(itemDataType: DataType): StructType =
StructType(
StructField("sketch", BinaryType, nullable = false) ::
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("itemDataTypeDDL", StringType, nullable = false) :: Nil)
+
+ def dataTypeToDDL(dataType: DataType): String = dataType match {
+ case _: StringType =>
+ // Hide collation information in DDL format
Review Comment:
Yes. IIUC, this test runs for all expressions, and assert that for string utf8Binary and string utf8Lcase the expression should have the same output, or throw the same exception.
If I force to output collation in toDDL/fromDDL, approx_top_k_accumulate has different outputs and fails on assertion
> ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string collate utf8_binary not null}") did not equal ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string collate utf8_lcase not null}")
If I simply use `StructField("item", dataType, nullable = false).toDDL` for string, approx_top_k_accumulate still has different outputs and fails on assertion
>ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string not null}") did not equal ArraySeq("{[04 01 0a 03 03 00 00 00 01 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 01 00 00 00 00 00 00 00 0c 00 00 00 64 75 6d 6d 79 20 73 74 72 69 6e 67], 5, null, item string collate utf8_lcase not null}")
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2417560164
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +507,203 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def setItemDataType(dataType: DataType): Unit = {
+ if (this.itemDataType == null) {
+ this.itemDataType = dataType
+ } else if (this.itemDataType != dataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(this.itemDataType, dataType)
+ }
+ }
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def setMaxItemsTracked(maxItemsTracked: Int): Unit = this.maxItemsTracked = maxItemsTracked
+}
+
+/**
+ * An aggregate function that combines multiple sketches into a single sketch.
+ *
+ * @param state the expression containing the sketches to combine
+ * @param maxItemsTracked the maximum number of items to track in the sketch
+ * @param mutableAggBufferOffset the offset for mutable aggregation buffer
+ * @param inputAggBufferOffset the offset for input aggregation buffer
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(state, maxItemsTracked) - Combines multiple sketches into a single sketch.
+ `maxItemsTracked` An optional positive INTEGER literal with upper limit of 1000000. If maxItemsTracked is not specified, it defaults to 10000.
+ """,
+ examples = """
+ Examples:
+ > SELECT approx_top_k_estimate(_FUNC_(sketch, 10000), 5) FROM (SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (0), (0), (1), (1) AS tab(expr) UNION ALL SELECT approx_top_k_accumulate(expr) AS sketch FROM VALUES (2), (3), (4), (4) AS tab(expr));
+ [{"item":0,"count":2},{"item":4,"count":2},{"item":1,"count":2},{"item":2,"count":1},{"item":3,"count":1}]
+ """,
+ group = "agg_funcs",
+ since = "4.1.0")
+// scalastyle:on line.size.limit
+case class ApproxTopKCombine(
+ state: Expression,
+ maxItemsTracked: Expression,
+ mutableAggBufferOffset: Int = 0,
+ inputAggBufferOffset: Int = 0)
+ extends TypedImperativeAggregate[CombineInternal[Any]]
+ with ImplicitCastInputTypes
+ with BinaryLike[Expression] {
+
+ def this(child: Expression, maxItemsTracked: Expression) = {
+ this(child, maxItemsTracked, 0, 0)
+ ApproxTopK.checkExpressionNotNull(maxItemsTracked, "maxItemsTracked")
+ ApproxTopK.checkMaxItemsTracked(maxItemsTrackedVal)
+ }
+
+ def this(child: Expression, maxItemsTracked: Int) = this(child, Literal(maxItemsTracked))
+
+ def this(child: Expression) = this(child, Literal(ApproxTopK.VOID_MAX_ITEMS_TRACKED), 0, 0)
+
+ private lazy val uncheckedItemDataType: DataType =
+ state.dataType.asInstanceOf[StructType](1).dataType
+ private lazy val maxItemsTrackedVal: Int = maxItemsTracked.eval().asInstanceOf[Int]
+ private lazy val combineSizeSpecified: Boolean =
+ maxItemsTrackedVal != ApproxTopK.VOID_MAX_ITEMS_TRACKED
+
+ override def left: Expression = state
+
+ override def right: Expression = maxItemsTracked
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegerType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val defaultCheck = super.checkInputDataTypes()
+ if (defaultCheck.isFailure) {
+ defaultCheck
+ } else {
+ val stateCheck = ApproxTopK.checkStateFieldAndType(state)
+ if (stateCheck.isFailure) {
+ stateCheck
+ } else if (!maxItemsTracked.foldable) {
+ TypeCheckFailure("Number of items tracked must be a constant literal")
+ } else {
+ TypeCheckSuccess
+ }
+ }
+ }
+
+ override def dataType: DataType = ApproxTopK.getSketchStateDataType(uncheckedItemDataType)
+
+ override def createAggregationBuffer(): CombineInternal[Any] = {
+ if (combineSizeSpecified) {
+ val maxMapSize = ApproxTopK.calMaxMapSize(maxItemsTrackedVal)
+ new CombineInternal[Any](
+ new ItemsSketch[Any](maxMapSize),
+ null,
+ maxItemsTrackedVal)
+ } else {
+ new CombineInternal[Any](
+ new ItemsSketch[Any](ApproxTopK.SKETCH_SIZE_PLACEHOLDER),
+ null,
+ ApproxTopK.VOID_MAX_ITEMS_TRACKED)
Review Comment:
Shall we simply set it as `MAX_ITEMS_TRACKED_LIMIT: Int = 1000000` here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "cboumalh (via GitHub)" <gi...@apache.org>.
cboumalh commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430661588
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -435,3 +491,283 @@ case class ApproxTopKAccumulate(
override def prettyName: String =
getTagValue(FunctionRegistry.FUNC_ALIAS).getOrElse("approx_top_k_accumulate")
}
+
+/**
+ * In internal class used as the aggregation buffer for ApproxTopKCombine.
+ *
+ * @param sketch the ItemsSketch instance
+ * @param itemDataType the data type of items in the sketch
+ * @param maxItemsTracked the maximum number of items tracked in the sketch
+ */
+class CombineInternal[T](
+ sketch: ItemsSketch[T],
+ var itemDataType: DataType,
+ var maxItemsTracked: Int) {
+ def getSketch: ItemsSketch[T] = sketch
+
+ def getItemDataType: DataType = itemDataType
+
+ def getMaxItemsTracked: Int = maxItemsTracked
+
+ def updateMaxItemsTracked(combineSizeSpecified: Boolean, newMaxItemsTracked: Int): Unit = {
+ if (!combineSizeSpecified) {
+ // check size
+ if (this.maxItemsTracked == ApproxTopK.VOID_MAX_ITEMS_TRACKED) {
+ // If buffer's maxItemsTracked VOID_MAX_ITEMS_TRACKED, it means the buffer is a placeholder
+ // sketch that has not beed updated by any input sketch yet.
+ // So we can set it to the input sketch's max items tracked.
+ this.maxItemsTracked = newMaxItemsTracked
+ } else {
+ if (this.maxItemsTracked != newMaxItemsTracked) {
+ // If buffer's maxItemsTracked is not VOID_MAX_ITEMS_TRACKED, it means the buffer has been
+ // updated by some input sketch. So if buffer and input sketch have different
+ // maxItemsTracked values, it means at least two of the input sketches have different
+ // maxItemsTracked values. In this case, we should throw an error.
+ throw QueryExecutionErrors.approxTopKSketchSizeNotMatch(
+ this.maxItemsTracked, newMaxItemsTracked)
+ }
+ }
+ }
+ }
+
+ def updateItemDataType(inputItemDataType: DataType): Unit = {
+ // When the buffer's dataType hasn't been set, set it to the input sketch's item data type
+ // When input sketch's item data type is null, buffer's item data type will remain null
+ if (this.itemDataType == null) {
+ this.itemDataType = inputItemDataType
+ } else {
+ // When the buffer's dataType has been set, throw an error
+ // if the input sketch's item data type is not null the two data types don't match
+ if (inputItemDataType != null && this.itemDataType != inputItemDataType) {
+ throw QueryExecutionErrors.approxTopKSketchTypeNotMatch(
+ this.itemDataType, inputItemDataType)
+ }
+ }
+ }
+
+ /**
+ * Serialize the CombineInternal instance to a byte array.
+ * Serialization format:
+ * maxItemsTracked (4 bytes int) +
+ * itemDataTypeDDL length n in byte (4 bytes int) +
+ * itemDataTypeDDL (n bytes) +
+ * sketchBytes
+ */
+ def serialize(): Array[Byte] = {
+ val sketchBytes = sketch.toByteArray(
+ ApproxTopK.genSketchSerDe(itemDataType).asInstanceOf[ArrayOfItemsSerDe[T]])
+ val itemDataTypeDDL = ApproxTopK.dataTypeToDDL(itemDataType)
+ val ddlBytes: Array[Byte] = itemDataTypeDDL.getBytes(StandardCharsets.UTF_8)
+ val byteArray = new Array[Byte](sketchBytes.length + 4 + 4 + ddlBytes.length)
Review Comment:
super nit, we could consider `Integer.BYTES` instead of 4 here for clarity about what the extra bytes represent. Not necessary
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "yhuang-db (via GitHub)" <gi...@apache.org>.
yhuang-db commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2427275479
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -318,7 +325,72 @@ object ApproxTopK {
StructType(
StructField("sketch", BinaryType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
+ StructField("typeCode", BinaryType, nullable = false) :: Nil)
Review Comment:
Need an explicit encoding of the dataType, because the `null` of `StructField("itemDataType", itemDataType)` might be implicitly widen during union.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-52798] [SQL] Add function approx_top_k_combine [spark]
Posted by "gengliangwang (via GitHub)" <gi...@apache.org>.
gengliangwang commented on code in PR #51505:
URL: https://github.com/apache/spark/pull/51505#discussion_r2430767557
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ApproxTopKAggregates.scala:
##########
@@ -317,8 +322,51 @@ object ApproxTopK {
def getSketchStateDataType(itemDataType: DataType): StructType =
StructType(
StructField("sketch", BinaryType, nullable = false) ::
+ StructField("maxItemsTracked", IntegerType, nullable = false) ::
StructField("itemDataType", itemDataType) ::
- StructField("maxItemsTracked", IntegerType, nullable = false) :: Nil)
+ StructField("itemDataTypeDDL", StringType, nullable = false) :: Nil)
+
+ def dataTypeToDDL(dataType: DataType): String = dataType match {
+ case _: StringType =>
+ // Hide collation information in DDL format
Review Comment:
I am confused. Why the test failed? Does it run with approx_top_k_*?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org