# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""
Three-layer correctness tests for predicate-driven bucket pruning.

Mirrors Java's ``BucketSelectConverter`` contract: PK Equal/In queries on
HASH_FIXED tables must touch only the bucket(s) the writer would have
placed those keys in. Two correctness obligations:

  1. Sound: every bucket retained by the selector contains AT MOST a
     superset of matching rows. Buckets that DO contain matching rows
     are NEVER dropped — false-negative-free.
  2. Hash-consistent with writers: ``RowKeyExtractor`` (writer) and
     ``BucketSelectConverter`` (reader) must agree on every literal.
     This is what makes ``pk = 'X'`` read the bucket holding 'X'.

Layered:
  * Unit       — direct calls to ``create_bucket_selector`` with crafted
                 predicates, asserting selector behaviour.
  * Integration — real PK tables with multiple buckets; queries; assert
                 (a) result correctness, (b) bucket pruning happened.
  * Property   — randomly-seeded PK tables, random Equal/In predicates,
                 result == oracle. No hypothesis dependency (keeps
                 Python 3.6 compat).
"""

import os
import random
import shutil
import tempfile
import unittest
from typing import Any, Dict, List

import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.common.predicate_builder import PredicateBuilder
from pypaimon.read.scanner.bucket_select_converter import (
    MAX_VALUES, create_bucket_selector)
from pypaimon.schema.data_types import AtomicType, DataField
from pypaimon.write.row_key_extractor import (FixedBucketRowKeyExtractor,
                                              _bucket_from_hash,
                                              _hash_bytes_by_words)
from pypaimon.table.row.generic_row import GenericRow, GenericRowSerializer
from pypaimon.table.row.internal_row import RowKind


def _bigint_field(idx: int, name: str) -> DataField:
    return DataField(idx, name, AtomicType('BIGINT', nullable=False))


def _field(idx: int, name: str, type_name: str) -> DataField:
    return DataField(idx, name, AtomicType(type_name, nullable=False))


def _hash_bucket(values: List[Any], fields: List[DataField], total: int) -> int:
    """Re-implement the writer's hash so unit tests can compute the
    expected bucket without spinning up a real table."""
    row = GenericRow(values, fields, RowKind.INSERT)
    serialized = GenericRowSerializer.to_bytes(row)
    h = _hash_bytes_by_words(serialized[4:])
    return _bucket_from_hash(h, total)


# ---------------------------------------------------------------------------
# Layer 1 — Unit: drive ``create_bucket_selector`` with crafted predicates.
# ---------------------------------------------------------------------------
class BucketSelectConverterUnitTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.id_field = _bigint_field(0, 'id')
        cls.val_field = _bigint_field(1, 'val')
        cls.k1 = _bigint_field(0, 'k1')
        cls.k2 = _bigint_field(1, 'k2')
        cls.pb_id_val = PredicateBuilder([cls.id_field, cls.val_field])
        cls.pb_k1_k2 = PredicateBuilder([cls.k1, cls.k2])

    # -- Equal / In on single bucket key ---------------------------------
    def test_equal_on_single_bucket_key_yields_single_bucket(self):
        sel = create_bucket_selector(
            self.pb_id_val.equal('id', 42), [self.id_field])
        self.assertIsNotNone(sel, "PK Equal must produce a selector")
        expected = _hash_bucket([42], [self.id_field], total=8)
        for b in range(8):
            self.assertEqual(
                sel(b, 8), b == expected,
                "only bucket {} must be kept (got {})".format(expected, b))

    def test_in_on_single_bucket_key_unions_buckets(self):
        sel = create_bucket_selector(
            self.pb_id_val.is_in('id', [1, 2, 3, 100]), [self.id_field])
        expected = {_hash_bucket([v], [self.id_field], 8)
                    for v in (1, 2, 3, 100)}
        for b in range(8):
            self.assertEqual(sel(b, 8), b in expected)

    def test_or_of_equals_on_same_field_unions_buckets(self):
        # ``id = 1 OR id = 2`` must equal ``id IN (1, 2)``.
        pred = PredicateBuilder.or_predicates([
            self.pb_id_val.equal('id', 1),
            self.pb_id_val.equal('id', 2),
        ])
        sel = create_bucket_selector(pred, [self.id_field])
        expected = {_hash_bucket([v], [self.id_field], 8) for v in (1, 2)}
        for b in range(8):
            self.assertEqual(sel(b, 8), b in expected)

    # -- Composite bucket keys ------------------------------------------
    def test_composite_bucket_key_intersects_via_cartesian(self):
        pred = PredicateBuilder.and_predicates([
            self.pb_k1_k2.is_in('k1', [1, 2]),
            self.pb_k1_k2.equal('k2', 99),
        ])
        sel = create_bucket_selector(pred, [self.k1, self.k2])
        expected = {
            _hash_bucket([k1, 99], [self.k1, self.k2], 4)
            for k1 in (1, 2)
        }
        for b in range(4):
            self.assertEqual(sel(b, 4), b in expected)

    def test_composite_bucket_key_missing_one_field_returns_none(self):
        pred = self.pb_k1_k2.equal('k1', 1)  # k2 unconstrained
        sel = create_bucket_selector(pred, [self.k1, self.k2])
        self.assertIsNone(sel,
                          "all bucket keys must be constrained or fall back")

    # -- Predicates that can't be reduced -------------------------------
    def test_non_bucket_key_predicate_returns_none(self):
        sel = create_bucket_selector(
            self.pb_id_val.equal('val', 5), [self.id_field])
        self.assertIsNone(sel, "predicate not on bucket key -> no selector")

    def test_range_predicate_on_bucket_key_returns_none(self):
        sel = create_bucket_selector(
            self.pb_id_val.greater_than('id', 100), [self.id_field])
        self.assertIsNone(sel, "ranges can't be turned into a finite bucket set")

    def test_or_with_non_bucket_key_returns_none(self):
        # ``id = 1 OR val = 5`` — ``val`` isn't a bucket key, so the OR
        # is not a pure bucket-key constraint.
        pred = PredicateBuilder.or_predicates([
            self.pb_id_val.equal('id', 1),
            self.pb_id_val.equal('val', 5),
        ])
        sel = create_bucket_selector(pred, [self.id_field])
        self.assertIsNone(sel)

    def test_repeated_equal_on_same_key_with_empty_intersection_returns_none(self):
        # ``id = 1 AND id = 2``: literal sets {1} and {2} intersect to
        # empty; Java's ``retainAll`` would also bail here, since the
        # predicate is unsatisfiable.
        pred = PredicateBuilder.and_predicates([
            self.pb_id_val.equal('id', 1),
            self.pb_id_val.equal('id', 2),
        ])
        sel = create_bucket_selector(pred, [self.id_field])
        self.assertIsNone(sel)

    def test_repeated_in_on_same_key_intersects_literals(self):
        # ``id IN (1,2,3) AND id IN (2,3,4)`` should now keep the
        # intersection {2, 3} and prune to those buckets only. Used to
        # bail with no selector before the Java parity fix.
        pred = PredicateBuilder.and_predicates([
            self.pb_id_val.is_in('id', [1, 2, 3]),
            self.pb_id_val.is_in('id', [2, 3, 4]),
        ])
        sel = create_bucket_selector(pred, [self.id_field])
        self.assertIsNotNone(sel)
        expected = {_hash_bucket([v], [self.id_field], 8) for v in (2, 3)}
        for b in range(8):
            self.assertEqual(sel(b, 8), b in expected)

    def test_and_with_unrelated_clause_is_unaffected(self):
        # ``id = 7 AND val > 100`` — the ``val > 100`` part doesn't
        # constrain buckets, but mustn't disqualify the ``id = 7`` part.
        pred = PredicateBuilder.and_predicates([
            self.pb_id_val.equal('id', 7),
            self.pb_id_val.greater_than('val', 100),
        ])
        sel = create_bucket_selector(pred, [self.id_field])
        self.assertIsNotNone(sel)
        expected = _hash_bucket([7], [self.id_field], 4)
        for b in range(4):
            self.assertEqual(sel(b, 4), b == expected)

    # -- Cap & degenerate edge cases ------------------------------------
    def test_cartesian_above_max_values_returns_none(self):
        # Two columns of size > sqrt(MAX_VALUES) → product > MAX_VALUES.
        size = 33  # 33 * 33 = 1089 > 1000
        pred = PredicateBuilder.and_predicates([
            self.pb_k1_k2.is_in('k1', list(range(size))),
            self.pb_k1_k2.is_in('k2', list(range(size))),
        ])
        self.assertGreater(size * size, MAX_VALUES)
        sel = create_bucket_selector(pred, [self.k1, self.k2])
        self.assertIsNone(sel)

    def test_null_only_literal_drops_everything(self):
        # ``id IN (NULL)`` after null-stripping has zero literals; the
        # cartesian product is empty → selector matches no buckets. Same
        # behaviour as Java.
        pred = self.pb_id_val.is_in('id', [None])
        sel = create_bucket_selector(pred, [self.id_field])
        self.assertIsNotNone(sel)
        for b in range(4):
            self.assertFalse(sel(b, 4),
                             "all-null literal collapses bucket set to empty")

    def test_no_predicate_returns_none(self):
        self.assertIsNone(create_bucket_selector(None, [self.id_field]))

    def test_no_bucket_keys_returns_none(self):
        self.assertIsNone(
            create_bucket_selector(self.pb_id_val.equal('id', 1), []))

    # -- Selector cache + rescale -------------------------------------
    def test_selector_caches_per_total_buckets(self):
        """Selector must answer correctly when the same query applies to
        different ``total_buckets`` values (the rescale scenario)."""
        sel = create_bucket_selector(
            self.pb_id_val.equal('id', 42), [self.id_field])
        for total in (4, 8, 16, 32):
            expected = _hash_bucket([42], [self.id_field], total)
            self.assertTrue(sel(expected, total))
            other = (expected + 1) % total
            self.assertFalse(sel(other, total))

    def test_non_positive_total_buckets_fails_open(self):
        """Manifest entries can carry ``total_buckets <= 0`` for legacy /
        special bucket modes. Pruning MUST fail open — returning False
        would silently drop rows the writer placed in those entries.
        This is correctness, not performance: the soundness contract
        forbids false-negatives."""
        sel = create_bucket_selector(
            self.pb_id_val.equal('id', 1), [self.id_field])
        for total in (0, -1, -2):
            self.assertTrue(sel(0, total),
                            "total_buckets={} must be kept (fail open)".format(total))
            self.assertTrue(sel(-1, total))
            self.assertTrue(sel(99, total))

    # -- Bucket-key column types beyond BIGINT --------------------------
    def test_string_bucket_key_yields_correct_bucket(self):
        """STRING uses a different ``GenericRowSerializer`` path (utf-8
        encode + variable-part offset) — verify writer/reader agree on
        its byte layout independent of the BIGINT happy path."""
        sf = _field(0, 'sk', 'STRING')
        vf = _bigint_field(1, 'val')
        pb = PredicateBuilder([sf, vf])
        sel = create_bucket_selector(pb.equal('sk', 'hello'), [sf])
        self.assertIsNotNone(sel)
        expected = _hash_bucket(['hello'], [sf], total=8)
        for b in range(8):
            self.assertEqual(sel(b, 8), b == expected)

    def test_int_bucket_key_yields_correct_bucket(self):
        """INT (32-bit) and BIGINT (64-bit) hit different struct.pack
        paths in the serializer — guard the smaller width."""
        intf = _field(0, 'i', 'INT')
        vf = _bigint_field(1, 'val')
        pb = PredicateBuilder([intf, vf])
        sel = create_bucket_selector(pb.equal('i', 7), [intf])
        self.assertIsNotNone(sel)
        expected = _hash_bucket([7], [intf], total=4)
        for b in range(4):
            self.assertEqual(sel(b, 4), b == expected)

    # -- Hash-divergence-prone types refuse to build a selector --------
    def test_decimal_bucket_key_disables_pruning(self):
        """DECIMAL columns risk silent hash divergence between writer
        (Decimal) and reader-supplied ``float`` literals. Soundness
        contract demands fail-open: refuse to build a selector at all."""
        df = _field(0, 'd', 'DECIMAL(10, 2)')
        vf = _bigint_field(1, 'val')
        pb = PredicateBuilder([df, vf])
        from decimal import Decimal
        sel = create_bucket_selector(pb.equal('d', Decimal('1.50')), [df])
        self.assertIsNone(
            sel, "DECIMAL bucket-key column must disable pruning")

    def test_array_bucket_key_disables_pruning(self):
        """Composite types (ARRAY/MAP/ROW/MULTISET/VARIANT/BLOB) have no
        cross-validated byte alignment with Java's ``BinaryRow`` — until
        that exists, refuse to prune on them."""
        # Hand-roll a DataField whose AtomicType reports an ARRAY type
        # name; the converter inspects ``field.type.type`` only.
        af = DataField(0, 'arr', AtomicType('ARRAY<BIGINT>'))
        vf = _bigint_field(1, 'val')
        pb = PredicateBuilder([af, vf])
        sel = create_bucket_selector(pb.equal('arr', [1]), [af])
        self.assertIsNone(
            sel, "ARRAY bucket-key column must disable pruning")

    def test_timestamp_bucket_key_disables_pruning(self):
        """TIMESTAMP columns serialise via ``value.timestamp()`` whose
        result depends on the process timezone for naive datetimes —
        writer and reader running in different TZs would disagree."""
        tf = _field(0, 't', 'TIMESTAMP(3)')
        vf = _bigint_field(1, 'val')
        pb = PredicateBuilder([tf, vf])
        from datetime import datetime
        sel = create_bucket_selector(
            pb.equal('t', datetime(2026, 1, 1)), [tf])
        self.assertIsNone(
            sel, "TIMESTAMP bucket-key column must disable pruning")

    def test_type_mismatched_literal_fails_open_not_crash(self):
        """If the user constructs a predicate whose literal type doesn't
        match the bucket-key column's atomic type — e.g. a STRING literal
        on a BIGINT column — ``GenericRowSerializer`` raises during the
        deferred hash inside ``_Selector``. The selector MUST swallow the
        exception and fail open (return True for every bucket) rather
        than propagate it. Crashing the entire scan with an opaque
        ``struct.error`` is a worse user experience than silently
        skipping bucket pruning, and the soundness contract still
        forbids false-negatives."""
        sel = create_bucket_selector(
            self.pb_id_val.equal('id', 'not-an-int'), [self.id_field])
        # Construction itself succeeds (no eager hashing).
        self.assertIsNotNone(sel)
        # Calling the selector must NOT raise; instead it returns True
        # for every (bucket, total_buckets), preserving soundness.
        for total in (4, 8):
            for b in range(total):
                self.assertTrue(sel(b, total),
                                "type-mismatched literal must fail open, "
                                "not crash (bucket={}, total={})".format(b, total))


class PartitionAwareBucketSelectorUnitTest(unittest.TestCase):
    """Unit tests for the per-partition predicate specialisation path.

    Covers ``replace_partition_predicate`` (the AND/OR fold walker) and
    the partition-aware ``_Selector.__call__(partition, bucket,
    total_buckets)`` 3-arg form that ``FileScanner._filter_manifest_entry``
    will use after wiring."""

    @classmethod
    def setUpClass(cls):
        cls.id_field = _bigint_field(0, 'id')
        cls.part_field = DataField(2, 'part', AtomicType('STRING'))
        cls.pb = PredicateBuilder([cls.id_field, cls.part_field])

    # ----- replace_partition_predicate --------------------------------

    def test_replace_partition_leaf_to_true_drops_constraint(self):
        from pypaimon.read.scanner.bucket_select_converter import \
            replace_partition_predicate
        # ``part = 'a' AND id = 1`` against partition {part: 'a'} →
        # part leaf becomes True → AND fold removes it → only ``id = 1``
        pred = PredicateBuilder.and_predicates([
            self.pb.equal('part', 'a'),
            self.pb.equal('id', 1),
        ])
        result = replace_partition_predicate(
            pred, {'part'}, {'part': 'a'})
        self.assertTrue(isinstance(result, type(pred)),
                        "AND should fold to a remaining single leaf")
        self.assertEqual(result.method, 'equal')
        self.assertEqual(result.field, 'id')

    def test_replace_partition_leaf_to_false_collapses_and(self):
        from pypaimon.read.scanner.bucket_select_converter import \
            replace_partition_predicate
        # ``part = 'a' AND id = 1`` against partition {part: 'b'} →
        # part leaf becomes False → AND collapses to AlwaysFalse (False).
        pred = PredicateBuilder.and_predicates([
            self.pb.equal('part', 'a'),
            self.pb.equal('id', 1),
        ])
        result = replace_partition_predicate(
            pred, {'part'}, {'part': 'b'})
        self.assertIs(result, False)

    def test_replace_partition_leaf_in_or_keeps_other_branch(self):
        from pypaimon.read.scanner.bucket_select_converter import \
            replace_partition_predicate
        # ``(part='a' AND id=1) OR (part='b' AND id=2)`` against
        # partition {part: 'a'} → first OR child becomes ``id=1``, second
        # collapses to AlwaysFalse and is dropped. Result is just ``id=1``.
        pred = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'a'),
                self.pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'b'),
                self.pb.equal('id', 2),
            ]),
        ])
        result = replace_partition_predicate(
            pred, {'part'}, {'part': 'a'})
        # OR with a single surviving child unwraps to that child.
        self.assertEqual(result.method, 'equal')
        self.assertEqual(result.field, 'id')
        self.assertEqual(result.literals, [1])

    def test_replace_partition_leaf_in_or_other_partition(self):
        from pypaimon.read.scanner.bucket_select_converter import \
            replace_partition_predicate
        # Same predicate, partition {part: 'b'} → second branch survives
        # as ``id=2``.
        pred = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'a'),
                self.pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'b'),
                self.pb.equal('id', 2),
            ]),
        ])
        result = replace_partition_predicate(
            pred, {'part'}, {'part': 'b'})
        self.assertEqual(result.method, 'equal')
        self.assertEqual(result.field, 'id')
        self.assertEqual(result.literals, [2])

    def test_replace_partition_leaf_unrelated_predicate_unchanged(self):
        from pypaimon.read.scanner.bucket_select_converter import \
            replace_partition_predicate
        # No partition leaf → predicate returned as-is.
        pred = self.pb.equal('id', 42)
        result = replace_partition_predicate(
            pred, {'part'}, {'part': 'a'})
        self.assertIs(result, pred)

    # ----- _Selector partition-aware path -----------------------------

    def test_selector_3arg_specialises_per_partition(self):
        # ``(part='a' AND id=1) OR (part='b' AND id=2)`` should hit
        # bucket(1) only when partition='a' and bucket(2) only when
        # partition='b'. Master without this fix returns "all buckets".
        pred = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'a'),
                self.pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'b'),
                self.pb.equal('id', 2),
            ]),
        ])
        sel = create_bucket_selector(
            pred, [self.id_field], partition_fields=[self.part_field])
        self.assertIsNotNone(sel)
        bucket_for_1 = _hash_bucket([1], [self.id_field], total=8)
        bucket_for_2 = _hash_bucket([2], [self.id_field], total=8)

        part_a = GenericRow(['a'], [self.part_field], RowKind.INSERT)
        part_b = GenericRow(['b'], [self.part_field], RowKind.INSERT)

        for b in range(8):
            self.assertEqual(sel(part_a, b, 8), b == bucket_for_1,
                             "partition a only keeps bucket %d" % bucket_for_1)
            self.assertEqual(sel(part_b, b, 8), b == bucket_for_2,
                             "partition b only keeps bucket %d" % bucket_for_2)

    def test_selector_falls_through_when_partition_unknown(self):
        """Early manifest filter passes ``partition=None`` (or uses the
        2-arg form) — no specialisation runs, bucket set falls back to a
        sound over-approximation: all buckets accept."""
        pred = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'a'),
                self.pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'b'),
                self.pb.equal('id', 2),
            ]),
        ])
        sel = create_bucket_selector(
            pred, [self.id_field], partition_fields=[self.part_field])
        self.assertIsNotNone(sel)
        # 2-arg form (legacy callsite) — partition unknown, all buckets keep.
        for b in range(8):
            self.assertTrue(sel(b, 8),
                            "partition-unknown call must accept all buckets")
        # 3-arg form with partition=None has the same semantics.
        for b in range(8):
            self.assertTrue(sel(None, b, 8))

    def test_selector_partition_not_matching_returns_empty_bucket_set(self):
        # ``part = 'a' AND id = 1`` on partition {part: 'c'} simplifies to
        # AlwaysFalse — the selector returns False for every bucket since
        # no row in this partition can possibly match. Sound: dropping a
        # partition that *can't* contain matches doesn't lose data.
        pred = PredicateBuilder.and_predicates([
            self.pb.equal('part', 'a'),
            self.pb.equal('id', 1),
        ])
        sel = create_bucket_selector(
            pred, [self.id_field], partition_fields=[self.part_field])
        self.assertIsNotNone(sel)
        part_c = GenericRow(['c'], [self.part_field], RowKind.INSERT)
        for b in range(8):
            self.assertFalse(sel(part_c, b, 8),
                             "partition c can't satisfy part='a', "
                             "drop every bucket (b=%d)" % b)

    def test_selector_partition_only_constraint_drops_partition(self):
        # ``part='a' AND id IN (1,2)`` — same partition value 'a'
        # specialises ``part='a'`` to True, leaving ``id IN (1,2)``.
        pred = PredicateBuilder.and_predicates([
            self.pb.equal('part', 'a'),
            self.pb.is_in('id', [1, 2]),
        ])
        sel = create_bucket_selector(
            pred, [self.id_field], partition_fields=[self.part_field])
        self.assertIsNotNone(sel)
        part_a = GenericRow(['a'], [self.part_field], RowKind.INSERT)
        expected = {_hash_bucket([v], [self.id_field], 8) for v in (1, 2)}
        for b in range(8):
            self.assertEqual(sel(part_a, b, 8), b in expected)

    def test_selector_caches_per_partition(self):
        pred = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'a'),
                self.pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                self.pb.equal('part', 'b'),
                self.pb.equal('id', 2),
            ]),
        ])
        sel = create_bucket_selector(
            pred, [self.id_field], partition_fields=[self.part_field])
        part_a = GenericRow(['a'], [self.part_field], RowKind.INSERT)
        part_b = GenericRow(['b'], [self.part_field], RowKind.INSERT)
        # Drive the cache.
        for _ in range(5):
            sel(part_a, 0, 8)
            sel(part_b, 0, 8)
        # Cache keyed by (partition tuple, total_buckets); two distinct
        # partitions × one total → exactly two entries.
        self.assertEqual(len(sel._cache), 2)


# ---------------------------------------------------------------------------
# Layer 2 — Integration: real tables, public API, assert correctness AND
# that pruning actually fired (otherwise we're not testing the optimisation,
# only that we didn't break full-scan).
# ---------------------------------------------------------------------------
class BucketPruningIntegrationTest(unittest.TestCase):

    NUM_BUCKETS = 8

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', False)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _create_pk_table(self, name: str, num_buckets: int = NUM_BUCKETS,
                         bucket_key: str = None) -> Any:
        opts = {'bucket': str(num_buckets), 'file.format': 'parquet'}
        if bucket_key is not None:
            opts['bucket-key'] = bucket_key
        pa_schema = pa.schema([
            pa.field('id', pa.int64(), nullable=False),
            ('val', pa.int64()),
        ])
        schema = Schema.from_pyarrow_schema(
            pa_schema, primary_keys=['id'], options=opts)
        full = 'default.{}'.format(name)
        self.catalog.create_table(full, schema, False)
        return self.catalog.get_table(full)

    def _write(self, table, rows: List[Dict]):
        pa_schema = pa.schema([
            pa.field('id', pa.int64(), nullable=False),
            ('val', pa.int64()),
        ])
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        try:
            w.write_arrow(pa.Table.from_pylist(rows, schema=pa_schema))
            c.commit(w.prepare_commit())
        finally:
            w.close()
            c.close()

    def _read_with(self, table, predicate=None):
        rb = table.new_read_builder()
        if predicate is not None:
            rb = rb.with_filter(predicate)
        splits = rb.new_scan().plan().splits()
        if not splits:
            return [], splits
        return rb.new_read().to_arrow(splits).to_pylist(), splits

    @staticmethod
    def _split_buckets(splits) -> set:
        """Collect the distinct bucket numbers actually returned in a plan."""
        return {s.bucket for s in splits}

    @staticmethod
    def _expected_buckets(table, ids, value_field='val') -> set:
        """Use the writer's RowKeyExtractor to compute the bucket(s) the
        rows for ``ids`` were written into. Cross-check against the
        reader's selector — divergence indicates read/write hash drift."""
        ext = FixedBucketRowKeyExtractor(table.table_schema)
        pa_schema = pa.schema([
            pa.field('id', pa.int64(), nullable=False),
            (value_field, pa.int64()),
        ])
        out = set()
        for i in ids:
            arr = pa.RecordBatch.from_pylist(
                [{'id': i, value_field: 0}], schema=pa_schema)
            out.update(ext._extract_buckets_batch(arr))
        return out

    # -- Equal on PK -----------------------------------------------------
    def test_pk_equal_only_reads_target_bucket(self):
        table = self._create_pk_table('int_eq')
        rows = [{'id': i, 'val': i * 11} for i in range(100)]
        self._write(table, rows)

        target_id = 42
        pred = table.new_read_builder().new_predicate_builder().equal(
            'id', target_id)
        got, splits = self._read_with(table, pred)

        # Correctness: row for id=42 returned (and only that).
        self.assertEqual(got, [{'id': 42, 'val': 42 * 11}])

        # Pruning effectiveness AND hash correctness: the touched bucket
        # must equal the bucket the writer placed id=42 into. Asserting
        # only ``len == 1`` would mask a hash drift that picks the wrong
        # single bucket.
        self.assertEqual(self._split_buckets(splits),
                         self._expected_buckets(table, [target_id]),
                         "PK equal must touch exactly the writer's bucket")

    def test_pk_in_reads_only_target_buckets(self):
        table = self._create_pk_table('int_in')
        rows = [{'id': i, 'val': i * 7} for i in range(200)]
        self._write(table, rows)

        ids = [3, 17, 99, 150]
        pred = table.new_read_builder().new_predicate_builder().is_in(
            'id', ids)
        got, splits = self._read_with(table, pred)
        got_sorted = sorted(got, key=lambda r: r['id'])
        self.assertEqual(got_sorted,
                         [{'id': i, 'val': i * 7} for i in sorted(ids)])

        actual = self._split_buckets(splits)
        expected_buckets = self._expected_buckets(table, ids)
        # Equality (not subset): under the single-commit setup every
        # target bucket actually has a file, so the planner must produce
        # exactly the writer's bucket set. ``issubset`` would mask a
        # selector that's overly aggressive on a subset of the IN list.
        self.assertEqual(actual, expected_buckets,
                         "got {}, expected {}".format(actual, expected_buckets))

    # -- Predicates that should NOT prune -------------------------------
    def test_value_only_predicate_falls_back_to_full_scan(self):
        """``val < X`` doesn't constrain the PK → selector must be None
        and no bucket pruning may fire. Both checked: result correctness
        AND the explicit "selector is None" property."""
        table = self._create_pk_table('val_only')
        rows = [{'id': i, 'val': i} for i in range(100)]
        self._write(table, rows)

        pred = table.new_read_builder().new_predicate_builder().less_than(
            'val', 30)
        got, splits = self._read_with(table, pred)
        self.assertEqual(sorted([r['id'] for r in got]), list(range(30)))

        # Inspect the scanner's bucket selector to prove pruning DIDN'T
        # fire — without this assertion the test would also pass under a
        # buggy selector that prunes wrongly but happens to keep the
        # rows we picked.
        rb = table.new_read_builder().with_filter(pred)
        scan = rb.new_scan()
        self.assertIsNone(scan.file_scanner._bucket_selector,
                          "value-only predicate must NOT produce a selector")

    def test_range_on_pk_falls_back_to_full_scan(self):
        """``id > X`` is a range, not Equal/In, so cannot derive a bucket
        set. Selector returns None — result must still be exact."""
        table = self._create_pk_table('pk_range')
        rows = [{'id': i, 'val': i} for i in range(50)]
        self._write(table, rows)

        pred = table.new_read_builder().new_predicate_builder().greater_or_equal(
            'id', 40)
        got, _ = self._read_with(table, pred)
        self.assertEqual(sorted([r['id'] for r in got]), list(range(40, 50)))

    # -- Mixed predicate: Equal on PK AND range on val ------------------
    def test_pk_equal_with_unrelated_value_predicate_still_prunes(self):
        table = self._create_pk_table('int_eq_with_val')
        rows = [{'id': i, 'val': i} for i in range(50)]
        self._write(table, rows)

        pb = table.new_read_builder().new_predicate_builder()
        pred = pb.and_predicates([
            pb.equal('id', 25),
            pb.greater_than('val', 20),
        ])
        got, splits = self._read_with(table, pred)
        self.assertEqual(got, [{'id': 25, 'val': 25}])
        self.assertEqual(self._split_buckets(splits),
                         self._expected_buckets(table, [25]),
                         "Equal on PK still narrows to the writer's bucket "
                         "even when AND'd with a non-bucket-key predicate")

    def test_early_filter_skips_full_entry_decode_for_pruned_buckets(self):
        """Entries the bucket selector rejects must never reach
        ``GenericRowDeserializer.from_bytes`` for their partition / key
        stats. Without the early filter the count would scale with the
        manifest entry count; with it, only the surviving entries pay
        the deserialisation cost."""
        from unittest import mock

        from pypaimon.table.row import generic_row

        table = self._create_pk_table('early_filter')
        # 8 separate single-row commits → 8 manifest entries each touching
        # a different bucket. ``pk = X`` should reach exactly one of them.
        for i in range(self.NUM_BUCKETS):
            self._write(table, [{'id': i, 'val': i * 11}])

        pred = table.new_read_builder().new_predicate_builder().equal('id', 0)
        rb = table.new_read_builder().with_filter(pred)

        real_from_bytes = generic_row.GenericRowDeserializer.from_bytes
        calls = {'n': 0}

        def counting(*args, **kwargs):
            calls['n'] += 1
            return real_from_bytes(*args, **kwargs)

        with mock.patch.object(generic_row.GenericRowDeserializer,
                               'from_bytes',
                               side_effect=counting):
            splits = rb.new_scan().plan().splits()
            got = rb.new_read().to_arrow(splits).to_pylist() if splits else []

        self.assertEqual(got, [{'id': 0, 'val': 0}])
        # Each surviving entry decodes partition + min_key + max_key
        # (3 ``from_bytes`` calls). Allow a small slack in case the planner
        # touches extras, but assert it is well below 8 entries × 3 = 24.
        self.assertLess(
            calls['n'], 3 * self.NUM_BUCKETS,
            "early filter should skip from_bytes for pruned entries; "
            "got {} calls (would be {}+ without the filter)".format(
                calls['n'], 3 * self.NUM_BUCKETS))

    def test_init_bucket_selector_fails_open_when_bucket_keys_raises(self):
        """``TableSchema.bucket_keys`` raises if ``bucket-key`` references
        an unknown column. The pre-Java-alignment selector path used to
        catch ``Exception`` from instantiating ``FixedBucketRowKeyExtractor``
        and silently skip pruning; that property must survive the move
        of bucket-key resolution onto ``TableSchema``. Crashing the scan
        on a misconfiguration would be worse than skipping the
        optimisation."""
        table = self._create_pk_table('init_fails_open')
        self._write(table, [{'id': 1, 'val': 1}])
        # Mutate the in-memory schema options to a broken value to
        # simulate a corrupted/migrated catalog without rewriting it.
        table.table_schema.options['bucket-key'] = 'nope_no_such_column'

        rb = table.new_read_builder().with_filter(
            table.new_read_builder().new_predicate_builder().equal('id', 1))
        scanner = rb.new_scan().file_scanner
        # Must NOT raise: the broken option falls back to "no pruning",
        # and the scan still finds the row.
        self.assertIsNone(scanner._init_bucket_selector())
        got, _ = self._read_with(table, scanner.predicate)
        self.assertEqual(got, [{'id': 1, 'val': 1}])

    # -- Explicit bucket-key option ------------------------------------
    def test_bucket_key_option_overrides_pk_for_pruning(self):
        """When the ``bucket-key`` option is set explicitly, the bucket
        derivation must use it — not the trimmed primary keys. This is
        the path that catches read/write hash divergence if a refactor
        forgets the option."""
        # PK = id, bucket-key = id explicitly (single key but exercises
        # the explicit-config branch in ``_init_bucket_selector``).
        table = self._create_pk_table('explicit_bk', bucket_key='id')
        rows = [{'id': i, 'val': i * 3} for i in range(40)]
        self._write(table, rows)

        pred = table.new_read_builder().new_predicate_builder().equal('id', 17)
        got, splits = self._read_with(table, pred)
        self.assertEqual(got, [{'id': 17, 'val': 51}])
        self.assertEqual(self._split_buckets(splits),
                         self._expected_buckets(table, [17]))

    def test_per_partition_pruning_with_mixed_or(self):
        """``(part='a' AND id=1) OR (part='b' AND id=2)``: each partition
        sees only the bucket for its own ``id`` literal. Without
        per-partition predicate specialisation this query falls through
        to "all buckets in both partitions"."""
        opts = {'bucket': '4', 'file.format': 'parquet'}
        pa_schema = pa.schema([
            pa.field('part', pa.string(), nullable=False),
            pa.field('id', pa.int64(), nullable=False),
            ('val', pa.int64()),
        ])
        schema = Schema.from_pyarrow_schema(
            pa_schema, primary_keys=['part', 'id'],
            partition_keys=['part'], options=opts)
        identifier = 'default.per_part_mixed_or'
        self.catalog.create_table(identifier, schema, False)
        table = self.catalog.get_table(identifier)
        # Two partitions × three id values each → up to 6 (part, bucket)
        # combinations after the writer hashes.
        rows = []
        for p in ('a', 'b'):
            for i in (1, 2, 3):
                rows.append({'part': p, 'id': i, 'val': i * 7})
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        try:
            w.write_arrow(pa.Table.from_pylist(rows, schema=pa_schema))
            c.commit(w.prepare_commit())
        finally:
            w.close()
            c.close()

        pb = table.new_read_builder().new_predicate_builder()
        from pypaimon.common.predicate_builder import PredicateBuilder
        mixed = PredicateBuilder.or_predicates([
            PredicateBuilder.and_predicates([
                pb.equal('part', 'a'),
                pb.equal('id', 1),
            ]),
            PredicateBuilder.and_predicates([
                pb.equal('part', 'b'),
                pb.equal('id', 2),
            ]),
        ])
        got, splits = self._read_with(table, mixed)
        # Correctness: only the two matching rows.
        got_sorted = sorted(got, key=lambda r: (r['part'], r['id']))
        self.assertEqual(
            got_sorted,
            [{'part': 'a', 'id': 1, 'val': 7},
             {'part': 'b', 'id': 2, 'val': 14}])
        # Pruning effectiveness: across both partitions we should see at
        # most two distinct (partition, bucket) splits — one per branch.
        # Without per-partition pruning we'd see every (partition, bucket)
        # combo that exists on disk for the predicate's id literals.
        self.assertLessEqual(len(splits), 2,
                             "per-partition pruning should keep ≤ 2 splits, "
                             "got %d" % len(splits))


# ---------------------------------------------------------------------------
# Layer 3 — Property: random PK tables, random Equal/In predicates,
# correctness vs oracle.
# ---------------------------------------------------------------------------
class BucketPruningPropertyTest(unittest.TestCase):

    SEED = 0xB0CC
    TRIALS = 30

    @classmethod
    def setUpClass(cls):
        cls.tempdir = tempfile.mkdtemp()
        cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
        cls.catalog = CatalogFactory.create({'warehouse': cls.warehouse})
        cls.catalog.create_database('default', False)
        cls.rnd = random.Random(cls.SEED)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tempdir, ignore_errors=True)

    def _make_table(self, idx: int, num_buckets: int):
        pa_schema = pa.schema([
            pa.field('k', pa.int64(), nullable=False),
            ('v', pa.int64()),
        ])
        schema = Schema.from_pyarrow_schema(
            pa_schema,
            primary_keys=['k'],
            options={'bucket': str(num_buckets), 'file.format': 'parquet'},
        )
        name = 'default.bp_{}'.format(idx)
        self.catalog.create_table(name, schema, False)
        return self.catalog.get_table(name)

    def _write(self, table, rows):
        pa_schema = pa.schema([
            pa.field('k', pa.int64(), nullable=False),
            ('v', pa.int64()),
        ])
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        c = wb.new_commit()
        try:
            w.write_arrow(pa.Table.from_pylist(rows, schema=pa_schema))
            c.commit(w.prepare_commit())
        finally:
            w.close()
            c.close()

    @staticmethod
    def _expected_buckets(table, keys) -> set:
        """Independent oracle: writer's bucket placement for the given keys."""
        ext = FixedBucketRowKeyExtractor(table.table_schema)
        pa_schema = pa.schema([
            pa.field('k', pa.int64(), nullable=False),
            ('v', pa.int64()),
        ])
        out = set()
        for k in keys:
            arr = pa.RecordBatch.from_pylist(
                [{'k': k, 'v': 0}], schema=pa_schema)
            out.update(ext._extract_buckets_batch(arr))
        return out

    def test_property_pk_equal_correctness(self):
        for trial in range(self.TRIALS):
            num_buckets = self.rnd.choice([2, 4, 8, 16])
            table = self._make_table(trial, num_buckets)
            keys = self.rnd.sample(range(1000), self.rnd.randint(20, 100))
            rows = [{'k': k, 'v': k * 13} for k in keys]
            self._write(table, rows)

            target = self.rnd.choice(keys)
            pb = table.new_read_builder().new_predicate_builder()
            pred = pb.equal('k', target)
            rb = table.new_read_builder().with_filter(pred)
            splits = rb.new_scan().plan().splits()
            if splits:
                got = rb.new_read().to_arrow(splits).to_pylist()
            else:
                got = []
            self.assertEqual(got, [{'k': target, 'v': target * 13}],
                             "trial {} buckets={} target={}: result mismatch"
                             .format(trial, num_buckets, target))
            # Pruning fired AND picked the writer's bucket. Without this
            # cross-check a fail-open selector (i.e. no pruning) would
            # still pass the result-equality assertion above.
            self.assertEqual(self._split_buckets(splits),
                             self._expected_buckets(table, [target]),
                             "trial {}: bucket set != writer's placement"
                             .format(trial))

    def test_property_pk_in_correctness(self):
        for trial in range(self.TRIALS):
            num_buckets = self.rnd.choice([2, 4, 8, 16])
            offset = self.TRIALS + trial  # avoid name clash with prev test
            table = self._make_table(offset, num_buckets)
            keys = self.rnd.sample(range(1000), self.rnd.randint(20, 100))
            rows = [{'k': k, 'v': k * 13} for k in keys]
            self._write(table, rows)

            target_n = self.rnd.randint(1, min(10, len(keys)))
            targets = self.rnd.sample(keys, target_n)
            pb = table.new_read_builder().new_predicate_builder()
            pred = pb.is_in('k', targets)
            rb = table.new_read_builder().with_filter(pred)
            splits = rb.new_scan().plan().splits()
            if splits:
                got = rb.new_read().to_arrow(splits).to_pylist()
            else:
                got = []
            got_sorted = sorted(got, key=lambda r: r['k'])
            want = sorted(
                [{'k': k, 'v': k * 13} for k in targets],
                key=lambda r: r['k'])
            self.assertEqual(got_sorted, want,
                             "trial {}: IN result mismatch".format(trial))
            self.assertEqual(self._split_buckets(splits),
                             self._expected_buckets(table, targets),
                             "trial {}: IN bucket set != writer's placement"
                             .format(trial))

    @staticmethod
    def _split_buckets(splits) -> set:
        return {s.bucket for s in splits}


if __name__ == '__main__':
    unittest.main()
