From a34a4f4c8380825ac660759323b01678bfa57901 Mon Sep 17 00:00:00 2001 From: AuroraVoyage Date: Mon, 20 Jul 2026 10:26:07 +0800 Subject: [PATCH 1/4] [python] Introduce theta_sketch aggregator function. --- paimon-python/dev/requirements.txt | 2 + .../pypaimon/read/merge_engine_support.py | 1 + .../read/reader/aggregate/aggregators.py | 45 +++++++++++++++++++ .../pypaimon/tests/test_field_aggregators.py | 29 ++++++++++++ 4 files changed, 77 insertions(+) diff --git a/paimon-python/dev/requirements.txt b/paimon-python/dev/requirements.txt index dc9b4b4e911a..494082c5fbfc 100644 --- a/paimon-python/dev/requirements.txt +++ b/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.8" +datasketches>=5,<6; python_version >= "3.8" diff --git a/paimon-python/pypaimon/read/merge_engine_support.py b/paimon-python/pypaimon/read/merge_engine_support.py index cefcee2e6701..95d9010800d7 100644 --- a/paimon-python/pypaimon/read/merge_engine_support.py +++ b/paimon-python/pypaimon/read/merge_engine_support.py @@ -64,6 +64,7 @@ "bool_or", "bool_and", "listagg", "nested_update", + "theta_sketch", ]) _FIELDS_PREFIX = "fields." _FIELD_SEQUENCE_GROUP_SUFFIX = ".sequence-group" diff --git a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py b/paimon-python/pypaimon/read/reader/aggregate/aggregators.py index 11b182edaf44..00996c9987ee 100644 --- a/paimon-python/pypaimon/read/reader/aggregate/aggregators.py +++ b/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 @@ -60,6 +62,7 @@ NAME_BOOL_AND = "bool_and" NAME_LISTAGG = "listagg" NAME_NESTED_UPDATE = "nested_update" +NAME_THETA_SKETCH = "theta_sketch" # Base SQL type names treated as numeric for sum/product-style @@ -664,6 +667,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 @@ -740,3 +782,6 @@ def _factory(field_type, field_name, options): register_aggregator( NAME_NESTED_UPDATE, _build_field_options(FieldNestedUpdateAgg, NAME_NESTED_UPDATE) ) +register_aggregator( + NAME_THETA_SKETCH, _build_no_type_check(FieldThetaSketchAgg, NAME_THETA_SKETCH) +) diff --git a/paimon-python/pypaimon/tests/test_field_aggregators.py b/paimon-python/pypaimon/tests/test_field_aggregators.py index 9c9d5876fe9f..b106fcc2cb23 100644 --- a/paimon-python/pypaimon/tests/test_field_aggregators.py +++ b/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 @@ -46,6 +48,7 @@ FieldSumAgg, FieldListaggAgg, FieldNestedUpdateAgg, + FieldThetaSketchAgg, ) from pypaimon.schema.data_types import AtomicType, DataField, RowType, ArrayType from pypaimon.table.row.generic_row import GenericRow @@ -1558,6 +1561,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 From 090d3626cf7da1bf845e3f0d61014f4fc630b4dd Mon Sep 17 00:00:00 2001 From: AuroraVoyage Date: Mon, 20 Jul 2026 11:39:09 +0800 Subject: [PATCH 2/4] [python] Fix datasketches dependency installation in CI. --- .github/workflows/paimon-python-checks.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index af7c7883fa68..56687e661feb 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -130,7 +130,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. @@ -142,7 +142,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' pypaimon-rust==0.2.0 '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' pypaimon-rust==0.2.0 'datafusion>=52' datasketches python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION }}' -i https://pypi.org/simple/ if python -c "import sys; sys.exit(0 if sys.version_info >= (3, 11) else 1)"; then python -m pip install vortex-data==0.70.0 @@ -289,7 +289,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: | From a1612152981bdd464fa2d253257015fa40023809 Mon Sep 17 00:00:00 2001 From: AuroraVoyage Date: Mon, 20 Jul 2026 12:23:41 +0800 Subject: [PATCH 3/4] [python] Fix datasketches dependency installation in CI. --- .github/workflows/paimon-python-checks.yml | 2 +- paimon-python/dev/requirements.txt | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/paimon-python-checks.yml b/.github/workflows/paimon-python-checks.yml index 56687e661feb..23685a0822be 100755 --- a/.github/workflows/paimon-python-checks.yml +++ b/.github/workflows/paimon-python-checks.yml @@ -190,7 +190,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 diff --git a/paimon-python/dev/requirements.txt b/paimon-python/dev/requirements.txt index 494082c5fbfc..2725f24cfefa 100644 --- a/paimon-python/dev/requirements.txt +++ b/paimon-python/dev/requirements.txt @@ -42,5 +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.8" -datasketches>=5,<6; python_version >= "3.8" +datasketches>=4,<5; python_version <= "3.8" +datasketches>=5,<6; python_version > "3.8" From e1b90c0762cc3957298d2721cf5aa4eb058c736a Mon Sep 17 00:00:00 2001 From: AuroraVoyage Date: Tue, 21 Jul 2026 09:39:42 +0800 Subject: [PATCH 4/4] [python] Fix datasketches dependency installation in CI. --- paimon-python/dev/requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/paimon-python/dev/requirements.txt b/paimon-python/dev/requirements.txt index 2725f24cfefa..e2697f0b3953 100644 --- a/paimon-python/dev/requirements.txt +++ b/paimon-python/dev/requirements.txt @@ -42,5 +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.8" -datasketches>=5,<6; python_version > "3.8" +datasketches>=4,<5; python_version < "3.9" +datasketches>=5,<6; python_version >= "3.9"