[python] Introduce theta_sketch aggregator function

.github/workflows/paimon-python-checks.yml

@@ -131,7 +131,7 @@ jobs:
           if [[ "${{ matrix.python-version }}" == "3.6.15" ]]; then
             python -m pip install --upgrade pip==21.3.1
             python --version
-            python -m pip install --no-cache-dir pyroaring readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 2>&1 >/dev/null
+            python -m pip install --no-cache-dir pyroaring readerwriterlock==1.0.9 'fsspec==2021.10.1' 'cachetools==4.2.4' 'ossfs==2021.8.0' pyarrow==6.0.1 pandas==1.1.5 'polars==0.9.12' 'fastavro==1.4.7' zstandard==0.19.0 dataclasses==0.8.0 flake8 pytest py4j==0.10.9.9 requests parameterized==0.8.1 datasketches==4.1.0 2>&1 >/dev/null
             python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
           elif [[ "${{ matrix.python-version }}" == "3.7" ]]; then
             # 3.7 installs the version-pinned set declared for 3.7 in dev/requirements.txt.
@@ -143,7 +143,7 @@ jobs:
           else
             python -m pip install --upgrade pip
             pip install torch --index-url https://download.pytorch.org/whl/cpu
-            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 'isal>=1.8,<2' pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 cramjam flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 'daft>=0.7.6' 'datafusion>=52'
+            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 'isal>=1.8,<2' pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 cramjam flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 'daft>=0.7.6' 'datafusion>=52' datasketches
             if [[ "${{ matrix.python-version }}" == "3.11" ]]; then
               # Exercise the split-planning API in one lane until the compatible 0.3 wheel is published.
               python -m pip install "git+https://github.com/apache/paimon-rust.git@${PYPAIMON_RUST_REV}#subdirectory=bindings/python"
@@ -198,7 +198,7 @@ jobs:
         run: |
             python -m pip install --upgrade pip
             pip install torch --index-url https://download.pytorch.org/whl/cpu
-            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0
+            python -m pip install pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0 fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0 datasketches
             python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
       - name: Run lint-python.sh
         shell: bash
@@ -297,7 +297,7 @@ jobs:
             pyroaring readerwriterlock==1.0.9 fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 \
             fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2 \
             numpy==1.24.3 pandas==2.0.3 cramjam pytest~=7.0 py4j==0.10.9.9 requests \
-            parameterized==0.9.0 packaging
+            parameterized==0.9.0 packaging datasketches
           python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/
       - name: Test Ray version compatibility
         run: |

paimon-python/dev/requirements.txt

@@ -42,3 +42,5 @@ zstandard>=0.19,<1
 backports.zstd>=1.0.0,<1.4.0; python_version >= "3.9" and python_version < "3.14"
 cramjam>=1.3.0,<3; python_version>="3.7"
 pyyaml>=5.4,<7
+datasketches>=4,<5; python_version < "3.9"
+datasketches>=5,<6; python_version >= "3.9"

paimon-python/pypaimon/read/merge_engine_support.py

@@ -64,7 +64,8 @@
     "bool_or", "bool_and",
     "listagg",
     "nested_update",
-    "collect"
+    "collect",
+    "theta_sketch",
 ])
 _FIELDS_PREFIX = "fields."
 _FIELD_SEQUENCE_GROUP_SUFFIX = ".sequence-group"

paimon-python/pypaimon/read/reader/aggregate/aggregators.py

@@ -35,6 +35,8 @@
 
 from typing import Any, List, Dict, Optional, Tuple, Union
 
+from _datasketches import compact_theta_sketch, theta_union
+
 from pypaimon.common.options import CoreOptions
 from pypaimon.common.options.core_options import NestedKeyNullStrategy
 from pypaimon.read.reader.aggregate import register_aggregator
@@ -61,6 +63,7 @@
 NAME_LISTAGG = "listagg"
 NAME_NESTED_UPDATE = "nested_update"
 NAME_COLLECT = "collect"
+NAME_THETA_SKETCH = "theta_sketch"
 
 
 # Base SQL type names treated as numeric for sum/product-style
@@ -730,6 +733,45 @@ def _apply_nested_key_null_strategy(self, key: Tuple[Any, ...]) -> bool:
         )
 
 
+class FieldThetaSketchAgg(FieldAggregator):
+    """Aggregator for ThetaSketch."""
+
+    def __init__(self, name: str, field_type: DataType):
+        super().__init__(name, field_type)
+        if _atomic_base_name(field_type) not in ("VARBINARY", "BYTES"):
+            raise ValueError(
+                "Data type for theta sketch column must be 'VarBinaryType' but was '{}'.".format(field_type)
+            )
+
+    def agg(self, accumulator: Any, input_field: Any) -> Any:
+        if accumulator is None or input_field is None:
+            return input_field if accumulator is None else accumulator
+
+        if not isinstance(accumulator, (bytes, bytearray)):
+            raise TypeError(
+                "ThetaSketch accumulator must be bytes, got {}".format(type(accumulator))
+            )
+
+        if not isinstance(input_field, (bytes, bytearray)):
+            raise TypeError(
+                "ThetaSketch input must be bytes, got {}".format(type(input_field))
+            )
+
+        if isinstance(accumulator, bytearray):
+            accumulator = bytes(accumulator)
+        if isinstance(input_field, bytearray):
+            input_field = bytes(input_field)
+
+        sketch1 = compact_theta_sketch.deserialize(accumulator)
+        sketch2 = compact_theta_sketch.deserialize(input_field)
+
+        union = theta_union()
+        union.update(sketch1)
+        union.update(sketch2)
+
+        return union.get_result().serialize()
+
+
 # ---------------------------------------------------------------------------
 # Registration. Each builder binds an identifier to a factory that
 # optionally validates the column DataType before constructing the
@@ -809,3 +851,6 @@ def _factory(field_type, field_name, options):
 register_aggregator(
     NAME_COLLECT, _build_field_options(FieldCollectAgg, NAME_COLLECT)
 )
+register_aggregator(
+    NAME_THETA_SKETCH, _build_no_type_check(FieldThetaSketchAgg, NAME_THETA_SKETCH)
+)

paimon-python/pypaimon/tests/test_field_aggregators.py

@@ -30,6 +30,8 @@
 from functools import reduce
 from typing import List
 
+from _datasketches import update_theta_sketch
+
 from pypaimon.common.options import CoreOptions, Options
 from pypaimon.data import Timestamp
 from pypaimon.read.reader.aggregate import create_field_aggregator
@@ -47,6 +49,7 @@
     FieldListaggAgg,
     FieldNestedUpdateAgg,
     FieldCollectAgg,
+    FieldThetaSketchAgg,
 )
 from pypaimon.schema.data_types import AtomicType, DataField, RowType, ArrayType, MapType
 from pypaimon.table.row.generic_row import GenericRow
@@ -1765,6 +1768,32 @@ def test_field_nested_update_with_non_sequential_field_ids(self):
         self.assertCountEqual(accumulator, [self.row(0, 0, "A", 1), ])
 
 
+class FieldThetaSketchAggTest(unittest.TestCase):
+    @staticmethod
+    def sketch_of(*values: int) -> bytes:
+        sketch = update_theta_sketch()
+
+        for value in values:
+            sketch.update(value)
+
+        return sketch.compact().serialize()
+
+    def test_field_theta_sketch_agg(self):
+        agg = _make("theta_sketch", "VARBINARY(20)")
+        self.assertIsInstance(agg, FieldThetaSketchAgg)
+
+        input_val = self.sketch_of(1)
+        acc1 = self.sketch_of(2, 3)
+        acc2 = self.sketch_of(1, 2, 3)
+
+        self.assertIsNone(agg.agg(None, None))
+
+        self.assertEqual(agg.agg(None, input_val), input_val)
+        self.assertEqual(agg.agg(acc1, None), acc1)
+        self.assertEqual(agg.agg(acc1, input_val), acc2)
+        self.assertEqual(agg.agg(acc2, input_val), acc2)
+
+
 class RegistrationTest(unittest.TestCase):
     """Sanity check that all 10 expected aggregators (the primary-key
     placeholder plus 9 value aggregators) are registered when the