[SPARK-52407][SQL] Add support for Theta Sketch #51298

cboumalh · 2025-06-27T01:06:27Z

What changes were proposed in this pull request?

This PR's goal is to provide 7 new functions which utilize Theta Sketches.
The functions are:

theta_sketch_agg(sketch_col)
theta_union(sketch1, sketch2)
theta_union_agg(sketch_col)
theta_difference(sketch1, sketch2)
theta_intersection(sketch1, sketch2)
theta_intersection_agg(sketch_col)
theta_sketch_estimate(sketch)

Why are the changes needed?

Today, Spark supports HyperLogLog. HLL only supports the union set operation. Theta on the other hand allows us to find the difference and intersection between two sketches.

Does this PR introduce any user-facing change?

Yes, this PR introduces three new aggregate functions:

theta_sketch_agg
theta_intersection_agg
theta_union_agg

And four new scalar functions:

theta_sketch_estimate
theta_intersection
theta_difference
theta_union

How was this patch tested?

I've included some tests in the test suites. I also tested it locally on my machine.

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

No

Jira

https://issues.apache.org/jira/browse/SPARK-52407

cboumalh · 2025-06-27T01:13:55Z

I know there is probably a lot missing in the PR before it's ready to merge. Please let me know!
@mkaravel @RyanBerti

…ketch

…n examples in ExpressionDescription and unit tests

…nditions

…and python builtin.py

cboumalh · 2025-07-21T05:56:59Z

Below are some of the results I came up with. This run was using spark shell on a c6a.2xlarge instance.

Sketch Type	Data Type	Size	Unique Count	Estimate	Error %	Time (ms)
Theta	Integer	12	20000000	19706898	1.47	315
Theta	String	12	20000000	20158220	0.79	679
HLL	Integer	12	20000000	20084065	0.42	297
HLL	String	12	20000000	20413618	2.07	628
Theta	Integer	20	20000000	20001631	0.01	1313
Theta	String	20	20000000	20005869	0.03	1884
HLL	Integer	20	20000000	20029729	0.15	366
HLL	String	20	20000000	20012835	0.06	694
Theta	Integer	12	10000000	10101680	1.02	154
Theta	String	12	10000000	9984868	0.15	361
HLL	Integer	12	10000000	10146710	1.47	131
HLL	String	12	10000000	10129938	1.30	296
Theta	Integer	20	10000000	9995024	0.05	924
Theta	String	20	10000000	9994589	0.05	1161
HLL	Integer	20	10000000	9995662	0.04	208
HLL	String	20	10000000	10004646	0.05	379
Theta	Integer	12	5000000	5068868	1.38	103
Theta	String	12	5000000	5026543	0.53	202
HLL	Integer	12	5000000	5057766	1.16	128
HLL	String	12	5000000	4906794	1.86	207
Theta	Integer	20	5000000	4999593	0.01	592
Theta	String	20	5000000	5001495	0.03	693
HLL	Integer	20	5000000	4996974	0.06	175
HLL	String	20	5000000	5004207	0.08	284
Theta	Integer	12	1000000	1003909	0.39	75
Theta	String	12	1000000	990056	0.99	91
HLL	Integer	12	1000000	976680	2.33	56
HLL	String	12	1000000	1002235	0.22	92
Theta	Integer	20	1000000	1000000	0.00	227
Theta	String	20	1000000	1000000	0.00	227
HLL	Integer	20	1000000	999921	0.01	113
HLL	String	20	1000000	1000244	0.02	138
Theta	Integer	12	500000	493167	1.37	78
Theta	String	12	500000	491379	1.72	112
HLL	Integer	12	500000	502392	0.48	107
HLL	String	12	500000	488070	2.39	85
Theta	Integer	20	500000	500000	0.00	163
Theta	String	20	500000	500000	0.00	233
HLL	Integer	20	500000	500639	0.13	91
HLL	String	20	500000	500114	0.02	128
Theta	Integer	12	100000	98046	1.95	195
Theta	String	12	100000	98214	1.79	172
HLL	Integer	12	100000	102878	2.88	115
HLL	String	12	100000	100078	0.08	97
Theta	Integer	20	100000	100000	0.00	120
Theta	String	20	100000	100000	0.00	142
HLL	Integer	20	100000	100013	0.01	102
HLL	String	20	100000	100013	0.01	93

Results comparison by group

100000 entries (size 20):

Integer: Theta slower by 1.18x, Error diff: 0.01%
String: Theta slower by 1.53x, Error diff: 0.01%

100000 entries (size 12):

Integer: Theta slower by 1.70x, Error diff: 0.92%
String: Theta slower by 1.77x, Error diff: 1.71%

500000 entries (size 20):

Integer: Theta slower by 1.79x, Error diff: 0.13%
String: Theta slower by 1.82x, Error diff: 0.02%

500000 entries (size 12):

Integer: Theta faster by 1.37x, Error diff: 0.89%
String: Theta slower by 1.32x, Error diff: 0.66%

1000000 entries (size 20):

Integer: Theta slower by 2.01x, Error diff: 0.01%
String: Theta slower by 1.64x, Error diff: 0.02%

1000000 entries (size 12):

Integer: Theta slower by 1.34x, Error diff: 1.94%
String: Theta faster by 1.01x, Error diff: 0.77%

5000000 entries (size 20):

Integer: Theta slower by 3.38x, Error diff: 0.05%
String: Theta slower by 2.44x, Error diff: 0.05%

5000000 entries (size 12):

Integer: Theta faster by 1.24x, Error diff: 0.22%
String: Theta faster by 1.02x, Error diff: 1.33%

10000000 entries (size 20):

Integer: Theta slower by 4.44x, Error diff: 0.01%
String: Theta slower by 3.06x, Error diff: 0.01%

10000000 entries (size 12):

Integer: Theta slower by 1.18x, Error diff: 0.45%
String: Theta slower by 1.22x, Error diff: 1.15%

20000000 entries (size 20):

Integer: Theta slower by 3.59x, Error diff: 0.14%
String: Theta slower by 2.71x, Error diff: 0.03%

20000000 entries (size 12):

Integer: Theta slower by 1.06x, Error diff: 1.05%
String: Theta slower by 1.08x, Error diff: 1.28%

Summary and comments

Theta is competitive. It wins or is very close in many cases. Where HLL dominates is with larger sketch sizes, consistently 3-5x faster than Theta. Theta's major weakness is that it becomes significantly slower at higher precision, but on the plus side it is very accurate. Choose Theta sketches when: You need set operations (intersection, difference, union), size 12 is acceptable for your accuracy needs, and you prioritize functionality over raw speed.

https://gist.github.com/cboumalh/81ae7e1362a35d2a9968184fa1a52e0d

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

…o utils

xiongbo-sjtu · 2025-07-24T03:33:43Z

LGTM, but I don't have permission to approve it.

@dongjoon-hyun Can you please help find someone to review this PR? Thanks!

…afety

dtenedor · 2025-08-26T21:16:44Z

I can help review this, I reviewed the original HLL sketch aggregate function contributions by Ryan Berti earlier.

...lyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala

sql/core/src/test/resources/sql-tests/inputs/thetasketch.sql

...main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ThetaSketchUtils.scala

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala

dtenedor

This is looking good, thanks again for putting in all the work to make this implementation available for Spark!!

common/utils/src/main/resources/error/error-conditions.json

python/pyspark/sql/connect/functions/builtin.py

python/pyspark/sql/functions/builtin.py

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/ThetaSketchUtilsSuite.scala

...lyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/thetasketchesExpressions.scala

...main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala

dtenedor

This is looking almost done, just a few comments left and we can merge it soon.

sql/core/src/test/resources/sql-tests/inputs/thetasketch.sql

dtenedor

I reviewed everything again and the functionality looks clear and thoroughly-tested. We can leave the PR open for another day or so if anyone else has any concerns, then proceed to merge it.

cboumalh · 2025-09-04T23:58:18Z

Sounds good! Thanks a lot for the thorough review. I’ll be available in case others want to chime in, and I’m happy to address any last-minute feedback. Really appreciate your help moving this forward and excited to make this feature available for the Spark community.

python/pyspark/sql/functions/builtin.py

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

...main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala

...c/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/ThetasketchesAggSuite.scala

sql/core/src/test/resources/sql-tests/results/thetasketch.sql.out

mkaravel

@cboumalh Thank you for addressing the review comments!

One more quick round of review. Please take a look.
I intend to do another pass.

python/pyspark/sql/functions/builtin.py

.../main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/datasketchesAggregates.scala

...main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala

github-actions bot added the SQL label Jun 27, 2025

[SPARK-52407][SQL] Add support for Theta Sketch

e9a2cbe

cboumalh force-pushed the SPARK-52407_add_datasketches_thetasketch branch from b301ac8 to e9a2cbe Compare June 27, 2025 01:12

cboumalh marked this pull request as draft June 27, 2025 01:16

cboumalh and others added 6 commits June 26, 2025 21:22

Merge branch 'apache:master' into SPARK-52407_add_datasketches_thetas…

5185c7f

…ketch

added ExpressionDescription to ThetaIntersectionAgg, fixed some broke…

a44a6be

…n examples in ExpressionDescription and unit tests

wrap array unit tests with array data to avoid casting errors test run

7c8e4f5

Update merge function parameter name in thesketchesAggregates

07a0c87

Added Dataframes Suite, Theta Sketch Utils, sql expressions, error co…

1b3bfc8

…nditions

generate golden files

edae21a

github-actions bot added the CONNECT label Jul 15, 2025

python functionality added

db75ca3

github-actions bot added DOCS PYTHON labels Jul 15, 2025

cboumalh and others added 6 commits July 15, 2025 10:21

Merge branch 'master' into SPARK-52407_add_datasketches_thetasketch

a892e0c

added missing builtin.py code for pyspark

9e9886e

fix typos and formatting errors

b0d57a9

fix broken tests in ExpressionsSchemaSuite, PlanGenerationTestSuite, …

5cca851

…and python builtin.py

add .explain files from ProtoToParsedPlanTestSuite

a00fdd6

typos, formatting, and naming fixes

b3cb40a

cboumalh marked this pull request as ready for review July 16, 2025 17:22

xiongbo-sjtu reviewed Jul 21, 2025

View reviewed changes

sql/api/src/main/scala/org/apache/spark/sql/functions.scala Outdated Show resolved Hide resolved

sql/api/src/main/scala/org/apache/spark/sql/functions.scala Outdated Show resolved Hide resolved

changed theta sketch spark since version, added default sketch size t…

2334aa1

…o utils

Chris Boumalhab added 4 commits July 24, 2025 17:13

remove unnecessary query from sql query tests

134dd29

update to use toByteArrayCompressed, heapify in deserialization for s…

4928059

…afety

format

32a4d66

optimize aggregating performance

30fd361

Chris Boumalhab added 2 commits August 10, 2025 13:41

fix

34ad644

comments fix

a80b7d4

comment fix

e91d702

dtenedor reviewed Aug 26, 2025

View reviewed changes

...main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/thetasketchesAggregates.scala Outdated Show resolved Hide resolved

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/ThetaSketchUtils.scala Outdated Show resolved Hide resolved

dtenedor reviewed Aug 26, 2025

View reviewed changes

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala Outdated Show resolved Hide resolved

resolving comments

e792ac2

dtenedor reviewed Aug 28, 2025

View reviewed changes

Chris Boumalhab added 5 commits August 29, 2025 03:41

resolving comments

8905ea9

comment fix

1075f6b

fix

ffd2b30

fix

49c88fc

comments fix

ae1f1bd

dtenedor reviewed Aug 29, 2025

View reviewed changes

sql/core/src/test/resources/sql-tests/inputs/thetasketch.sql Outdated Show resolved Hide resolved

Chris Boumalhab added 5 commits August 30, 2025 00:18

resolve comments

8ef59e8

fix

0ae8315

fix

526d791

fix

cc452e8

fix

c86c4bf

dtenedor approved these changes Sep 4, 2025

View reviewed changes

mkaravel reviewed Sep 8, 2025

View reviewed changes

Chris Boumalhab added 4 commits September 9, 2025 00:02

fix

f1669ec

fix

93f4a0a

fix

44fd9d4

fix

613ee05

mkaravel reviewed Sep 11, 2025

View reviewed changes

fix

6c3b9bc

cboumalh mentioned this pull request Sep 11, 2025

[SPARK-53559][SQL][CATALYST] Fix HLL sketch updates to use raw collation key bytes #52316

Open

[SPARK-52407][SQL] Add support for Theta Sketch #51298

Are you sure you want to change the base?

[SPARK-52407][SQL] Add support for Theta Sketch #51298

Conversation

cboumalh commented Jun 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

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

Jira

Uh oh!

cboumalh commented Jun 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cboumalh commented Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Results comparison by group

Summary and comments

Uh oh!

Uh oh!

Uh oh!

xiongbo-sjtu commented Jul 24, 2025

Uh oh!

dtenedor commented Aug 26, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dtenedor left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dtenedor left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dtenedor left a comment

Choose a reason for hiding this comment

Uh oh!

cboumalh commented Sep 4, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mkaravel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cboumalh commented Jun 27, 2025 •

edited

Loading

cboumalh commented Jun 27, 2025 •

edited

Loading

cboumalh commented Jul 21, 2025 •

edited

Loading