# 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.

import os
import shutil
import sys
import tempfile
import time
import unittest

import numpy as np
import pandas as pd
import pyarrow as pa

from pypaimon import CatalogFactory, Schema
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.write.file_store_commit import RetryResult


class AoReaderTest(unittest.TestCase):
    @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', True)

        cls.pa_schema = pa.schema([
            ('user_id', pa.int32()),
            ('item_id', pa.int64()),
            ('behavior', pa.string()),
            ('dt', pa.string())
        ])
        cls.expected = pa.Table.from_pydict({
            'user_id': [1, 2, 3, 4, 5, 6, 7, 8],
            'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008],
            'behavior': ['a', 'b', 'c', None, 'e', 'f', 'g', 'h'],
            'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p2'],
        }, schema=cls.pa_schema)

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

    def test_parquet_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_append_only_parquet', schema, False)
        table = self.catalog.get_table('default.test_append_only_parquet')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_target_file_row_num_fails_fast(self):
        schema = Schema.from_pyarrow_schema(
            self.pa_schema, partition_keys=['dt'],
            options={'target-file-row-num': '5'})
        self.catalog.create_table('default.test_target_file_row_num', schema, False)
        table = self.catalog.get_table('default.test_target_file_row_num')
        with self.assertRaisesRegex(
                NotImplementedError, 'row-count based file rolling'):
            write_builder = table.new_batch_write_builder()
            table_write = write_builder.new_write()
            table_write.write_arrow(self.expected)

    def test_orc_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'orc'})
        self.catalog.create_table('default.test_append_only_orc', schema, False)
        table = self.catalog.get_table('default.test_append_only_orc')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_avro_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'})
        self.catalog.create_table('default.test_append_only_avro', schema, False)
        table = self.catalog.get_table('default.test_append_only_avro')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_lance_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'})
        self.catalog.create_table('default.test_append_only_lance', schema, False)
        table = self.catalog.get_table('default.test_append_only_lance')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_isnull_partition_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.isnull_partition_filter', schema, False)
        table = self.catalog.get_table('default.isnull_partition_filter')
        data = pa.Table.from_pydict({
            'user_id': [1, 2],
            'item_id': [1001, 1002],
            'behavior': ['a', 'b'],
            'dt': [None, 'p1'],
        }, schema=self.pa_schema)
        wb = table.new_batch_write_builder()
        w, c = wb.new_write(), wb.new_commit()
        w.write_arrow(data)
        c.commit(w.prepare_commit())
        w.close()
        c.close()

        pb = table.new_read_builder().new_predicate_builder()
        rb = table.new_read_builder().with_filter(pb.is_null('dt'))
        result = self._read_test_table(rb)

        self.assertEqual(result.column('dt').to_pylist(), [None])
        self.assertEqual(result.column('user_id').to_pylist(), [1])

    def test_plan_snapshot_id_for_empty_and_non_empty_scan(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_plan_snapshot_id', schema, False)
        table = self.catalog.get_table('default.test_plan_snapshot_id')

        empty_plan = table.new_read_builder().new_scan().plan()
        self.assertIsNone(empty_plan.snapshot_id)
        self.assertEqual(len(empty_plan.splits()), 0)

        self._write_test_table(table)

        plan = table.new_read_builder().new_scan().plan()
        self.assertEqual(plan.snapshot_id, 2)
        self.assertGreater(len(plan.splits()), 0)

    def test_incremental_timestamp_empty_range_keeps_end_snapshot_id(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_incremental_empty_range_snapshot', schema, False)
        table = self.catalog.get_table('default.test_incremental_empty_range_snapshot')

        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        pa_table = pa.Table.from_pydict({
            'user_id': [1],
            'item_id': [1001],
            'behavior': ['a'],
            'dt': ['p1'],
        }, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        snapshot_manager = table.snapshot_manager()
        snapshot = snapshot_manager.get_latest_snapshot()
        table_inc = table.copy({
            CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key():
                "{},{}".format(snapshot.time_millis, snapshot.time_millis + 1)
        })

        plan = table_inc.new_read_builder().new_scan().plan()
        self.assertEqual(plan.snapshot_id, snapshot.id)
        self.assertEqual(len(plan.splits()), 0)

    @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11")
    def test_vortex_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'})
        self.catalog.create_table('default.test_append_only_vortex', schema, False)
        table = self.catalog.get_table('default.test_append_only_vortex')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11")
    def test_vortex_ao_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'})
        self.catalog.create_table('default.test_append_only_vortex_filter', schema, False)
        table = self.catalog.get_table('default.test_append_only_vortex_filter')
        self._write_test_table(table)

        predicate_builder = table.new_read_builder().new_predicate_builder()
        p1 = predicate_builder.less_than('user_id', 7)
        p2 = predicate_builder.greater_or_equal('user_id', 2)
        p3 = predicate_builder.between('user_id', 0, 6)  # [2/b, 3/c, 4/d, 5/e, 6/f] left
        p6 = predicate_builder.is_not_null('behavior')  # exclude 4/d -> [2/b, 3/c, 5/e, 6/f]
        g1 = predicate_builder.and_predicates([p1, p2, p3, p6])
        read_builder = table.new_read_builder().with_filter(g1)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(1, 2),  # 2/b, 3/c
            self.expected.slice(4, 2),  # 5/e, 6/f
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

        # OR predicates with startswith, endswith, contains, equal, is_null
        p7 = predicate_builder.startswith('behavior', 'a')
        p10 = predicate_builder.equal('item_id', 1002)
        p11 = predicate_builder.is_null('behavior')
        p9 = predicate_builder.contains('behavior', 'f')
        p8 = predicate_builder.endswith('dt', 'p2')
        g2 = predicate_builder.or_predicates([p7, p8, p9, p10, p11])
        read_builder = table.new_read_builder().with_filter(g2)
        actual = self._read_test_table(read_builder)
        self.assertEqual(actual.sort_by('user_id'), self.expected)

        # Combined AND + OR
        g3 = predicate_builder.and_predicates([g1, g2])
        read_builder = table.new_read_builder().with_filter(g3)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(1, 2),  # 2/b, 3/c
            self.expected.slice(4, 2),  # 5/e, 6/f
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

        # not_equal also filters None values
        p12 = predicate_builder.not_equal('behavior', 'f')
        read_builder = table.new_read_builder().with_filter(p12)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(0, 1),  # 1/a
            self.expected.slice(1, 1),  # 2/b
            self.expected.slice(2, 1),  # 3/c
            self.expected.slice(4, 1),  # 5/e
            self.expected.slice(6, 1),  # 7/g
            self.expected.slice(7, 1),  # 8/h
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

    @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11")
    def test_vortex_ao_reader_with_projection(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'})
        self.catalog.create_table('default.test_vortex_append_only_projection', schema, False)
        table = self.catalog.get_table('default.test_vortex_append_only_projection')
        self._write_test_table(table)

        read_builder = table.new_read_builder().with_projection(['dt', 'user_id'])
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = self.expected.select(['dt', 'user_id'])
        self.assertEqual(actual, expected)

    @unittest.skipIf(sys.version_info < (3, 11), "vortex-data requires Python >= 3.11")
    def test_vortex_ao_reader_with_shard(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'vortex'})
        self.catalog.create_table('default.test_append_only_vortex_shard', schema, False)
        table = self.catalog.get_table('default.test_append_only_vortex_shard')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        shard_tables = []
        total_shards = 3
        for i in range(total_shards):
            splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits()
            shard_tables.append(table_read.to_arrow(splits))

        actual = pa.concat_tables(shard_tables).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_lance_ao_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'})
        self.catalog.create_table('default.test_append_only_lance_filter', schema, False)
        table = self.catalog.get_table('default.test_append_only_lance_filter')
        self._write_test_table(table)

        predicate_builder = table.new_read_builder().new_predicate_builder()
        p1 = predicate_builder.less_than('user_id', 7)
        p2 = predicate_builder.greater_or_equal('user_id', 2)
        p3 = predicate_builder.between('user_id', 0, 6)  # [2/b, 3/c, 4/d, 5/e, 6/f] left
        p4 = predicate_builder.is_not_in('behavior', ['b', 'e'])  # [3/c, 4/d, 6/f] left
        p5 = predicate_builder.is_in('dt', ['p1'])  # exclude 3/c
        p6 = predicate_builder.is_not_null('behavior')  # exclude 4/d
        g1 = predicate_builder.and_predicates([p1, p2, p3, p4, p5, p6])
        read_builder = table.new_read_builder().with_filter(g1)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(5, 1)  # 6/f
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

    def test_lance_ao_reader_with_shard(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'})
        self.catalog.create_table('default.test_append_only_lance_shard', schema, False)
        table = self.catalog.get_table('default.test_append_only_lance_shard')
        self._write_test_table(table)

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()

        shard_tables = []
        total_shards = 3
        for i in range(total_shards):
            splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits()
            shard_tables.append(table_read.to_arrow(splits))

        actual = pa.concat_tables(shard_tables).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_lance_sliced_split_row_range_pushdown(self):
        """
        SlicedSplit with Lance format calls read_range() instead of read_all(),
        reading only the requested row slice from disk rather than the full file.
        """
        import unittest.mock as mock
        try:
            import lance as _lance
            import lance.file  # ensure submodule is loaded before write_lance uses it
        except ImportError:
            self.skipTest("lance not installed")

        schema = Schema.from_pyarrow_schema(
            pa.schema([('id', pa.int64()), ('value', pa.string())]),
            options={'file.format': 'lance'})
        self.catalog.create_table('default.test_lance_sliced_split', schema, False)
        table = self.catalog.get_table('default.test_lance_sliced_split')

        # Write 1000 rows in a single shot so they land in one file
        n = 1000
        pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]})
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        scan = table.new_read_builder().new_scan()
        splits = scan.plan().splits()
        self.assertEqual(len(splits), 1, "Expected single split (no partition, single bucket)")
        data_split = splits[0]
        self.assertEqual(len(data_split.files), 1, "Expected single file in split")

        file_name = data_split.files[0].file_name
        slice_start, slice_end = 200, 700  # request 500 of 1000 rows

        from pypaimon.read.sliced_split import SlicedSplit
        sliced = SlicedSplit(data_split, {file_name: (slice_start, slice_end)})

        # Spy on lance.file.LanceFileReader to verify read_range is used
        OrigReader = _lance.file.LanceFileReader
        read_all_calls = []
        read_range_calls = []

        class SpyLanceFileReader:
            def __init__(self, *args, **kwargs):
                self._r = OrigReader(*args, **kwargs)

            def read_all(self, *args, **kwargs):
                read_all_calls.append(())
                return self._r.read_all(*args, **kwargs)

            def read_range(self, start, count, **kwargs):
                read_range_calls.append((start, count))
                return self._r.read_range(start, count, **kwargs)

            def __getattr__(self, name):
                return getattr(self._r, name)

        table_read = table.new_read_builder().new_read()
        with mock.patch.object(_lance.file, 'LanceFileReader', SpyLanceFileReader):
            result = table_read.to_arrow([sliced])

        # Verify that read_range was used, not read_all
        self.assertEqual(len(read_all_calls), 0,
                         "read_all() must not be called when SlicedSplit row range is available")
        self.assertEqual(len(read_range_calls), 1)
        actual_start, actual_count = read_range_calls[0]
        self.assertEqual(actual_start, slice_start)
        self.assertEqual(actual_count, slice_end - slice_start,
                         f"read_range should request exactly {slice_end - slice_start} rows, "
                         f"not the full {n} rows of the file")

        # Verify functional correctness: correct rows returned
        self.assertEqual(result.num_rows, slice_end - slice_start)
        result_ids = sorted(result.column('id').to_pylist())
        self.assertEqual(result_ids, list(range(slice_start, slice_end)))

    def test_lance_shard_row_id_correctness(self):
        """
        with_shard splits a file into contiguous ranges via SlicedSplit.
        _ROW_ID across all shards must equal the full table's _ROW_ID.
        """
        try:
            import lance  # noqa: F401
        except ImportError:
            self.skipTest("lance not installed")

        schema = Schema.from_pyarrow_schema(
            pa.schema([('id', pa.int64()), ('value', pa.string())]),
            options={'file.format': 'lance',
                     'row-tracking.enabled': 'true',
                     'data-evolution.enabled': 'true'})
        self.catalog.create_table('default.test_lance_shard_row_id', schema, False)
        table = self.catalog.get_table('default.test_lance_shard_row_id')

        n = 1000
        pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]})
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Read full table _ROW_ID as ground truth
        read_builder = table.new_read_builder().with_projection(['id', '_ROW_ID'])
        full_splits = read_builder.new_scan().plan().splits()
        table_read = read_builder.new_read()
        full_result = table_read.to_arrow(full_splits)
        expected_row_ids = sorted(full_result.column('_ROW_ID').to_pylist())

        # Read via with_shard and collect _ROW_ID from all shards
        total_shards = 3
        all_row_ids = []
        for i in range(total_shards):
            shard_splits = read_builder.new_scan().with_shard(i, total_shards).plan().splits()
            shard_result = table_read.to_arrow(shard_splits)
            all_row_ids.extend(shard_result.column('_ROW_ID').to_pylist())

        self.assertEqual(sorted(all_row_ids), expected_row_ids)

    def test_lance_indexed_split_take_rows_pushdown(self):
        """
        IndexedSplit row ranges (from ANN global index results) are converted to
        local file indices and pushed down to lance.file.LanceFileReader.take_rows(),
        so only the matched rows are physically read instead of the full file.
        """
        import unittest.mock as mock
        try:
            import lance as _lance
            import lance.file
        except ImportError:
            self.skipTest("lance not installed")

        schema = Schema.from_pyarrow_schema(
            pa.schema([('id', pa.int64()), ('value', pa.string())]),
            options={'file.format': 'lance',
                     'row-tracking.enabled': 'true',
                     'data-evolution.enabled': 'true'})
        self.catalog.create_table('default.test_lance_indexed_split', schema, False)
        table = self.catalog.get_table('default.test_lance_indexed_split')

        n = 1000
        pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]})
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        splits = table.new_read_builder().new_scan().plan().splits()
        self.assertEqual(len(splits), 1)
        data_split = splits[0]
        self.assertEqual(len(data_split.files), 1)

        file_meta = data_split.files[0]
        self.assertIsNotNone(file_meta.first_row_id, "row tracking must assign first_row_id")
        first_row_id = file_meta.first_row_id

        # Simulate ANN result: 3 scattered global row IDs
        from pypaimon.globalindex.indexed_split import IndexedSplit
        from pypaimon.utils.range import Range
        target_global_ids = [first_row_id + 50, first_row_id + 300, first_row_id + 700]
        row_ranges = [Range(g, g) for g in target_global_ids]
        indexed = IndexedSplit(data_split, row_ranges)

        OrigReader = _lance.file.LanceFileReader
        take_rows_calls = []
        read_all_calls = []

        class SpyLanceFileReader:
            def __init__(self, *args, **kwargs):
                self._r = OrigReader(*args, **kwargs)

            def take_rows(self, indices, **kwargs):
                take_rows_calls.append(list(indices))
                return self._r.take_rows(indices, **kwargs)

            def read_all(self, *args, **kwargs):
                read_all_calls.append(())
                return self._r.read_all(*args, **kwargs)

            def read_range(self, *args, **kwargs):
                return self._r.read_range(*args, **kwargs)

            def __getattr__(self, name):
                return getattr(self._r, name)

        table_read = table.new_read_builder().new_read()
        with mock.patch.object(_lance.file, 'LanceFileReader', SpyLanceFileReader):
            result = table_read.to_arrow([indexed])

        self.assertEqual(len(read_all_calls), 0,
                         "read_all() must not be called when IndexedSplit row ranges are available")
        self.assertEqual(len(take_rows_calls), 1)
        # local indices = global_ids - first_row_id
        expected_local = [g - first_row_id for g in target_global_ids]
        self.assertEqual(sorted(take_rows_calls[0]), sorted(expected_local),
                         f"take_rows should request local indices {expected_local}, "
                         f"not read the full {n} rows")

        self.assertEqual(result.num_rows, len(target_global_ids))
        result_ids = sorted(result.column('id').to_pylist())
        self.assertEqual(result_ids, [50, 300, 700])

    def test_lance_indexed_split_row_id_correctness(self):
        """
        IndexedSplit (ANN vector search) with native take_rows pushdown must
        return correct _ROW_ID values matching the original global row IDs.
        """
        try:
            import lance  # noqa: F401
        except ImportError:
            self.skipTest("lance not installed")

        schema = Schema.from_pyarrow_schema(
            pa.schema([('id', pa.int64()), ('value', pa.string())]),
            options={'file.format': 'lance',
                     'row-tracking.enabled': 'true',
                     'data-evolution.enabled': 'true'})
        self.catalog.create_table('default.test_lance_indexed_row_id', schema, False)
        table = self.catalog.get_table('default.test_lance_indexed_row_id')

        n = 1000
        pa_table = pa.table({'id': list(range(n)), 'value': [f'v{i}' for i in range(n)]})
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # Read full table to get ground-truth _ROW_ID for specific rows
        read_builder = table.new_read_builder().with_projection(['id', '_ROW_ID'])
        full_splits = read_builder.new_scan().plan().splits()
        table_read = read_builder.new_read()
        full_result = table_read.to_arrow(full_splits)
        # Build id -> _ROW_ID mapping from full scan
        id_to_row_id = dict(zip(
            full_result.column('id').to_pylist(),
            full_result.column('_ROW_ID').to_pylist()))

        # Construct IndexedSplit targeting specific rows
        data_split = full_splits[0]
        file_meta = data_split.files[0]
        first_row_id = file_meta.first_row_id

        from pypaimon.globalindex.indexed_split import IndexedSplit
        from pypaimon.utils.range import Range
        target_local_offsets = [50, 300, 700]
        target_global_ids = [first_row_id + o for o in target_local_offsets]
        row_ranges = [Range(g, g) for g in target_global_ids]
        indexed = IndexedSplit(data_split, row_ranges)

        result = table_read.to_arrow([indexed])

        self.assertEqual(result.num_rows, len(target_global_ids))
        for i in range(result.num_rows):
            row_id = result.column('_ROW_ID')[i].as_py()
            data_id = result.column('id')[i].as_py()
            self.assertEqual(row_id, id_to_row_id[data_id],
                             f"row id={data_id}: _ROW_ID should be {id_to_row_id[data_id]}, got {row_id}")

    def test_append_only_multi_write_once_commit(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_append_only_multi_once_commit', schema, False)
        table = self.catalog.get_table('default.test_append_only_multi_once_commit')
        write_builder = table.new_batch_write_builder()

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table1 = pa.Table.from_pydict(data1, schema=self.pa_schema)
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
        }
        pa_table2 = pa.Table.from_pydict(data2, schema=self.pa_schema)

        table_write.write_arrow(pa_table1)
        table_write.write_arrow(pa_table2)

        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_commit_retry_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_commit_retry_filter', schema, False)
        table = self.catalog.get_table('default.test_commit_retry_filter')
        write_builder = table.new_batch_write_builder()

        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table1 = pa.Table.from_pydict(data1, schema=self.pa_schema)
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
        }
        pa_table2 = pa.Table.from_pydict(data2, schema=self.pa_schema)

        table_write.write_arrow(pa_table1)
        table_write.write_arrow(pa_table2)

        messages = table_write.prepare_commit()
        table_commit.commit(messages)
        table_write.close()

        snapshot_manager = table.snapshot_manager()
        latest_snapshot = snapshot_manager.get_latest_snapshot()
        commit_entries = []
        for msg in messages:
            partition = GenericRow(list(msg.partition), table.partition_keys_fields)
            for file in msg.new_files:
                commit_entries.append(ManifestEntry(
                    kind=0,
                    partition=partition,
                    bucket=msg.bucket,
                    total_buckets=table.total_buckets,
                    file=file
                ))
        # mock retry
        success = table_commit.file_store_commit._try_commit_once(
            RetryResult(None),
            "APPEND",
            commit_entries,
            [],
            BATCH_COMMIT_IDENTIFIER,
            latest_snapshot)
        self.assertTrue(success.is_success())
        table_commit.close()
        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.expected)

    def test_over_1000_cols_read(self):
        num_rows = 1
        num_cols = 10
        table_name = "default.testBug"
        # Generate dynamic schema based on column count
        schema_fields = []
        for i in range(1, num_cols + 1):
            col_name = f'c{i:03d}'
            if i == 1:
                schema_fields.append((col_name, pa.string()))  # ID column
            elif i == 2:
                schema_fields.append((col_name, pa.string()))  # Name column
            elif i == 3:
                schema_fields.append((col_name, pa.string()))  # Category column (partition key)
            elif i % 4 == 0:
                schema_fields.append((col_name, pa.float64()))  # Float columns
            elif i % 4 == 1:
                schema_fields.append((col_name, pa.int32()))  # Int columns
            elif i % 4 == 2:
                schema_fields.append((col_name, pa.string()))  # String columns
            else:
                schema_fields.append((col_name, pa.int64()))  # Long columns

        pa_schema = pa.schema(schema_fields)
        schema = Schema.from_pyarrow_schema(
            pa_schema,
            partition_keys=['c003'],  # Use c003 as partition key
        )

        # Create table
        self.catalog.create_table(table_name, schema, False)
        table = self.catalog.get_table(table_name)

        # Generate test data
        np.random.seed(42)  # For reproducible results
        categories = ['Electronics', 'Clothing', 'Books', 'Home', 'Sports', 'Food', 'Toys', 'Beauty', 'Health', 'Auto']
        statuses = ['Active', 'Inactive', 'Pending', 'Completed']

        # Generate data dictionary
        test_data = {}
        for i in range(1, num_cols + 1):
            col_name = f'c{i:03d}'
            if i == 1:
                test_data[col_name] = [f'Product_{j}' for j in range(1, num_rows + 1)]
            elif i == 2:
                test_data[col_name] = [f'Product_{j}' for j in range(1, num_rows + 1)]
            elif i == 3:
                test_data[col_name] = np.random.choice(categories, num_rows)
            elif i % 4 == 0:
                test_data[col_name] = np.random.uniform(1.0, 1000.0, num_rows).round(2)
            elif i % 4 == 1:
                test_data[col_name] = np.random.randint(1, 100, num_rows)
            elif i % 4 == 2:
                test_data[col_name] = np.random.choice(statuses, num_rows)
            else:
                test_data[col_name] = np.random.randint(1640995200, 1672531200, num_rows)

        test_df = pd.DataFrame(test_data)

        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()

        table_write.write_pandas(test_df)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        result = table_read.to_pandas(table_scan.plan().splits())
        self.assertEqual(result.to_dict(), test_df.to_dict())

    def test_ao_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_append_only_filter', schema, False)
        table = self.catalog.get_table('default.test_append_only_filter')
        self._write_test_table(table)

        predicate_builder = table.new_read_builder().new_predicate_builder()
        p1 = predicate_builder.less_than('user_id', 7)
        p2 = predicate_builder.greater_or_equal('user_id', 2)
        p3 = predicate_builder.between('user_id', 0, 6)  # [2/b, 3/c, 4/d, 5/e, 6/f] left
        p4 = predicate_builder.is_not_in('behavior', ['b', 'e'])  # [3/c, 4/d, 6/f] left
        p5 = predicate_builder.is_in('dt', ['p1'])  # exclude 3/c
        p6 = predicate_builder.is_not_null('behavior')  # exclude 4/d
        g1 = predicate_builder.and_predicates([p1, p2, p3, p4, p5, p6])
        read_builder = table.new_read_builder().with_filter(g1)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(5, 1)  # 6/f
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

        p7 = predicate_builder.startswith('behavior', 'a')
        p10 = predicate_builder.equal('item_id', 1002)
        p11 = predicate_builder.is_null('behavior')
        p9 = predicate_builder.contains('behavior', 'f')
        p8 = predicate_builder.endswith('dt', 'p2')
        g2 = predicate_builder.or_predicates([p7, p8, p9, p10, p11])
        read_builder = table.new_read_builder().with_filter(g2)
        actual = self._read_test_table(read_builder)
        self.assertEqual(actual.sort_by('user_id'), self.expected)

        g3 = predicate_builder.and_predicates([g1, g2])
        read_builder = table.new_read_builder().with_filter(g3)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            self.expected.slice(5, 1)  # 6/f
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

        # Same as java, 'not_equal' will also filter records of 'None' value
        p12 = predicate_builder.not_equal('behavior', 'f')
        read_builder = table.new_read_builder().with_filter(p12)
        actual = self._read_test_table(read_builder)
        expected = pa.concat_tables([
            # not only 6/f, but also 4/d will be filtered
            self.expected.slice(0, 1),  # 1/a
            self.expected.slice(1, 1),  # 2/b
            self.expected.slice(2, 1),  # 3/c
            self.expected.slice(4, 1),  # 5/e
            self.expected.slice(6, 1),  # 7/g
            self.expected.slice(7, 1),  # 8/h
        ])
        self.assertEqual(actual.sort_by('user_id'), expected)

    def test_ao_reader_with_projection(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_append_only_projection', schema, False)
        table = self.catalog.get_table('default.test_append_only_projection')
        self._write_test_table(table)

        read_builder = table.new_read_builder().with_projection(['dt', 'user_id'])
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = self.expected.select(['dt', 'user_id'])
        self.assertEqual(actual, expected)

    def test_avro_ao_reader_with_projection(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'})
        self.catalog.create_table('default.test_avro_append_only_projection', schema, False)
        table = self.catalog.get_table('default.test_avro_append_only_projection')
        self._write_test_table(table)

        read_builder = table.new_read_builder().with_projection(['dt', 'user_id'])
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = self.expected.select(['dt', 'user_id'])
        self.assertEqual(actual, expected)

    def test_ao_reader_with_limit(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_append_only_limit', schema, False)
        table = self.catalog.get_table('default.test_append_only_limit')
        self._write_test_table(table)

        # Row-level limit: the reader stops at exactly N rows (not "first
        # split's full row count"). Scan still keeps the first split that
        # covers the limit; the reader short-circuits inside it.
        read_builder = table.new_read_builder().with_limit(1)
        actual = self._read_test_table(read_builder)
        self.assertEqual(actual.num_rows, 1)

    def test_incremental_timestamp(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_incremental_parquet', schema, False)
        table = self.catalog.get_table('default.test_incremental_parquet')
        timestamp = int(time.time() * 1000)
        self._write_test_table(table)

        snapshot_manager = table.snapshot_manager()
        t1 = snapshot_manager.get_snapshot_by_id(1).time_millis
        t2 = snapshot_manager.get_snapshot_by_id(2).time_millis
        # test 1
        table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp - 1) + ',' + str(timestamp)})
        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder)
        self.assertEqual(len(actual), 0)
        # test 2
        table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(timestamp) + ',' + str(t2)})
        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(self.expected, actual)
        # test 3
        table = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): str(t1) + ',' + str(t2)})
        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = self.expected.slice(4, 4)
        self.assertEqual(expected, actual)

    def test_incremental_read_multi_snapshots(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.catalog.create_table('default.test_incremental_100', schema, False)
        table = self.catalog.get_table('default.test_incremental_100')

        write_builder = table.new_batch_write_builder()
        for i in range(1, 101):
            table_write = write_builder.new_write()
            table_commit = write_builder.new_commit()
            pa_table = pa.Table.from_pydict({
                'user_id': [i],
                'item_id': [1000 + i],
                'behavior': [f'snap{i}'],
                'dt': ['p1' if i % 2 == 1 else 'p2'],
            }, schema=self.pa_schema)
            table_write.write_arrow(pa_table)
            table_commit.commit(table_write.prepare_commit())
            table_write.close()
            table_commit.close()

        snapshot_manager = table.snapshot_manager()
        t10 = snapshot_manager.get_snapshot_by_id(10).time_millis
        t20 = snapshot_manager.get_snapshot_by_id(20).time_millis

        table_inc = table.copy({CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP.key(): f"{t10},{t20}"})
        read_builder = table_inc.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')

        expected = pa.Table.from_pydict({
            'user_id': list(range(11, 21)),
            'item_id': [1000 + i for i in range(11, 21)],
            'behavior': [f'snap{i}' for i in range(11, 21)],
            'dt': ['p1' if i % 2 == 1 else 'p2' for i in range(11, 21)],
        }, schema=self.pa_schema).sort_by('user_id')
        self.assertEqual(expected, actual)

    def test_concurrent_writes_with_retry(self):
        """Test concurrent writes to verify retry mechanism works correctly."""
        import threading

        # Run the test 10 times to verify stability
        iter_num = 5
        for test_iteration in range(iter_num):
            # Create a unique table for each iteration
            table_name = f'default.test_concurrent_writes_{test_iteration}'
            # Concurrent commits are expected here; enlarge the retry budget so the
            # default (commit.max-retries=10, commit.max-retry-wait=1s) does not
            # exhaust under heavy CI load and produce a flaky failure.
            schema = Schema.from_pyarrow_schema(
                self.pa_schema,
                options={
                    'commit.max-retries': '50',
                    'commit.max-retry-wait': '30s',
                },
            )
            self.catalog.create_table(table_name, schema, False)
            table = self.catalog.get_table(table_name)

            write_results = []
            write_errors = []

            def write_data(thread_id, start_user_id):
                """Write data in a separate thread."""
                try:
                    threading.current_thread().name = f"Iter{test_iteration}-Thread-{thread_id}"
                    write_builder = table.new_batch_write_builder()
                    table_write = write_builder.new_write()
                    table_commit = write_builder.new_commit()

                    # Create unique data for this thread
                    data = {
                        'user_id': list(range(start_user_id, start_user_id + 5)),
                        'item_id': [1000 + i for i in range(start_user_id, start_user_id + 5)],
                        'behavior': [f'thread{thread_id}_{i}' for i in range(5)],
                        'dt': ['p1' if i % 2 == 0 else 'p2' for i in range(5)],
                    }
                    pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)

                    table_write.write_arrow(pa_table)
                    commit_messages = table_write.prepare_commit()

                    table_commit.commit(commit_messages)
                    table_write.close()
                    table_commit.close()

                    write_results.append({
                        'thread_id': thread_id,
                        'start_user_id': start_user_id,
                        'success': True
                    })
                except Exception as e:
                    write_errors.append({
                        'thread_id': thread_id,
                        'error': str(e)
                    })

            # Create and start multiple threads. Keep this modest (3 vs. the
            # original 10) because GHA runners under load can't drain 10
            # simultaneously-conflicting commits even with
            # ``commit.max-retries=50`` (50 attempts * 30s back-off ~25 min,
            # still timing out in CI). Three threads exercises the retry path
            # without pushing each iteration past the per-test wall-time
            # budget.
            threads = []
            num_threads = 3
            for i in range(num_threads):
                thread = threading.Thread(
                    target=write_data,
                    args=(i, i * 10)
                )
                threads.append(thread)
                thread.start()

            # Wait for all threads to complete
            for thread in threads:
                thread.join()

            # Verify all writes succeeded (retry mechanism should handle conflicts)
            self.assertEqual(num_threads, len(write_results),
                             f"Iteration {test_iteration}: Expected {num_threads} successful writes, "
                             f"got {len(write_results)}. Errors: {write_errors}")
            self.assertEqual(0, len(write_errors),
                             f"Iteration {test_iteration}: Expected no errors, but got: {write_errors}")

            read_builder = table.new_read_builder()
            actual = self._read_test_table(read_builder).sort_by('user_id')

            # Verify data rows
            self.assertEqual(num_threads * 5, actual.num_rows,
                             f"Iteration {test_iteration}: Expected {num_threads * 5} rows")

            # Verify user_id
            user_ids = actual.column('user_id').to_pylist()
            expected_user_ids = []
            for i in range(num_threads):
                expected_user_ids.extend(range(i * 10, i * 10 + 5))
            expected_user_ids.sort()

            self.assertEqual(user_ids, expected_user_ids,
                             f"Iteration {test_iteration}: User IDs mismatch")

            # Verify snapshot count (should have num_threads snapshots)
            snapshot_manager = table.snapshot_manager()
            latest_snapshot = snapshot_manager.get_latest_snapshot()
            self.assertIsNotNone(latest_snapshot,
                                 f"Iteration {test_iteration}: Latest snapshot should not be None")
            self.assertEqual(latest_snapshot.id, num_threads,
                             f"Iteration {test_iteration}: Expected snapshot ID {num_threads}, "
                             f"got {latest_snapshot.id}")

            print(f"✓ Iteration {test_iteration + 1}/{iter_num} completed successfully")

    def test_is_in_with_partitions(self):
        from pypaimon.manifest.manifest_file_manager import ManifestFileManager
        from io import BytesIO
        import fastavro

        num_partitions = 5000
        in_size = 2000

        schema = Schema.from_pyarrow_schema(
            pa.schema([('pk', pa.int32()), ('val', pa.string()),
                       ('pt', pa.string())]),
            partition_keys=['pt'])
        self.catalog.create_table(
            'default.test_is_in_partition_bench', schema, False)
        table = self.catalog.get_table(
            'default.test_is_in_partition_bench')

        wb = table.new_batch_write_builder()
        w = wb.new_write()
        bench_schema = pa.schema(
            [('pk', pa.int32()), ('val', pa.string()),
             ('pt', pa.string())])
        for i in range(num_partitions):
            w.write_arrow(pa.Table.from_pydict(
                {'pk': [i], 'val': [f'v_{i}'], 'pt': [f'pt_{i}']},
                schema=bench_schema))
        wb.new_commit().commit(w.prepare_commit())
        w.close()

        builder = table.new_read_builder().new_predicate_builder()
        in_values = [f'pt_{i}' for i in range(in_size)]
        pred = builder.is_in('pt', in_values)

        splits = table.new_read_builder().with_filter(
            pred).new_scan().plan().splits()
        self.assertEqual(len(splits), in_size)

        from pypaimon.manifest.schema.data_file_meta import DataFileMeta

        entry_counts = {'avro_total': 0, 'constructed': 0}
        original_read = ManifestFileManager.read
        original_dfm_init = DataFileMeta.__init__

        def counting_read(self_mgr, manifest_file_name,
                          manifest_entry_filter=None,
                          drop_stats=True, early_entry_filter=None,
                          early_record_filter=None, partition_filter=None):
            # avro_total = every entry in the manifest (no manifest-file pruning
            # here: single file, is_in spans its partition stats).
            path = f"{self_mgr.manifest_path}/{manifest_file_name}"
            with self_mgr.file_io.new_input_stream(path) as s:
                for _ in fastavro.reader(BytesIO(s.read())):
                    entry_counts['avro_total'] += 1
            return original_read(
                self_mgr, manifest_file_name,
                manifest_entry_filter, drop_stats,
                early_entry_filter, early_record_filter, partition_filter)

        def counting_dfm_init(self_dfm, *args, **kwargs):
            entry_counts['constructed'] += 1
            original_dfm_init(self_dfm, *args, **kwargs)

        ManifestFileManager.read = counting_read
        DataFileMeta.__init__ = counting_dfm_init
        try:
            table.new_read_builder().with_filter(
                pred).new_scan().plan()
        finally:
            ManifestFileManager.read = original_read
            DataFileMeta.__init__ = original_dfm_init

        # all entries scanned, but only matching ones build a DataFileMeta
        # (_FILE + drop_stats copy = 2 each); non-matching skipped pre-construct.
        self.assertEqual(entry_counts['avro_total'], num_partitions)
        self.assertEqual(entry_counts['constructed'], in_size * 2)

    def test_partition_filter_test_handles_isnull(self):
        # The early filter keeps/drops an entry via Predicate.test on its exact
        # partition value: isNull keeps null, drops non-null; equal is the reverse.
        from pypaimon.table.row.generic_row import GenericRow
        from pypaimon.table.row.row_kind import RowKind

        schema = Schema.from_pyarrow_schema(
            pa.schema([('pk', pa.int32()), ('pt', pa.string())]),
            partition_keys=['pt'])
        self.catalog.create_table('default.test_isnull_pt_pred', schema, False)
        table = self.catalog.get_table('default.test_isnull_pt_pred')
        fields = table.table_schema.fields  # predicate index is table-schema space
        row_null = GenericRow([1, None], fields, RowKind.INSERT)
        row_x = GenericRow([2, 'x'], fields, RowKind.INSERT)

        pb = table.new_read_builder().new_predicate_builder()
        self.assertTrue(pb.is_null('pt').test(row_null))
        self.assertFalse(pb.is_null('pt').test(row_x))
        self.assertFalse(pb.equal('pt', 'x').test(row_null))
        self.assertTrue(pb.equal('pt', 'x').test(row_x))

    def test_early_partition_filter_isnull_scan(self):
        # Black-box: drive the early manifest partition filter through an actual
        # scan (not just Predicate.test), over a null + non-null partition mix.
        # isNull must keep only the null partition, and the non-matching entries
        # must be skipped before their _FILE block is deserialized.
        from pypaimon.manifest.schema.data_file_meta import DataFileMeta

        schema = Schema.from_pyarrow_schema(
            pa.schema([('pk', pa.int32()), ('pt', pa.string())]),
            partition_keys=['pt'])
        self.catalog.create_table('default.early_isnull_pt', schema, False)
        table = self.catalog.get_table('default.early_isnull_pt')
        ws = pa.schema([('pk', pa.int32()), ('pt', pa.string())])
        wb = table.new_batch_write_builder()
        w = wb.new_write()
        w.write_arrow(pa.Table.from_pydict({'pk': [0], 'pt': [None]}, schema=ws))
        for i in range(1, 6):
            w.write_arrow(pa.Table.from_pydict(
                {'pk': [i], 'pt': [f'p{i}']}, schema=ws))
        wb.new_commit().commit(w.prepare_commit())
        w.close()

        pb = table.new_read_builder().new_predicate_builder()

        # isNull keeps only the null partition; count DataFileMeta constructions to
        # prove the 5 non-null entries were skipped by the early filter pre-_FILE
        # (each kept entry builds 2: _FILE + drop_stats copy).
        constructed = {'n': 0}
        orig_init = DataFileMeta.__init__

        def counting_init(self_dfm, *args, **kwargs):
            constructed['n'] += 1
            orig_init(self_dfm, *args, **kwargs)

        DataFileMeta.__init__ = counting_init
        try:
            null_splits = table.new_read_builder().with_filter(
                pb.is_null('pt')).new_scan().plan().splits()
        finally:
            DataFileMeta.__init__ = orig_init
        self.assertEqual([s.partition.values for s in null_splits], [[None]])
        self.assertEqual(constructed['n'], 2)

        # complement: equal keeps only the matching non-null partition
        p3_splits = table.new_read_builder().with_filter(
            pb.equal('pt', 'p3')).new_scan().plan().splits()
        self.assertEqual([s.partition.values for s in p3_splits], [['p3']])

    def _write_test_table(self, table):
        write_builder = table.new_batch_write_builder()

        # first write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data1 = {
            'user_id': [1, 2, 3, 4],
            'item_id': [1001, 1002, 1003, 1004],
            'behavior': ['a', 'b', 'c', None],
            'dt': ['p1', 'p1', 'p2', 'p1'],
        }
        pa_table = pa.Table.from_pydict(data1, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        # second write
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        data2 = {
            'user_id': [5, 6, 7, 8],
            'item_id': [1005, 1006, 1007, 1008],
            'behavior': ['e', 'f', 'g', 'h'],
            'dt': ['p2', 'p1', 'p2', 'p2'],
        }
        pa_table = pa.Table.from_pydict(data2, schema=self.pa_schema)
        table_write.write_arrow(pa_table)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def _read_test_table(self, read_builder):
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        return table_read.to_arrow(splits)
