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

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

@@ -5560,6 +5560,18 @@
     ],
     "sqlState" : "428EK"
   },
+  "THETA_INVALID_INPUT_SKETCH_BUFFER" : {
+    "message" : [
+      "Invalid call to <function>; only valid Theta sketch buffers are supported as inputs (such as those produced by the `theta_sketch_agg` function)."
+    ],
+    "sqlState" : "22546"
+  },
+  "THETA_INVALID_LG_NOM_ENTRIES" : {
+    "message" : [
+      "Invalid Theta sketch call; the `lgNomEntries` value must be between <min> and <max>, inclusive: <value>."
+    ],
+    "sqlState" : "22546"
+  },
   "TRAILING_COMMA_IN_SELECT" : {
     "message" : [
       "Trailing comma detected in SELECT clause. Remove the trailing comma before the FROM clause."

python/docs/source/reference/pyspark.sql/functions.rst

@@ -490,6 +490,9 @@ Aggregate Functions
     string_agg_distinct
     sum
     sum_distinct
+    theta_intersection_agg
+    theta_sketch_agg
+    theta_union_agg
     try_avg
     try_sum
     var_pop
@@ -635,6 +638,10 @@ Misc Functions
     reflect
     session_user
     spark_partition_id
+    theta_difference
+    theta_intersection
+    theta_sketch_estimate
+    theta_union
     try_aes_decrypt
     try_reflect
     typeof

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

@@ -4229,6 +4229,90 @@ def hll_union(
 hll_union.__doc__ = pysparkfuncs.hll_union.__doc__
 
 
+def theta_sketch_agg(
+    col: "ColumnOrName",
+    lgNomEntries: Optional[Union[int, Column]] = None,
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function_over_columns("theta_sketch_agg", col)
+    else:
+        return _invoke_function_over_columns("theta_sketch_agg", col, lit(lgNomEntries))
+
+
+theta_sketch_agg.__doc__ = pysparkfuncs.theta_sketch_agg.__doc__
+
+
+def theta_union_agg(
+    col: "ColumnOrName",
+    lgNomEntries: Optional[Union[int, Column]] = None,
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function_over_columns("theta_union_agg", col)
+    else:
+        return _invoke_function_over_columns("theta_union_agg", col, lit(lgNomEntries))
+
+
+theta_union_agg.__doc__ = pysparkfuncs.theta_union_agg.__doc__
+
+
+def theta_intersection_agg(
+    col: "ColumnOrName",
+    lgNomEntries: Optional[Union[int, Column]] = None,
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function_over_columns("theta_intersection_agg", col)
+    else:
+        return _invoke_function_over_columns("theta_intersection_agg", col, lit(lgNomEntries))
+
+
+theta_intersection_agg.__doc__ = pysparkfuncs.theta_intersection_agg.__doc__
+
+
+def theta_sketch_estimate(col: "ColumnOrName") -> Column:
+    return _invoke_function("theta_sketch_estimate", _to_col(col))
+
+
+theta_sketch_estimate.__doc__ = pysparkfuncs.theta_sketch_estimate.__doc__
+
+
+def theta_union(
+    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function("theta_union", _to_col(col1), _to_col(col2))
+    else:
+        return _invoke_function("theta_union", _to_col(col1), _to_col(col2), lit(lgNomEntries))
+
+
+theta_union.__doc__ = pysparkfuncs.theta_union.__doc__
+
+
+def theta_intersection(
+    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function("theta_intersection", _to_col(col1), _to_col(col2))
+    else:
+        return _invoke_function(
+            "theta_intersection", _to_col(col1), _to_col(col2), lit(lgNomEntries)
+        )
+
+
+theta_intersection.__doc__ = pysparkfuncs.theta_intersection.__doc__
+
+
+def theta_difference(
+    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
+) -> Column:
+    if lgNomEntries is None:
+        return _invoke_function("theta_difference", _to_col(col1), _to_col(col2))
+    else:
+        return _invoke_function("theta_difference", _to_col(col1), _to_col(col2), lit(lgNomEntries))
+
+
+theta_difference.__doc__ = pysparkfuncs.theta_difference.__doc__
+
+
 # Predicates Function
 
 

python/pyspark/sql/functions/__init__.py

@@ -403,6 +403,9 @@
     "string_agg_distinct",
     "sum",
     "sum_distinct",
+    "theta_intersection_agg",
+    "theta_sketch_agg",
+    "theta_union_agg",
     "try_avg",
     "try_sum",
     "var_pop",
@@ -494,6 +497,10 @@
     "reflect",
     "session_user",
     "spark_partition_id",
+    "theta_difference",
+    "theta_intersection",
+    "theta_sketch_estimate",
+    "theta_union",
     "try_aes_decrypt",
     "try_reflect",
     "typeof",