# 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 unittest
from datetime import date, datetime
from decimal import Decimal
import os
import struct
import sys
import types

import pyarrow as pa

from pypaimon.globalindex.build_plan import (
    filter_non_indexable_splits as _filter_non_indexable_splits,
    split_by_global_index_shard as _split_by_global_index_shard,
    split_one_by_contiguous_row_range as _split_one_by_contiguous_row_range,
)
from pypaimon.globalindex.create_global_index import GlobalIndexBuilder
from pypaimon.globalindex.key_serializer import create_serializer
from pypaimon.globalindex.full_text.native_full_text_global_index_reader import (
    FULL_TEXT_IDENTIFIER,
    NativeFullTextIndexOptions,
)
from pypaimon.globalindex.full_text.native_full_text_index_writer import (
    NativeFullTextIndexWriter,
)
from pypaimon.globalindex.vindex.vindex_vector_index_writer import (
    VindexVectorIndexWriter,
    _sample_training_vectors,
    native_options,
    train_sample_ratio,
)
from pypaimon.globalindex.data_evolution_global_index_scanner import DataEvolutionGlobalIndexScanner
from pypaimon.index.index_file_handler import IndexFileHandler
from pypaimon.schema.data_types import ArrayType, AtomicType, RowType
from pypaimon.tests.data_evolution_test_helpers import (
    BatchModeMixin,
    DataEvolutionTestBase,
)
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.utils.range import Range


class _FakeFile:

    def __init__(self, file_name, first_row_id, row_count, schema_id=0):
        self.file_name = file_name
        self.first_row_id = first_row_id
        self.row_count = row_count
        self.schema_id = schema_id

    def row_id_range(self):
        if self.first_row_id is None:
            return None
        return Range(self.first_row_id,
                     self.first_row_id + self.row_count - 1)


class _FakeSplit:

    def __init__(self, files):
        self.files = files
        self.partition = GenericRow([], [])
        self.bucket = 0
        self.raw_convertible = False


class _FakeVectorIndexTraining:

    def __init__(self, options, data):
        self.options = dict(options)
        self.trained = data.tolist()
        self.closed = False

    def close(self):
        self.closed = True


class _FakeVectorIndexTrainer:

    @classmethod
    def train(cls, options, data):
        return _FakeVectorIndexTraining(options, data)


class _FakeVectorIndexWriter:
    instances = []

    def __init__(self, training):
        self.options = dict(training.options)
        self.trained = training.trained
        self.added_ids = None
        self.added_vectors = None
        self.closed = False
        _FakeVectorIndexWriter.instances.append(self)

    def add_vectors(self, ids, data):
        self.added_ids = ids.tolist()
        self.added_vectors = data.tolist()

    def write(self, file):
        file.write(b"fake-vindex")

    def close(self):
        self.closed = True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False


class _FakeFullTextForBuild(types.SimpleNamespace):

    def __init__(self):
        super().__init__()
        self.writers = []
        parent = self

        class FullTextIndexWriter:

            def __init__(self_inner, options=None):
                self_inner.options = dict(options or {})
                self_inner.documents = []
                self_inner.written = False
                self_inner.closed = False
                parent.writers.append(self_inner)

            def add_document(self_inner, row_id, text):
                self_inner.documents.append({
                    "row_id": int(row_id),
                    "text": str(text),
                })

            def write(self_inner, output):
                self_inner.written = True
                output.write(b"fake-ftindex")

            def close(self_inner):
                self_inner.closed = True

        self.FullTextIndexWriter = FullTextIndexWriter


def _archive_file_names(file_io, file_path):
    stream = file_io.new_input_stream(file_path)
    try:
        file_count = struct.unpack(">i", stream.read(4))[0]
        names = []
        for _ in range(file_count):
            name_len = struct.unpack(">i", stream.read(4))[0]
            names.append(stream.read(name_len).decode("utf-8"))
            data_len = struct.unpack(">q", stream.read(8))[0]
            stream.seek(stream.tell() + data_len)
        return names
    finally:
        stream.close()


class _FakeSchemaManager:

    def __init__(self, fields_by_schema_id):
        self.fields_by_schema_id = fields_by_schema_id

    def get_schema(self, schema_id):
        return types.SimpleNamespace(
            fields=[
                types.SimpleNamespace(name=name)
                for name in self.fields_by_schema_id[schema_id]
            ]
        )


class GlobalIndexBuildTest(
        BatchModeMixin, DataEvolutionTestBase, unittest.TestCase):

    table_options = {
        'row-tracking.enabled': 'true',
        'data-evolution.enabled': 'true',
        'global-index.enabled': 'true',
        'bucket': '-1',
        'file.format': 'parquet',
    }

    def test_create_btree_global_index_from_python(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [3, 1, 2, 2],
                'name': ['c', 'a', 'b1', 'b2'],
                'age': [30, 10, 20, 21],
                'city': ['z', 'x', 'y', 'y2'],
            },
            schema=self.pa_schema,
        ))

        added = table.create_global_index(
            'id',
            options={'sorted-index.records-per-range': '2'},
        )

        self.assertEqual(2, added)
        snapshot = table.snapshot_manager().get_latest_snapshot()
        self.assertIsNotNone(snapshot.index_manifest)

        entries = IndexFileHandler(table).scan(snapshot)
        self.assertEqual(2, len(entries))
        self.assertEqual({'btree'}, {e.index_file.index_type for e in entries})
        self.assertEqual({0}, {e.index_file.global_index_meta.row_range_start for e in entries})
        self.assertEqual({3}, {e.index_file.global_index_meta.row_range_end for e in entries})

        read_builder = table.new_read_builder()
        predicate = read_builder.new_predicate_builder().equal('id', 2)
        with DataEvolutionGlobalIndexScanner.create(
                table,
                predicate=predicate,
                snapshot=snapshot) as scanner:
            result = scanner.scan(predicate)

        self.assertEqual(
            [Range(2, 3)],
            result.results().to_range_list(),
        )

        table_name = table.identifier.get_full_name()
        table_indexes = self.catalog.get_table(table_name + '$table_indexes')
        index_read_builder = table_indexes.new_read_builder().with_projection([
            'index_type',
            'index_field_name',
            'row_range_start',
            'row_range_end',
        ])
        index_table = index_read_builder.new_read().to_arrow(
            index_read_builder.new_scan().plan().splits())
        self.assertEqual(2, index_table.num_rows)
        self.assertEqual(['btree', 'btree'],
                         index_table.column('index_type').to_pylist())
        self.assertEqual(['id', 'id'],
                         index_table.column('index_field_name').to_pylist())
        self.assertEqual([0, 0],
                         index_table.column('row_range_start').to_pylist())
        self.assertEqual([3, 3],
                         index_table.column('row_range_end').to_pylist())

        key_ranges = self.catalog.get_table(table_name + '$file_key_ranges')
        range_read_builder = key_ranges.new_read_builder().with_projection([
            'file_path',
            'record_count',
            'first_row_id',
        ])
        range_table = range_read_builder.new_read().to_arrow(
            range_read_builder.new_scan().plan().splits())
        self.assertEqual(1, range_table.num_rows)
        file_path = range_table.column('file_path').to_pylist()[0]
        self.assertIn('/bucket-0/', file_path)
        self.assertEqual([4], range_table.column('record_count').to_pylist())
        self.assertEqual([0], range_table.column('first_row_id').to_pylist())

        dv_ranges_type = table_indexes.row_type().fields[6].type
        self.assertIsInstance(dv_ranges_type, ArrayType)
        self.assertIsInstance(dv_ranges_type.element, RowType)
        self.assertEqual(
            ['f0', 'f1', 'f2', '_CARDINALITY'],
            [field.name for field in dv_ranges_type.element.fields],
        )

        self.assertEqual(2, table.drop_global_index('id', dry_run=True))
        self.assertEqual(2, len(IndexFileHandler(table).scan(
            table.snapshot_manager().get_latest_snapshot())))

        self.assertEqual(2, table.drop_global_index('id'))
        latest_snapshot = table.snapshot_manager().get_latest_snapshot()
        self.assertEqual([], IndexFileHandler(table).scan(latest_snapshot))

        index_read_builder = table_indexes.new_read_builder()
        index_table = index_read_builder.new_read().to_arrow(
            index_read_builder.new_scan().plan().splits())
        self.assertEqual(0, index_table.num_rows)

    def test_stale_global_index_commit_conflicts_when_data_files_removed(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [1, 2, 3, 4],
                'name': ['a', 'b', 'c', 'd'],
                'age': [10, 20, 30, 40],
                'city': ['x', 'y', 'z', 'w'],
            },
            schema=self.pa_schema,
        ))

        messages = GlobalIndexBuilder(
            table,
            'id',
            options={'sorted-index.records-per-range': '2'},
        ).build()
        self.assertGreater(sum(len(message.index_adds) for message in messages), 0)

        overwrite_wb = table.new_batch_write_builder().overwrite({})
        overwrite_write = overwrite_wb.new_write()
        overwrite_commit = overwrite_wb.new_commit()
        overwrite_write.write_arrow(pa.table(
            {
                'id': [5, 6],
                'name': ['e', 'f'],
                'age': [50, 60],
                'city': ['new1', 'new2'],
            },
            schema=self.pa_schema,
        ))
        overwrite_commit.commit(overwrite_write.prepare_commit())
        overwrite_write.close()
        overwrite_commit.close()

        stale_commit = table.new_batch_write_builder().new_commit()
        with self.assertRaises(RuntimeError) as ctx:
            stale_commit.commit(messages)
        stale_commit.close()
        self.assertIn(
            'Global index row ID existence conflict',
            str(ctx.exception),
        )
        self.assertEqual(
            [],
            IndexFileHandler(table).scan(table.snapshot_manager().get_latest_snapshot()),
        )

    def test_create_bitmap_global_index_from_python(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [1, 2, 3, 4, 5],
                'name': ['a', 'b', 'c', 'd', 'e'],
                'age': [10, 20, 30, 40, 50],
                'city': ['vip', 'trial', None, 'vip', 'blocked'],
            },
            schema=self.pa_schema,
        ))

        added = table.create_global_index(
            'city',
            index_type='bitmap',
            options={
                'sorted-index.records-per-range': '2',
                'bitmap-index.dictionary-block-size': '1 b',
            },
        )

        self.assertEqual(3, added)
        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = IndexFileHandler(table).scan(snapshot)
        self.assertEqual(3, len(entries))
        self.assertEqual({'bitmap'}, {e.index_file.index_type for e in entries})
        self.assertEqual([1, 2, 2],
                         sorted(e.index_file.row_count for e in entries))
        self.assertEqual(
            {0},
            {e.index_file.global_index_meta.row_range_start for e in entries},
        )
        self.assertEqual(
            {4},
            {e.index_file.global_index_meta.row_range_end for e in entries},
        )

        read_builder = table.new_read_builder()
        predicate_builder = read_builder.new_predicate_builder()
        cases = [
            (predicate_builder.is_in('city', ['vip', 'trial']),
             [Range(0, 1), Range(3, 3)]),
            (predicate_builder.is_null('city'), [Range(2, 2)]),
            (predicate_builder.not_equal('city', 'blocked'),
             [Range(0, 1), Range(3, 3)]),
        ]
        for predicate, expected in cases:
            with DataEvolutionGlobalIndexScanner.create(
                    table,
                    predicate=predicate,
                    snapshot=snapshot) as scanner:
                result = scanner.scan(predicate)
            self.assertEqual(expected, result.results().to_range_list())

    def test_create_bitmap_global_index_rejects_unsupported_compression(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [1],
                'name': ['a'],
                'age': [10],
                'city': ['vip'],
            },
            schema=self.pa_schema,
        ))

        with self.assertRaisesRegex(ValueError, 'bitmap-index.compression=none'):
            table.create_global_index(
                'city',
                index_type='bitmap',
                options={'bitmap-index.compression': 'lz4'},
            )

    def test_sorted_index_records_per_range_matches_java_floating_factor(self):
        table = self._create_table()
        rows = list(range(12))
        self._write_arrow(table, pa.table(
            {
                'id': rows,
                'name': ['n%s' % i for i in rows],
                'age': rows,
                'city': ['c%s' % i for i in rows],
            },
            schema=self.pa_schema,
        ))

        added = table.create_global_index(
            'id',
            options={'sorted-index.records-per-range': '10'},
        )

        self.assertEqual(1, added)

    def test_create_global_index_uses_external_path(self):
        external_root = 'file://%s' % os.path.join(
            self.tempdir, 'global-index-external')
        options = dict(self.table_options)
        options['global-index.external-path'] = external_root
        table = self._create_table(options=options)
        self._write_arrow(table, pa.table(
            {
                'id': [1, 2],
                'name': ['a', 'b'],
                'age': [10, 20],
                'city': ['x', 'y'],
            },
            schema=self.pa_schema,
        ))

        self.assertEqual(1, table.create_global_index('id'))

        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = IndexFileHandler(table).scan(snapshot)
        self.assertEqual(1, len(entries))
        external_path = entries[0].index_file.external_path
        self.assertIsNotNone(external_path)
        self.assertTrue(external_path.startswith(external_root + '/'))
        self.assertTrue(table.file_io.exists(external_path))

    def test_create_global_index_skips_existing_ranges(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [3, 1, 2, 2],
                'name': ['c', 'a', 'b1', 'b2'],
                'age': [30, 10, 20, 21],
                'city': ['z', 'x', 'y', 'y2'],
            },
            schema=self.pa_schema,
        ))
        options = {'sorted-index.records-per-range': '2'}

        self.assertEqual(2, table.create_global_index('id', options=options))
        snapshot = table.snapshot_manager().get_latest_snapshot()
        self.assertEqual(2, len(IndexFileHandler(table).scan(snapshot)))

        self.assertEqual(0, table.create_global_index('id', options=options))

        latest_snapshot = table.snapshot_manager().get_latest_snapshot()
        self.assertEqual(2, len(IndexFileHandler(table).scan(latest_snapshot)))

        self._write_arrow(table, pa.table(
            {
                'id': [5, 4],
                'name': ['e', 'd'],
                'age': [50, 40],
                'city': ['v', 'u'],
            },
            schema=self.pa_schema,
        ))

        self.assertEqual(1, table.create_global_index('id', options=options))
        latest_snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = sorted(
            IndexFileHandler(table).scan(latest_snapshot),
            key=lambda entry: (
                entry.index_file.global_index_meta.row_range_start,
                entry.index_file.file_name,
            ),
        )
        self.assertEqual(
            [(0, 3), (0, 3), (4, 5)],
            [
                (
                    entry.index_file.global_index_meta.row_range_start,
                    entry.index_file.global_index_meta.row_range_end,
                )
                for entry in entries
            ],
        )
        self.assertEqual(0, table.create_global_index('id', options=options))
        self.assertEqual(
            3,
            len(IndexFileHandler(table).scan(
                table.snapshot_manager().get_latest_snapshot())),
        )

    def test_create_btree_global_index_for_java_scalar_types(self):
        schema = pa.schema([
            ('flag', pa.bool_()),
            ('amount', pa.decimal128(10, 2)),
            ('dt', pa.date32()),
            ('ts', pa.timestamp('us')),
            ('payload', pa.string()),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        self._write_arrow(table, pa.table(
            {
                'flag': [True, False, True],
                'amount': [
                    Decimal('10.25'), Decimal('20.50'), Decimal('30.75')],
                'dt': [
                    date(2026, 6, 18),
                    date(2026, 6, 19),
                    date(2026, 6, 20),
                ],
                'ts': [
                    datetime(2026, 6, 18, 10, 0, 0, 123456),
                    datetime(2026, 6, 19, 10, 0, 0, 123456),
                    datetime(2026, 6, 20, 10, 0, 0, 123456),
                ],
                'payload': ['a', 'b', 'c'],
            },
            schema=schema,
        ))

        for column in ['flag', 'amount', 'dt', 'ts']:
            self.assertEqual(1, table.create_global_index(column))

    def test_create_vindex_global_index_from_python(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('embedding', pa.list_(pa.float32())),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        vectors = pa.array(
            [[1.0, 0.0], [0.0, 1.0], None],
            type=pa.list_(pa.float32()),
        )
        self._write_arrow(table, pa.table(
            {'id': [1, 2, 3], 'embedding': vectors},
            schema=schema,
        ))

        old_module = sys.modules.get("paimon_vindex")
        sys.modules["paimon_vindex"] = types.SimpleNamespace(
            VectorIndexTrainer=_FakeVectorIndexTrainer,
            VectorIndexWriter=_FakeVectorIndexWriter)
        _FakeVectorIndexWriter.instances = []
        try:
            added = table.create_global_index(
                'embedding',
                index_type='ivf-flat',
                options={
                    'ivf-flat.dimension': '2',
                    'ivf-flat.distance.metric': 'l2',
                    'ivf-flat.nlist': '1',
                },
            )
        finally:
            if old_module is None:
                sys.modules.pop("paimon_vindex", None)
            else:
                sys.modules["paimon_vindex"] = old_module

        self.assertEqual(1, added)
        self.assertEqual(1, len(_FakeVectorIndexWriter.instances))
        fake_writer = _FakeVectorIndexWriter.instances[0]
        self.assertEqual('ivf_flat', fake_writer.options['index.type'])
        self.assertEqual('2', fake_writer.options['dimension'])
        self.assertEqual('l2', fake_writer.options['metric'])
        self.assertEqual('1', fake_writer.options['nlist'])
        self.assertEqual([[1.0, 0.0], [0.0, 1.0]], fake_writer.trained)
        self.assertEqual([0, 1], fake_writer.added_ids)
        self.assertEqual([[1.0, 0.0], [0.0, 1.0]], fake_writer.added_vectors)
        self.assertTrue(fake_writer.closed)

        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = IndexFileHandler(table).scan(snapshot)
        self.assertEqual(1, len(entries))
        entry = entries[0]
        self.assertEqual('ivf-flat', entry.index_file.index_type)
        self.assertEqual(3, entry.index_file.row_count)
        self.assertEqual(b'{}', bytes(entry.index_file.global_index_meta.index_meta))
        self.assertTrue(table.file_io.exists(
            table.path_factory().global_index_path_factory().to_path(
                entry.index_file.file_name)))

    def test_create_vindex_global_index_respects_row_count_per_shard(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('embedding', pa.list_(pa.float32())),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        vectors = pa.array(
            [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5], [0.2, 0.8], [0.9, 0.1]],
            type=pa.list_(pa.float32()),
        )
        self._write_arrow(table, pa.table(
            {'id': [1, 2, 3, 4, 5], 'embedding': vectors},
            schema=schema,
        ))

        old_module = sys.modules.get("paimon_vindex")
        sys.modules["paimon_vindex"] = types.SimpleNamespace(
            VectorIndexTrainer=_FakeVectorIndexTrainer,
            VectorIndexWriter=_FakeVectorIndexWriter)
        _FakeVectorIndexWriter.instances = []
        try:
            added = table.create_global_index(
                'embedding',
                index_type='ivf-flat',
                options={
                    'global-index.row-count-per-shard': '2',
                    'ivf-flat.dimension': '2',
                },
            )
        finally:
            if old_module is None:
                sys.modules.pop("paimon_vindex", None)
            else:
                sys.modules["paimon_vindex"] = old_module

        self.assertEqual(3, added)
        self.assertEqual(3, len(_FakeVectorIndexWriter.instances))
        self.assertEqual(
            [[0, 1], [0, 1], [0]],
            [writer.added_ids for writer in _FakeVectorIndexWriter.instances],
        )
        self.assertEqual(
            ['2', '2', '1'],
            [
                writer.options['expected-vector-count']
                for writer in _FakeVectorIndexWriter.instances
            ],
        )

        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = sorted(
            IndexFileHandler(table).scan(snapshot),
            key=lambda entry: entry.index_file.global_index_meta.row_range_start,
        )
        self.assertEqual(
            [(0, 1, 2), (2, 3, 2), (4, 4, 1)],
            [
                (
                    entry.index_file.global_index_meta.row_range_start,
                    entry.index_file.global_index_meta.row_range_end,
                    entry.index_file.row_count,
                )
                for entry in entries
            ],
        )

    def test_create_vindex_global_index_skips_existing_ranges(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('embedding', pa.list_(pa.float32())),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        self._write_arrow(table, pa.table(
            {
                'id': [1, 2, 3],
                'embedding': pa.array(
                    [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]],
                    type=pa.list_(pa.float32()),
                ),
            },
            schema=schema,
        ))

        old_module = sys.modules.get("paimon_vindex")
        sys.modules["paimon_vindex"] = types.SimpleNamespace(
            VectorIndexTrainer=_FakeVectorIndexTrainer,
            VectorIndexWriter=_FakeVectorIndexWriter)
        _FakeVectorIndexWriter.instances = []
        options = {
            'global-index.row-count-per-shard': '2',
            'ivf-flat.dimension': '2',
        }
        try:
            self.assertEqual(
                2,
                table.create_global_index(
                    'embedding',
                    index_type='ivf-flat',
                    options=options,
                ),
            )
            self.assertEqual(
                0,
                table.create_global_index(
                    'embedding',
                    index_type='ivf-flat',
                    options=options,
                ),
            )

            self._write_arrow(table, pa.table(
                {
                    'id': [4, 5],
                    'embedding': pa.array(
                        [[0.2, 0.8], [0.9, 0.1]],
                        type=pa.list_(pa.float32()),
                    ),
                },
                schema=schema,
            ))

            self.assertEqual(
                2,
                table.create_global_index(
                    'embedding',
                    index_type='ivf-flat',
                    options=options,
                ),
            )
            self.assertEqual(
                0,
                table.create_global_index(
                    'embedding',
                    index_type='ivf-flat',
                    options=options,
                ),
            )
        finally:
            if old_module is None:
                sys.modules.pop("paimon_vindex", None)
            else:
                sys.modules["paimon_vindex"] = old_module

        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = sorted(
            IndexFileHandler(table).scan(snapshot),
            key=lambda entry: entry.index_file.global_index_meta.row_range_start,
        )
        self.assertEqual(
            [(0, 1), (2, 2), (3, 3), (4, 4)],
            [
                (
                    entry.index_file.global_index_meta.row_range_start,
                    entry.index_file.global_index_meta.row_range_end,
                )
                for entry in entries
            ],
        )

    def test_create_native_fulltext_global_index_from_python(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('content', pa.string()),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        self._write_arrow(table, pa.table(
            {
                'id': [1, 2, 3],
                'content': [
                    'Apache Paimon full text',
                    None,
                    'Native full text search',
                ],
            },
            schema=schema,
        ))

        ftindex = _FakeFullTextForBuild()
        old_ftindex = sys.modules.get("paimon_ftindex")
        sys.modules["paimon_ftindex"] = ftindex
        try:
            added = table.create_global_index(
                'content',
                index_type=FULL_TEXT_IDENTIFIER,
                options={
                    'global-index.row-count-per-shard': '2',
                    'full-text.tokenizer': 'ngram',
                    'full-text.ngram.min-gram': '2',
                    'full-text.ngram.max-gram': '3',
                    'full-text.ngram.prefix-only': 'true',
                    'full-text.with-position': 'false',
                },
            )
        finally:
            if old_ftindex is None:
                sys.modules.pop("paimon_ftindex", None)
            else:
                sys.modules["paimon_ftindex"] = old_ftindex

        self.assertEqual(2, added)
        self.assertEqual(2, len(ftindex.writers))
        self.assertEqual(
            {
                "tokenizer": "ngram",
                "ngram.min-gram": "2",
                "ngram.max-gram": "3",
                "ngram.prefix-only": "true",
                "with-position": "false",
            },
            ftindex.writers[0].options,
        )
        self.assertEqual(
            [{'row_id': 0, 'text': 'Apache Paimon full text'}],
            ftindex.writers[0].documents,
        )
        self.assertEqual(
            [{'row_id': 0, 'text': 'Native full text search'}],
            ftindex.writers[1].documents,
        )
        self.assertTrue(all(writer.written for writer in ftindex.writers))
        self.assertTrue(all(writer.closed for writer in ftindex.writers))

        snapshot = table.snapshot_manager().get_latest_snapshot()
        entries = sorted(
            IndexFileHandler(table).scan(snapshot),
            key=lambda entry: entry.index_file.global_index_meta.row_range_start,
        )
        self.assertEqual(
            [(0, 1, 2), (2, 2, 1)],
            [
                (
                    entry.index_file.global_index_meta.row_range_start,
                    entry.index_file.global_index_meta.row_range_end,
                    entry.index_file.row_count,
                )
                for entry in entries
            ],
        )
        self.assertEqual(
            {FULL_TEXT_IDENTIFIER},
            {entry.index_file.index_type for entry in entries},
        )
        for entry in entries:
            index_options = NativeFullTextIndexOptions.deserialize(
                entry.index_file.global_index_meta.index_meta)
            self.assertEqual(
                {
                    "tokenizer": "ngram",
                    "ngram.min-gram": "2",
                    "ngram.max-gram": "3",
                    "ngram.prefix-only": "true",
                    "with-position": "false",
                },
                index_options.to_native_options(),
            )
            file_path = table.path_factory().global_index_path_factory().to_path(
                entry.index_file.file_name)
            self.assertEqual(
                b"fake-ftindex",
                table.file_io.new_input_stream(file_path).read(),
            )

    def test_native_fulltext_writer_accepts_jieba_tokenizer(self):
        table = self._create_table()
        index_path = (
            table.path_factory()
            .global_index_path_factory()
            .global_index_root_path()
        )
        ftindex = _FakeFullTextForBuild()
        old_ftindex = sys.modules.get("paimon_ftindex")
        sys.modules["paimon_ftindex"] = ftindex
        try:
            writer = NativeFullTextIndexWriter(
                table.file_io,
                index_path,
                AtomicType('STRING'),
                {'full-text.tokenizer': 'jieba'},
            )
            writer.write('北京大学支持全文检索', 0)
            entries = writer.finish()
        finally:
            if old_ftindex is None:
                sys.modules.pop("paimon_ftindex", None)
            else:
                sys.modules["paimon_ftindex"] = old_ftindex

        self.assertEqual(1, len(entries))
        self.assertEqual([{"row_id": 0, "text": "北京大学支持全文检索"}],
                         ftindex.writers[0].documents)
        self.assertEqual({"tokenizer": "jieba"}, ftindex.writers[0].options)
        self.assertTrue(ftindex.writers[0].closed)

    def test_native_fulltext_options_are_passed_through(self):
        ngram = NativeFullTextIndexOptions.from_options({
            'full-text.tokenizer': 'ngram',
            'full-text.ngram.min-gram': '2',
            'full-text.ngram.max-gram': '3',
            'full-text.ngram.prefix-only': 'true',
            'full-text.lower-case': 'false',
            'full-text.custom-future-option': 'future-value',
            'unrelated': 'ignored',
        })

        self.assertEqual(
            {
                "tokenizer": "ngram",
                "ngram.min-gram": "2",
                "ngram.max-gram": "3",
                "ngram.prefix-only": "true",
                "lower-case": "false",
                "custom-future-option": "future-value",
            },
            ngram.to_native_options(),
        )

    def test_create_native_fulltext_global_index_rejects_non_string_column(self):
        table = self._create_table()
        self._write_arrow(table, pa.table(
            {
                'id': [1],
                'name': ['a'],
                'age': [10],
                'city': ['x'],
            },
            schema=self.pa_schema,
        ))

        with self.assertRaisesRegex(ValueError, 'requires string type'):
            table.create_global_index(
                'id',
                index_type=FULL_TEXT_IDENTIFIER,
            )

    def test_create_vindex_global_index_rejects_generic_unsupported_tables(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('embedding', pa.list_(pa.float32())),
        ])
        bucket_options = dict(self.table_options)
        bucket_options['bucket'] = '1'
        bucket_table = self._create_table(pa_schema=schema, options=bucket_options)
        with self.assertRaisesRegex(ValueError, 'unaware-bucket'):
            bucket_table.create_global_index(
                'embedding',
                index_type='ivf-flat',
                options={'ivf-flat.dimension': '2'},
            )

        dv_options = dict(self.table_options)
        dv_options['deletion-vectors.enabled'] = 'true'
        dv_table = self._create_table(pa_schema=schema, options=dv_options)
        with self.assertRaisesRegex(ValueError, 'deletion vectors'):
            dv_table.create_global_index(
                'embedding',
                index_type='ivf-flat',
                options={'ivf-flat.dimension': '2'},
            )

    def test_vindex_native_options_follow_java_mapping(self):
        data_type = ArrayType(True, AtomicType('FLOAT'))
        options = {
            'ivf-pq.dimension': '128',
            'ivf-pq.distance.metric': 'cosine',
            'ivf-pq.nlist': '256',
            'ivf-pq.pq.m': '16',
            'fields.embedding.dimension': '64',
            'fields.embedding.nlist': '512',
            'fields.embedding.pq.use-opq': 'true',
        }

        result = native_options(data_type, options, 'ivf-pq', 'embedding')

        self.assertEqual('ivf_pq', result['index.type'])
        self.assertEqual('64', result['dimension'])
        self.assertEqual('cosine', result['metric'])
        self.assertEqual('512', result['nlist'])
        self.assertEqual('16', result['pq.m'])
        self.assertEqual('true', result['use-opq'])

    def test_vindex_native_options_support_030_indexes(self):
        data_type = ArrayType(True, AtomicType('FLOAT'))
        options = {
            'ivf-rq.rq.bits': '5',
            'ivf-rq.max-bytes-per-vector': '96',
            'diskann.build-preset': 'balanced',
            'diskann.pq.code-ratio': '0.0625',
            'diskann.raw-vector-encoding': 'f16',
        }

        rq_result = native_options(
            data_type, options, 'ivf-rq', 'embedding')
        self.assertEqual('ivf_rq', rq_result['index.type'])
        self.assertEqual('5', rq_result['rq.bits'])
        self.assertEqual('96', rq_result['max-bytes-per-vector'])
        self.assertEqual('inner_product', rq_result['metric'])

        diskann_result = native_options(
            data_type, options, 'diskann', 'embedding')
        self.assertEqual('diskann', diskann_result['index.type'])
        self.assertEqual(
            'balanced', diskann_result['diskann.build-preset'])
        self.assertEqual('0.0625', diskann_result['pq.code-ratio'])
        self.assertEqual(
            'f16', diskann_result['diskann.raw-vector-encoding'])

    def test_vindex_training_sample_ratio(self):
        options = {
            'ivf-rq.train.sample-ratio': '0.5',
            'fields.embedding.train.sample-ratio': '0.25',
        }
        self.assertEqual(
            0.25, train_sample_ratio(options, 'ivf-rq', 'embedding'))
        self.assertEqual(
            0.5, train_sample_ratio(options, 'ivf-rq', 'other'))

        import numpy as np
        vectors = np.arange(20, dtype=np.float32).reshape(10, 2)
        sampled = _sample_training_vectors(np, vectors, 0.4)
        self.assertEqual(
            [[0.0, 1.0], [4.0, 5.0], [10.0, 11.0], [14.0, 15.0]],
            sampled.tolist(),
        )

    def test_split_by_contiguous_row_range_matches_java_builder(self):
        split = _FakeSplit([
            _FakeFile('a', 0, 2),
            _FakeFile('b', 4, 2),
            _FakeFile('c', 6, 1),
            _FakeFile('d', 10, 1),
        ])

        splits = _split_one_by_contiguous_row_range(split)

        self.assertEqual(
            [['a'], ['b', 'c'], ['d']],
            [[file.file_name for file in s.files] for s in splits],
        )

    def test_split_by_global_index_shard_matches_java_default_builder(self):
        split = _FakeSplit([
            _FakeFile('a', 0, 3),
            _FakeFile('b', 3, 2),
            _FakeFile('c', 6, 3),
        ])

        shards = _split_by_global_index_shard([split], 4)

        self.assertEqual(
            [
                (['a', 'b'], 0, 3),
                (['b'], 4, 4),
                (['c'], 6, 7),
                (['c'], 8, 8),
            ],
            [
                ([file.file_name for file in shard.files], row_range.from_, row_range.to)
                for shard, row_range in shards
            ],
        )

    def test_split_by_global_index_shard_uses_unindexed_ranges(self):
        split = _FakeSplit([
            _FakeFile('a', 0, 100),
        ])

        shards = _split_by_global_index_shard(
            [split], 100, [Range(0, 9), Range(90, 99)])

        self.assertEqual(
            [
                (['a'], 0, 9),
                (['a'], 90, 99),
            ],
            [
                ([file.file_name for file in shard.files], row_range.from_, row_range.to)
                for shard, row_range in shards
            ],
        )

    def test_split_by_global_index_shard_skips_files_without_row_ids(self):
        split = _FakeSplit([
            _FakeFile('no-row-id', None, 2),
            _FakeFile('indexed', 4, 2),
        ])

        shards = _split_by_global_index_shard([split], 4)

        self.assertEqual(
            [(['indexed'], 4, 5)],
            [
                ([file.file_name for file in shard.files], row_range.from_, row_range.to)
                for shard, row_range in shards
            ],
        )

    def test_filter_non_indexable_splits_matches_java_generic_builder(self):
        split = _FakeSplit([
            _FakeFile('indexable-before', 0, 2, schema_id=0),
            _FakeFile('non-indexable', 2, 2, schema_id=1),
            _FakeFile('indexable-after-boundary', 4, 2, schema_id=0),
        ])
        table = types.SimpleNamespace(
            schema_manager=_FakeSchemaManager({
                0: ['id', 'embedding'],
                1: ['id'],
            })
        )

        splits = _filter_non_indexable_splits(
            table, [split], ['embedding'])

        self.assertEqual(1, len(splits))
        self.assertEqual(
            ['indexable-before'],
            [file.file_name for file in splits[0].files],
        )

    def test_vindex_writer_close_cleans_temp_files(self):
        schema = pa.schema([
            ('id', pa.int32()),
            ('embedding', pa.list_(pa.float32())),
        ])
        table = self._create_table(pa_schema=schema, options=self.table_options)
        writer = VindexVectorIndexWriter(
            table.file_io,
            table.path_factory().global_index_path_factory().global_index_root_path(),
            ArrayType(True, AtomicType('FLOAT')),
            'ivf-flat',
            {'ivf-flat.dimension': '2'},
            'embedding',
        )
        writer.write([1.0, 0.0], 0)
        row_id_temp_path = writer._row_id_temp_path
        vector_temp_path = writer._vector_temp_path
        self.assertTrue(os.path.exists(row_id_temp_path))
        self.assertTrue(os.path.exists(vector_temp_path))

        writer.close()

        self.assertFalse(os.path.exists(row_id_temp_path))
        self.assertFalse(os.path.exists(vector_temp_path))

    def test_java_scalar_key_serializers_round_trip(self):
        cases = [
            ('BOOLEAN', True),
            ('TINYINT', -7),
            ('SMALLINT', 1024),
            ('INT', 42),
            ('BIGINT', 1234567890123),
            ('FLOAT', 1.25),
            ('DOUBLE', 3.14159),
            ('DECIMAL(20, 5)', Decimal('-1234567890123.45678')),
            ('DATE', date(2026, 6, 18)),
            ('TIME(3)', datetime(2026, 6, 18, 1, 2, 3, 456000).time()),
            ('TIMESTAMP(6)', datetime(2026, 6, 18, 1, 2, 3, 456789)),
            ('VARCHAR(16)', 'abc'),
        ]

        for type_name, value in cases:
            with self.subTest(type_name=type_name):
                serializer = create_serializer(AtomicType(type_name))
                actual = serializer.deserialize(serializer.serialize(value))
                if type_name == 'FLOAT':
                    self.assertAlmostEqual(value, actual, places=6)
                elif type_name == 'DOUBLE':
                    self.assertAlmostEqual(value, actual, places=12)
                else:
                    self.assertEqual(value, actual)


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