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

import pandas as pd
import pyarrow as pa

from pypaimon.common.options import Options
from pypaimon.catalog.catalog_context import CatalogContext
from pypaimon import CatalogFactory
from pypaimon.catalog.catalog_exception import TableNotExistException
from pypaimon.catalog.rest.rest_catalog import RESTCatalog
from pypaimon.common.identifier import Identifier
from pypaimon import Schema
from pypaimon.tests.rest.rest_base_test import RESTBaseTest

import ray


class RESTTableReadWriteTest(RESTBaseTest):

    def test_overwrite(self):
        simple_pa_schema = pa.schema([
            ('f0', pa.int32()),
            ('f1', pa.string())
        ])
        schema = Schema.from_pyarrow_schema(simple_pa_schema, partition_keys=['f0'],
                                            options={'dynamic-partition-overwrite': 'false'})
        self.rest_catalog.create_table('default.test_overwrite', schema, False)
        table = self.rest_catalog.get_table('default.test_overwrite')
        read_builder = table.new_read_builder()

        # test normal write
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()

        df0 = pd.DataFrame({
            'f0': [1, 2],
            'f1': ['apple', 'banana'],
        })

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

        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual_df0 = table_read.to_pandas(table_scan.plan().splits()).sort_values(by='f0')
        df0['f0'] = df0['f0'].astype('int32')
        pd.testing.assert_frame_equal(
            actual_df0.reset_index(drop=True), df0.reset_index(drop=True))

        # test partially overwrite
        write_builder = table.new_batch_write_builder().overwrite({'f0': 1})
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()

        df1 = pd.DataFrame({
            'f0': [1],
            'f1': ['watermelon'],
        })

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

        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual_df1 = table_read.to_pandas(table_scan.plan().splits()).sort_values(by='f0')
        expected_df1 = pd.DataFrame({
            'f0': [1, 2],
            'f1': ['watermelon', 'banana']
        })
        expected_df1['f0'] = expected_df1['f0'].astype('int32')
        pd.testing.assert_frame_equal(
            actual_df1.reset_index(drop=True), expected_df1.reset_index(drop=True))

        # test fully overwrite
        write_builder = table.new_batch_write_builder().overwrite()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()

        df2 = pd.DataFrame({
            'f0': [3],
            'f1': ['Neo'],
        })

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

        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual_df2 = table_read.to_pandas(table_scan.plan().splits())
        df2['f0'] = df2['f0'].astype('int32')
        pd.testing.assert_frame_equal(
            actual_df2.reset_index(drop=True), df2.reset_index(drop=True))

    def test_dynamic_partition_overwrite(self):
        pa_schema = pa.schema([
            ('f0', pa.string()),
            ('f1', pa.string())
        ])
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['f0'])
        self.rest_catalog.create_table('default.test_dynamic_overwrite', schema, False)
        table = self.rest_catalog.get_table('default.test_dynamic_overwrite')
        read_builder = table.new_read_builder()

        # Write initial non-null and null partitions
        self._batch_write(table, pd.DataFrame({
            'f0': ['a', 'b', None],
            'f1': ['apple', 'banana', 'cherry'],
        }))

        # Dynamic overwrite partition f0='a' only; 'b' and null untouched
        self._batch_overwrite(table, pd.DataFrame({
            'f0': ['a'],
            'f1': ['watermelon'],
        }))

        self._assert_table_equals(read_builder, pd.DataFrame({
            'f0': ['a', 'b', None],
            'f1': ['watermelon', 'banana', 'cherry'],
        }), sort_by='f0')

        # Dynamic overwrite partitions f0='a' and f0=None; 'b' untouched
        self._batch_overwrite(table, pd.DataFrame({
            'f0': ['a', None],
            'f1': ['mango', 'grape'],
        }))

        self._assert_table_equals(read_builder, pd.DataFrame({
            'f0': ['a', 'b', None],
            'f1': ['mango', 'banana', 'grape'],
        }), sort_by='f0')

    def _batch_write(self, table, df):
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_pandas(df)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def _batch_overwrite(self, table, df, partition=None):
        write_builder = table.new_batch_write_builder().overwrite(partition)
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_pandas(df)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

    def _assert_table_equals(self, read_builder, expected_df, sort_by=None):
        table_scan = read_builder.new_scan()
        table_read = read_builder.new_read()
        actual_df = table_read.to_pandas(table_scan.plan().splits())
        if sort_by:
            actual_df = actual_df.sort_values(by=sort_by)
        pd.testing.assert_frame_equal(
            actual_df.reset_index(drop=True), expected_df.reset_index(drop=True))

    def test_parquet_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_append_only_parquet', schema, False)
        table = self.rest_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_orc_ao_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'orc'})
        self.rest_catalog.create_table('default.test_append_only_orc', schema, False)
        table = self.rest_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.rest_catalog.create_table('default.test_append_only_avro', schema, False)
        table = self.rest_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.rest_catalog.create_table('default.test_append_only_lance', schema, False)
        table = self.rest_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_ao_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_append_only_filter', schema, False)
        table = self.rest_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.rest_catalog.create_table('default.test_append_only_projection', schema, False)
        table = self.rest_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.rest_catalog.create_table('default.test_avro_append_only_projection', schema, False)
        table = self.rest_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.rest_catalog.create_table('default.test_append_only_limit', schema, False)
        table = self.rest_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_pk_parquet_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={'bucket': '2'})
        self.rest_catalog.create_table('default.test_pk_parquet', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_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.pk_expected)

    def test_pk_orc_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={
                                                'bucket': '1',
                                                'file.format': 'orc'
                                            })
        self.rest_catalog.create_table('default.test_pk_orc', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_orc')
        self._write_test_table(table)

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

        # when bucket=1, actual field name will contain 'not null', so skip comparing field name
        for i in range(len(actual.columns)):
            col_a = actual.column(i)
            col_b = self.expected.column(i)
            self.assertEqual(col_a, col_b)

    def test_pk_avro_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={
                                                'bucket': '2',
                                                'file.format': 'avro'
                                            })
        self.rest_catalog.create_table('default.test_pk_avro', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_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.pk_expected)

    def test_pk_lance_reader(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={
                                                'bucket': '2',
                                                'file.format': 'lance'
                                            })
        self.rest_catalog.drop_table('default.test_pk_lance', True)
        self.rest_catalog.create_table('default.test_pk_lance', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_lance')
        # Use table's schema for writing to ensure schema consistency
        from pypaimon.schema.data_types import PyarrowFieldParser
        table_pa_schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
        self._write_test_table_with_schema(table, table_pa_schema)

        read_builder = table.new_read_builder()
        actual = self._read_test_table(read_builder).sort_by('user_id')
        self.assertEqual(actual, self.pk_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.rest_catalog.create_table('default.test_append_only_lance_filter', schema, False)
        table = self.rest_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)
        p4 = predicate_builder.is_not_in('behavior', ['b', 'e'])
        p5 = predicate_builder.is_in('dt', ['p1'])
        p6 = predicate_builder.is_not_null('behavior')
        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_pk_lance_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={
                                                'bucket': '2',
                                                'file.format': 'lance'
                                            })
        self.rest_catalog.create_table('default.test_pk_lance_filter', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_lance_filter')
        from pypaimon.schema.data_types import PyarrowFieldParser
        table_pa_schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
        self._write_test_table_with_schema(table, table_pa_schema)

        predicate_builder = table.new_read_builder().new_predicate_builder()
        p1 = predicate_builder.is_in('dt', ['p1'])
        p2 = predicate_builder.between('user_id', 2, 7)
        p3 = predicate_builder.is_not_null('behavior')
        g1 = predicate_builder.and_predicates([p1, p2, p3])
        read_builder = table.new_read_builder().with_filter(g1)
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = pa.concat_tables([
            self.pk_expected.slice(1, 1),  # 2/b
            self.pk_expected.slice(5, 1)  # 7/g
        ])
        self.assertEqual(actual, expected)

    def test_pk_lance_reader_no_bucket(self):
        """Test Lance format with PrimaryKey table without specifying bucket."""
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={'file.format': 'lance'})
        self.rest_catalog.drop_table('default.test_pk_lance_no_bucket', True)
        self.rest_catalog.create_table('default.test_pk_lance_no_bucket', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_lance_no_bucket')
        from pypaimon.schema.data_types import PyarrowFieldParser
        table_pa_schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
        self._write_test_table_with_schema(table, table_pa_schema)

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

    def test_pk_reader_with_filter(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={'bucket': '2'})
        self.rest_catalog.create_table('default.test_pk_filter', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_filter')
        self._write_test_table(table)

        predicate_builder = table.new_read_builder().new_predicate_builder()
        p1 = predicate_builder.is_in('dt', ['p1'])
        p2 = predicate_builder.between('user_id', 2, 7)
        p3 = predicate_builder.is_not_null('behavior')
        g1 = predicate_builder.and_predicates([p1, p2, p3])
        read_builder = table.new_read_builder().with_filter(g1)
        actual = self._read_test_table(read_builder).sort_by('user_id')
        expected = pa.concat_tables([
            self.pk_expected.slice(1, 1),  # 2/b
            self.pk_expected.slice(5, 1)  # 7/g
        ])
        self.assertEqual(actual, expected)

    def test_pk_reader_with_projection(self):
        schema = Schema.from_pyarrow_schema(self.pa_schema,
                                            partition_keys=['dt'],
                                            primary_keys=['user_id', 'dt'],
                                            options={'bucket': '2'})
        self.rest_catalog.create_table('default.test_pk_projection', schema, False)
        table = self.rest_catalog.get_table('default.test_pk_projection')
        self._write_test_table(table)

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

    def test_write_wrong_schema(self):
        self.rest_catalog.create_table('default.test_wrong_schema',
                                       Schema.from_pyarrow_schema(self.pa_schema),
                                       False)
        table = self.rest_catalog.get_table('default.test_wrong_schema')

        data = {
            'f0': [1, 2, 3],
            'f1': ['a', 'b', 'c'],
        }
        df = pd.DataFrame(data)
        schema = pa.schema([
            ('f0', pa.int64()),
            ('f1', pa.string())
        ])
        record_batch = pa.RecordBatch.from_pandas(df, schema)

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

        with self.assertRaises(ValueError) as e:
            table_write.write_arrow_batch(record_batch)
        self.assertTrue(str(e.exception).startswith("Input schema isn't consistent with table schema and write cols."))

    def test_reader_iterator(self):
        read_builder = self.table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        iterator = table_read.to_iterator(splits)
        result = []
        value = next(iterator, None)
        while value is not None:
            result.append(value.get_field(1))
            value = next(iterator, None)
        self.assertEqual(result, [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008])

    def test_reader_duckdb(self):
        read_builder = self.table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        duckdb_con = table_read.to_duckdb(splits, 'duckdb_table')
        actual = duckdb_con.query("SELECT * FROM duckdb_table").fetchdf()
        expect = pd.DataFrame(self.raw_data)
        pd.testing.assert_frame_equal(actual.reset_index(drop=True), expect.reset_index(drop=True))

    def test_reader_ray_data(self):
        if not ray.is_initialized():
            ray.init(ignore_reinit_error=True, num_cpus=2)

        read_builder = self.table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()

        ray_dataset = table_read.to_ray(splits, parallelism=2)

        self.assertIsNotNone(ray_dataset, "Ray dataset should not be None")
        self.assertEqual(ray_dataset.count(), 8, "Should have 8 rows")

        df = ray_dataset.to_pandas()
        expect = pd.DataFrame(self.raw_data)
        pd.testing.assert_frame_equal(df.sort_values(by='user_id').reset_index(drop=True),
                                      expect.sort_values(by='user_id').reset_index(drop=True))

    def test_ray_data_write_and_read(self):
        if not ray.is_initialized():
            ray.init(ignore_reinit_error=True, num_cpus=2)

        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])
        schema = Schema.from_pyarrow_schema(pa_schema)
        self.rest_catalog.create_table('default.test_ray_write_read', schema, False)
        table = self.rest_catalog.get_table('default.test_ray_write_read')

        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3],
            'name': ['Alice', 'Bob', 'Charlie'],
            'value': [100, 200, 300],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        table_write.write_arrow(test_data)
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()

        ray_dataset = table_read.to_ray(splits, parallelism=2)

        self.assertIsNotNone(ray_dataset, "Ray dataset should not be None")
        self.assertEqual(ray_dataset.count(), 3, "Should have 3 rows")

        df = ray_dataset.to_pandas()
        expected_df = pd.DataFrame({
            'id': [1, 2, 3],
            'name': ['Alice', 'Bob', 'Charlie'],
            'value': [100, 200, 300],
        })
        expected_df['id'] = expected_df['id'].astype('int32')
        pd.testing.assert_frame_equal(df.sort_values(by='id').reset_index(drop=True),
                                      expected_df.sort_values(by='id').reset_index(drop=True))

    def test_write_wide_table_large_data(self):
        logging.basicConfig(level=logging.INFO)
        catalog = CatalogFactory.create(self.options)

        # Build table structure: 200 data columns + 1 partition column
        # Create PyArrow schema
        pa_fields = []

        # Create 200 data columns f0 to f199
        for i in range(200):
            pa_fields.append(pa.field(f"f{i}", pa.string(), metadata={"description": f"Column f{i}"}))

        # Add partition column dt
        pa_fields.append(pa.field("dt", pa.string(), metadata={"description": "Partition column dt"}))

        # Create PyArrow schema
        pa_schema = pa.schema(pa_fields)

        # Convert to Paimon Schema and specify partition key
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=["dt"])

        # Create table
        table_identifier = Identifier.create("default", "wide_table_200cols")
        try:
            # If table already exists, drop it first
            try:
                catalog.get_table(table_identifier)
                catalog.drop_table(table_identifier)
                print(f"Dropped existing table {table_identifier}")
            except Exception:
                # Table does not exist, continue creating
                pass

            # Create new table
            catalog.create_table(
                identifier=table_identifier,
                schema=schema,
                ignore_if_exists=False
            )

            print(
                f"Successfully created table {table_identifier} with {len(pa_fields) - 1} "
                f"data columns and 1 partition column")
            print(
                f"Table schema: {len([f for f in pa_fields if f.name != 'dt'])} data columns (f0-f199) + dt partition")

        except Exception as e:
            print(f"Error creating table: {e}")
            raise e
        import random

        table_identifier = Identifier.create("default", "wide_table_200cols")
        table = catalog.get_table(table_identifier)

        total_rows = 500000  # rows of data
        batch_size = 100000  # 100,000 rows per batch
        commit_batches = total_rows // batch_size

        for commit_batch in range(commit_batches):
            start_idx = commit_batch * batch_size
            end_idx = start_idx + batch_size

            print(f"Processing batch {commit_batch + 1}/{commit_batches} ({start_idx:,} - {end_idx:,})...")
            # Generate data for current batch - generate data for all 200 columns
            data = {}
            # Generate data for f0-f199
            for i in range(200):
                if i == 0:
                    data[f"f{i}"] = [f'value_{j}' for j in range(start_idx, end_idx)]
                elif i == 1:
                    data[f"f{i}"] = [random.choice(['A', 'B', 'C', 'D', 'E']) for _ in range(batch_size)]
                elif i == 2:
                    data[f"f{i}"] = [f'detail_{random.randint(1, 1000)}' for _ in range(batch_size)]
                elif i == 3:
                    data[f"f{i}"] = [f'id_{j:06d}' for j in range(start_idx, end_idx)]
                else:
                    # Generate random string data for other columns
                    data[f"f{i}"] = [f'col{i}_val_{random.randint(1, 10000)}' for _ in range(batch_size)]

            # Add partition column data
            data['dt'] = ['2025-09-01' for _ in range(batch_size)]
            # Convert dictionary to PyArrow RecordBatch
            arrow_batch = pa.RecordBatch.from_pydict(data)
            # Create new write and commit objects for each commit batch
            write_builder = table.new_batch_write_builder()
            table_write = write_builder.new_write()
            table_commit = write_builder.new_commit()

            try:
                # Write current batch data
                table_write.write_arrow_batch(arrow_batch)
                print("Batch data write completed, committing...")
                # Commit current batch
                commit_messages = table_write.prepare_commit()
                table_commit.commit(commit_messages)
                print(f"Batch {commit_batch + 1} committed successfully! Written {end_idx:,} rows of data")

            finally:
                # Ensure resource cleanup
                table_write.close()
                table_commit.close()

        print(
            f"All data writing completed! "
            f"Total written {total_rows:,} rows of data to 200-column wide table in {commit_batches} commits")
        rest_catalog = RESTCatalog(CatalogContext.create_from_options(Options(self.options)))
        table = rest_catalog.get_table('default.wide_table_200cols')
        predicate_builder = table.new_read_builder().new_predicate_builder()
        read_builder = (table.new_read_builder()
                        .with_projection(['f0', 'f1'])
                        .with_filter(predicate=predicate_builder.equal("dt", "2025-09-01")))
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        self.assertEqual(table_read.to_arrow(splits).num_rows, total_rows)

    def test_drop_partitions(self):
        pa_schema = pa.schema([
            ('col', pa.int32()),
            ('dt', pa.string()),
        ])
        schema = Schema.from_pyarrow_schema(pa_schema, partition_keys=['dt'])
        self.rest_catalog.create_table('default.test_drop_partitions_table', schema, False)
        table = self.rest_catalog.get_table('default.test_drop_partitions_table')

        partition_specs = [{'dt': '20250101'}, {'dt': '20250102'}]
        write_builder = table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        for spec in partition_specs:
            batch = pa.RecordBatch.from_pydict(
                {'col': [0], 'dt': [spec['dt']]},
                schema=pa_schema,
            )
            table_write.write_arrow(pa.Table.from_batches([batch]))
        table_commit.commit(table_write.prepare_commit())
        table_write.close()
        table_commit.close()

        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()
        before = table_read.to_arrow(splits)
        self.assertEqual(before.num_rows, 2, "Should have 2 rows before drop_partitions")

        self.rest_catalog.drop_partitions(
            'default.test_drop_partitions_table',
            partition_specs,
        )

        splits_after = read_builder.new_scan().plan().splits()
        after = table_read.to_arrow(splits_after)
        self.assertEqual(after.num_rows, 0, "Table should be empty after drop_partitions")

    def test_drop_partitions_table_not_exist(self):
        with self.assertRaises(TableNotExistException):
            self.rest_catalog.drop_partitions(
                'default.non_existing_table',
                [{'dt': '20250101'}],
            )

    def test_drop_partitions_empty_list(self):
        with self.assertRaises(ValueError):
            self.rest_catalog.drop_partitions('default.test_reader_iterator', [])
