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

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

@@ -4259,60 +4259,46 @@ def theta_union_agg(
 
 def theta_intersection_agg(
     col: "ColumnOrName",
-    lgNomEntries: Optional[Union[int, Column]] = None,
 ) -> Column:
     fn = "theta_intersection_agg"
-    if lgNomEntries is None:
-        return _invoke_function_over_columns(fn, col)
-    else:
-        return _invoke_function_over_columns(fn, col, lit(lgNomEntries))
+    return _invoke_function_over_columns(fn, col)
 
 
 theta_intersection_agg.__doc__ = pysparkfuncs.theta_intersection_agg.__doc__
 
 
 def theta_sketch_estimate(col: "ColumnOrName") -> Column:
     fn = "theta_sketch_estimate"
-    return _invoke_function(fn, _to_col(col))
+    return _invoke_function_over_columns(fn, col)
 
 
 theta_sketch_estimate.__doc__ = pysparkfuncs.theta_sketch_estimate.__doc__
 
 
 def theta_union(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
+    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[Union[int, Column]] = None
 ) -> Column:
     fn = "theta_union"
     if lgNomEntries is None:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2))
+        return _invoke_function_over_columns(fn, col1, col2)
     else:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2), lit(lgNomEntries))
+        return _invoke_function_over_columns(fn, col1, col2, lit(lgNomEntries))
 
 
 theta_union.__doc__ = pysparkfuncs.theta_union.__doc__
 
 
-def theta_intersection(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
-) -> Column:
+def theta_intersection(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
     fn = "theta_intersection"
-    if lgNomEntries is None:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2))
-    else:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2), lit(lgNomEntries))
+    return _invoke_function_over_columns(fn, col1, col2)
 
 
 theta_intersection.__doc__ = pysparkfuncs.theta_intersection.__doc__
 
 
-def theta_difference(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
-) -> Column:
+def theta_difference(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
     fn = "theta_difference"
-    if lgNomEntries is None:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2))
-    else:
-        return _invoke_function(fn, _to_col(col1), _to_col(col2), lit(lgNomEntries))
+    return _invoke_function_over_columns(fn, col1, col2)
 
 
 theta_difference.__doc__ = pysparkfuncs.theta_difference.__doc__

python/pyspark/sql/functions/builtin.py

@@ -15799,8 +15799,7 @@ def regexp_count(str: "ColumnOrName", regexp: "ColumnOrName") -> Column:
 
 @_try_remote_functions
 def regexp_extract(str: "ColumnOrName", pattern: str, idx: int) -> Column:
-    r"""Extract a specific group matched by the Java regex `regexp`, from the specified
-    string column.
+    r"""Extract a specific group matched by the Java regex `regexp`, from the specified string column.
     If the regex did not match, or the specified group did not match, an empty string is returned.
 
     .. versionadded:: 1.5.0
@@ -25712,7 +25711,7 @@ def theta_sketch_agg(
 ) -> Column:
     """
     Aggregate function: returns the compact binary representation of the Datasketches
-    ThetaSketch configured with lgNomEntries arg.
+    ThetaSketch with the values in the input column configured with lgNomEntries nominal entries.
 
     .. versionadded:: 4.1.0
 
@@ -25769,8 +25768,7 @@ def theta_union_agg(
 ) -> Column:
     """
     Aggregate function: returns the compact binary representation of the Datasketches
-    ThetaSketch, generated by merging previously created Datasketches ThetaSketch instances
-    via a Datasketches Union instance.
+    ThetaSketch that is the union of the Theta sketches in the input column.
 
     .. versionadded:: 4.1.0
 
@@ -25815,23 +25813,16 @@ def theta_union_agg(
 
 
 @_try_remote_functions
-def theta_intersection_agg(
-    col: "ColumnOrName",
-    lgNomEntries: Optional[Union[int, Column]] = None,
-) -> Column:
+def theta_intersection_agg(col: "ColumnOrName") -> Column:
     """
     Aggregate function: returns the compact binary representation of the Datasketches
-    ThetaSketch, generated by intersecting previously created Datasketches ThetaSketch
-    instances via a Datasketches Intersection instance.
+    ThetaSketch that is the intersection of the Theta sketches in the input column
 
     .. versionadded:: 4.1.0
 
     Parameters
     ----------
     col : :class:`~pyspark.sql.Column` or column name
-    lgNomEntries : :class:`~pyspark.sql.Column` or int, optional
-        The log-base-2 of nominal entries for the intersection operation
-        (must be between 4 and 26, defaults to 12)
 
     Returns
     -------
@@ -25853,17 +25844,14 @@ def theta_intersection_agg(
     >>> df2 = df2.agg(sf.theta_sketch_agg("value").alias("sketch"))
     >>> df3 = df1.union(df2)
     >>> df3.agg(sf.theta_sketch_estimate(sf.theta_intersection_agg("sketch"))).show()
-    +---------------------------------------------------------+
-    |theta_sketch_estimate(theta_intersection_agg(sketch, 12))|
-    +---------------------------------------------------------+
-    |                                                        2|
-    +---------------------------------------------------------+
+    +-----------------------------------------------------+
+    |theta_sketch_estimate(theta_intersection_agg(sketch))|
+    +-----------------------------------------------------+
+    |                                                    2|
+    +-----------------------------------------------------+
     """
     fn = "theta_intersection_agg"
-    if lgNomEntries is None:
-        return _invoke_function_over_columns(fn, col)
-    else:
-        return _invoke_function_over_columns(fn, col, lit(lgNomEntries))
+    return _invoke_function_over_columns(fn, col)
 
 
 @_try_remote_functions
@@ -25903,15 +25891,14 @@ def theta_sketch_estimate(col: "ColumnOrName") -> Column:
     |                                                 3|
     +--------------------------------------------------+
     """
-    from pyspark.sql.classic.column import _to_java_column
 
     fn = "theta_sketch_estimate"
-    return _invoke_function(fn, _to_java_column(col))
+    return _invoke_function_over_columns(fn, col)
 
 
 @_try_remote_functions
 def theta_union(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
+    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[Union[int, Column]] = None
 ) -> Column:
     """
     Merges two binary representations of Datasketches ThetaSketch objects, using a
@@ -25923,7 +25910,7 @@ def theta_union(
     ----------
     col1 : :class:`~pyspark.sql.Column` or column name
     col2 : :class:`~pyspark.sql.Column` or column name
-    lgNomEntries : int, optional
+    lgNomEntries : :class:`~pyspark.sql.Column` or int, optional
         The log-base-2 of nominal entries for the union operation
         (must be between 4 and 26, defaults to 12)
 
@@ -25953,24 +25940,21 @@ def theta_union(
     |                                                       6|
     +--------------------------------------------------------+
     """
-    from pyspark.sql.classic.column import _to_java_column
 
     fn = "theta_union"
     if lgNomEntries is not None:
-        return _invoke_function(
+        return _invoke_function_over_columns(
             fn,
-            _to_java_column(col1),
-            _to_java_column(col2),
-            _enum_to_value(lgNomEntries),
+            col1,
+            col2,
+            lit(lgNomEntries),
         )
     else:
-        return _invoke_function(fn, _to_java_column(col1), _to_java_column(col2))
+        return _invoke_function_over_columns(fn, col1, col2)
 
 
 @_try_remote_functions
-def theta_intersection(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
-) -> Column:
+def theta_intersection(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
     """
     Returns the intersection of two binary representations of Datasketches ThetaSketch
     objects, using a Datasketches Intersection object.
@@ -25981,9 +25965,6 @@ def theta_intersection(
     ----------
     col1 : :class:`~pyspark.sql.Column` or column name
     col2 : :class:`~pyspark.sql.Column` or column name
-    lgNomEntries : int, optional
-        The log-base-2 of nominal entries for the intersection operation
-        (must be between 4 and 26, defaults to 12)
 
     Returns
     -------
@@ -26005,30 +25986,19 @@ def theta_intersection(
     ...     sf.theta_sketch_agg("v2").alias("sketch2")
     ... )
     >>> df.select(sf.theta_sketch_estimate(sf.theta_intersection(df.sketch1, "sketch2"))).show()
-    +---------------------------------------------------------------+
-    |theta_sketch_estimate(theta_intersection(sketch1, sketch2, 12))|
-    +---------------------------------------------------------------+
-    |                                                              3|
-    +---------------------------------------------------------------+
+    +-----------------------------------------------------------+
+    |theta_sketch_estimate(theta_intersection(sketch1, sketch2))|
+    +-----------------------------------------------------------+
+    |                                                          3|
+    +-----------------------------------------------------------+
     """
-    from pyspark.sql.classic.column import _to_java_column
 
     fn = "theta_intersection"
-    if lgNomEntries is not None:
-        return _invoke_function(
-            fn,
-            _to_java_column(col1),
-            _to_java_column(col2),
-            _enum_to_value(lgNomEntries),
-        )
-    else:
-        return _invoke_function(fn, _to_java_column(col1), _to_java_column(col2))
+    return _invoke_function_over_columns(fn, col1, col2)
 
 
 @_try_remote_functions
-def theta_difference(
-    col1: "ColumnOrName", col2: "ColumnOrName", lgNomEntries: Optional[int] = None
-) -> Column:
+def theta_difference(col1: "ColumnOrName", col2: "ColumnOrName") -> Column:
     """
     Returns the set difference of two binary representations of Datasketches ThetaSketch
     objects (elements in first sketch but not in second), using a Datasketches ANotB object.
@@ -26039,9 +26009,6 @@ def theta_difference(
     ----------
     col1 : :class:`~pyspark.sql.Column` or column name
     col2 : :class:`~pyspark.sql.Column` or column name
-    lgNomEntries : int, optional
-        The log-base-2 of nominal entries for the difference operation
-        (must be between 4 and 26, defaults to 12)
 
     Returns
     -------
@@ -26064,24 +26031,15 @@ def theta_difference(
     ...     sf.theta_sketch_agg("v2").alias("sketch2")
     ... )
     >>> df.select(sf.theta_sketch_estimate(sf.theta_difference(df.sketch1, "sketch2"))).show()
-    +-------------------------------------------------------------+
-    |theta_sketch_estimate(theta_difference(sketch1, sketch2, 12))|
-    +-------------------------------------------------------------+
-    |                                                            3|
-    +-------------------------------------------------------------+
+    +---------------------------------------------------------+
+    |theta_sketch_estimate(theta_difference(sketch1, sketch2))|
+    +---------------------------------------------------------+
+    |                                                        3|
+    +---------------------------------------------------------+
     """
-    from pyspark.sql.classic.column import _to_java_column
 
     fn = "theta_difference"
-    if lgNomEntries is not None:
-        return _invoke_function(
-            fn,
-            _to_java_column(col1),
-            _to_java_column(col2),
-            _enum_to_value(lgNomEntries),
-        )
-    else:
-        return _invoke_function(fn, _to_java_column(col1), _to_java_column(col2))
+    return _invoke_function_over_columns(fn, col1, col2)
 
 
 # ---------------------- Predicates functions ------------------------------