[SPARK-58069][SQL] approx_top_k returns incorrect/garbage sort-key bytes instead of actual values for collated strings

[jiwen624]

What changes were proposed in this pull request?

For a non-UTF8_BINARY collated STRING column, approx_top_k (and approx_top_k_accumulate / _combine / _estimate) stored the collation sort key in the sketch and returned it verbatim as the item. This makes the sketch store CollatedString instead: it dedupes by the collation key (so collation-equal values still count as one) but retains an actual input value to return, the way mode() does. The serde writes only the original value and recomputes the key on read, so the UTF8_BINARY format is unchanged.

Note: while working on the fix, two new issues are found in the approx_top_k_ expressions. The following Jira issues are raised to address them separately - so this PR will only focus on the issue mentioned in SPARK-57177:

• https://issues.apache.org/jira/browse/SPARK-58096
• https://issues.apache.org/jira/browse/SPARK-58095

Why are the changes needed?

The returned items were wrong: raw ICU sort-key bytes (invalid UTF-8) for ICU collations, and the lower-cased form for UTF8_LCASE (a value that may never have appeared). Counts were correct, but the item labels were garbage. Silent wrong result, no error.

Before the fix:

scala> spark.sql("""SELECT approx_top_k(c, 5, 100) FROM (
     |   SELECT CAST(col AS STRING COLLATE UNICODE_CI) AS c
     |   FROM VALUES ('HELLO'), ('HELLO'), ('HELLO'), ('hello'), ('world') AS t(col)
     | )""").show(100, false)

+--------------------------------+
|approx_top_k(c, 5, 100)         |
+--------------------------------+
|[{93AAG\t, 4}, {WGMA1\t, 1}]|
+--------------------------------+

scala> spark.sql("""SELECT approx_top_k(c, 5, 100) FROM (
     |   SELECT CAST(col AS STRING COLLATE UNICODE) AS c
     |   FROM VALUES ('HELLO'), ('HELLO'), ('HELLO'), ('hello'), ('world') AS t(col)
     | )""").show(100, false)
+------------------------------------------------------------+
|approx_top_k(c, 5, 100)                                     |
+------------------------------------------------------------+
|[{93AAG\t�����, 3}, {93AAG\t\t, 1}, {WGMA1\t\t, 1}]|
+------------------------------------------------------------+

scala> spark.sql("""SELECT approx_top_k(c, 5, 100) FROM (
     |   SELECT CAST(col AS STRING COLLATE UTF8_LCASE) AS c
     |   FROM VALUES ('HELLO'), ('HELLO'), ('HELLO'), ('hello'), ('world') AS t(col)
     | )""").show(100, false)
+------------------------+
|approx_top_k(c, 5, 100) |
+------------------------+
|[{hello, 4}, {world, 1}]|
+------------------------+

Does this PR introduce any user-facing change?

Yes. approx_top_k over a collated string column now returns an actual input value instead of the internal collation key. Sketches persisted by approx_top_k_accumulate over a non-binary collated column before this change are not readable after it (they contained wrong data); the UTF8_BINARY format is unchanged.

How was this patch tested?

Added test cases.

Was this patch authored or co-authored using generative AI tooling?

Yes.

[jiwen624]

Hi @cloud-fan @gengliangwang @yhuang-db since you've worked on and/or reviewed code of this area, could you take a look at this bug fix when you get a chance? Thanks.
Note that as mentioned in the PR description, there are two follow-up Jiras/fixes that need to be addressed after this PR.

[cloud-fan]

0 blocking, 1 non-blocking, 0 nits.
Clean, well-structured fix that separates the dedup key from the returned value (the mode() pattern), with sound collation propagation across the shuffle boundary and good test coverage.

Correctness (1)

• ApproxTopKAggregates.scala:386: collated dedup key uses a lossy .toString of ICU sort-key bytes, which can over-merge collation-distinct values; the cited mode() analogue keys on the byte-exact UTF8String — see inline

Verification

Traced the collation dedup against the cited mode() analogue. mode() keys its buffer on the raw UTF8String collation key (Mode.scala:118), which is byte-exact. This PR keys CollatedString on getCollationKey(...).toString (ApproxTopKAggregates.scala:386), and UTF8String.toString decodes the bytes as UTF-8, mapping every invalid byte sequence to U+FFFD (UTF8String.java:2018). ICU sort keys are not valid UTF-8, so two collation-distinct keys can decode to the same String and over-merge. The counts, serde round-trip, and combine collation-propagation paths verified otherwise consistent; this is the one gap.

[cloud-fan]

The dedup key here is a lossy UTF-8 decode of the ICU sort-key bytes: getCollationKey(...) returns a UTF8String of the raw sort key, and .toString runs new String(bytes, UTF_8) (UTF8String.java:2018), which replaces every invalid-UTF-8 byte sequence with U+FFFD. ICU sort keys are arbitrary bytes, not valid UTF-8, so two collation-distinct values whose keys differ only within invalid-byte regions can decode to the same String, compare equal in CollatedString.equals/hashCode, and be over-merged into one item — inflating the survivor's count and dropping a distinct value from the top-K.

The mode() pattern this PR follows avoids that by keying on the UTF8String collation key directly (Mode.scala:118), which is byte-exact. Consider keying CollatedString on the UTF8String (or the raw key bytes) rather than the decoded String. Non-blocking — the lossy decode predates this PR — but since this path is being reworked and CollatedString is new, it's a clean spot to close the gap.

[jiwen624]

Thanks for the review @cloud-fan ! As mentioned in the PR description this issue is found while working on this PR and a follow-up ticket is raised at https://issues.apache.org/jira/browse/SPARK-58096 . The fix is not introduced to this PR to avoid confusion since the lossy decode is a different issue that predates this PR and the PR itself is already big.
Does it make sense to you?

[cloud-fan]

0 addressed, 1 remaining, 0 new.

Re-reviewed at the same commit as my prior review (no new commits). A full fresh pass surfaced no new findings.

The one remaining item (over-merge from the lossy .toString dedup key at ApproxTopKAggregates.scala:386) is pre-existing in master and you've scoped the fix to follow-up SPARK-58096 — agreed, deferring it is reasonable; it's orthogonal to this PR's key-vs-value separation, which is correct regardless.

Verification

Traced the fix against the cited mode() analogue and the persistence paths. CollatedString correctly separates the dedup key (equals/hashCode on the collation key) from the returned value (original), matching mode()'s key-vs-value pattern; eval returns original for both CollatedString and plain String, with a clear internalError on the impossible third case. The new ArrayOfCollatedStringsSerDe writes only original via the delegate ArrayOfStringsSerDe and recomputes the key on read, so the UTF8_BINARY on-wire format is unchanged — the same structure as ArrayOfDecimalsSerDe. Across the shuffle boundary, CombineInternal now persists the item type as collation-preserving JSON and recovers collation from the state struct's static field-2 type via withCollationOf, so a collated sketch merges by collation key consistently on the accumulate/combine/estimate paths. LGTM.

[cloud-fan]

Merge Summary:

• merged into master 384d5e8
• merged into branch-4.x addcf3e

Posted by merge_spark_pr.py