# 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 tempfile
import unittest
import shutil

import pyarrow as pa
import pyarrow.types as pa_types
import ray

from pypaimon import CatalogFactory, Schema
from pypaimon.common.options.core_options import CoreOptions
from pypaimon.schema.data_types import PyarrowFieldParser


class RayDataTest(unittest.TestCase):
    """Tests for Ray Data integration with PyPaimon."""

    @classmethod
    def setUpClass(cls):
        """Set up test environment."""
        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)

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

        try:
            from ray.data import DataContext
            context = DataContext.get_current()
            if hasattr(context, 'shuffle_strategy'):
                try:
                    from ray.data._internal.execution.interfaces.execution_options import ShuffleStrategy
                    context.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED
                except (ImportError, AttributeError):
                    pass
            if hasattr(context, 'use_polars_sort'):
                context.use_polars_sort = True
        except (ImportError, AttributeError):
            pass

    @classmethod
    def tearDownClass(cls):
        """Clean up test environment."""
        try:
            if ray.is_initialized():
                ray.shutdown()
        except Exception:
            pass
        try:
            shutil.rmtree(cls.tempdir)
        except OSError:
            pass

    def setUp(self):
        """Set up test method."""
        pass

    def test_basic_ray_data_read(self):
        """Test basic Ray Data read from PyPaimon table."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_basic', schema, False)
        table = self.catalog.get_table('default.test_ray_basic')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            'value': [100, 200, 300, 400, 500],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read using Ray Data
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

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

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

        # Test basic operations
        sample_data = ray_dataset.take(3)
        self.assertEqual(len(sample_data), 3, "Should have 3 sample rows")

        # Convert to pandas for verification
        df = ray_dataset.to_pandas()
        self.assertEqual(len(df), 5, "DataFrame should have 5 rows")
        # Sort by id to ensure order-independent comparison
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['id']), [1, 2, 3, 4, 5], "ID column should match")
        self.assertEqual(
            list(df_sorted['name']),
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            "Name column should match"
        )

    def test_basic_ray_data_write(self):
        """Test basic Ray Data write from PyPaimon table."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_write', schema, False)
        table = self.catalog.get_table('default.test_ray_write')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            'value': [100, 200, 300, 400, 500],
        }, schema=pa_schema)
        
        from ray.data.read_api import from_arrow
        ds = from_arrow(test_data)
        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_ray(ds, concurrency=2)
        # Read using Ray Data
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

        arrow_result = table_read.to_arrow(splits)

        # Verify PyArrow table
        self.assertIsNotNone(arrow_result, "Arrow table should not be None")
        self.assertEqual(arrow_result.num_rows, 5, "Should have 5 rows")

        # Test basic operations - get first 3 rows
        sample_table = arrow_result.slice(0, 3)
        self.assertEqual(sample_table.num_rows, 3, "Should have 3 sample rows")

        # Convert to pandas for verification
        df = arrow_result.to_pandas()
        self.assertEqual(len(df), 5, "DataFrame should have 5 rows")
        # Sort by id to ensure order-independent comparison
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['id']), [1, 2, 3, 4, 5], "ID column should match")
        self.assertEqual(
            list(df_sorted['name']),
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            "Name column should match"
        )

    def test_ray_data_with_predicate(self):
        """Test Ray Data read with predicate filtering."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('category', pa.string()),
            ('amount', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_predicate', schema, False)
        table = self.catalog.get_table('default.test_ray_predicate')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'category': ['A', 'B', 'A', 'C', 'B'],
            'amount': [100, 200, 150, 300, 250],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read with predicate
        read_builder = table.new_read_builder()
        predicate_builder = read_builder.new_predicate_builder()
        predicate = predicate_builder.equal('category', 'A')
        read_builder = read_builder.with_filter(predicate)

        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

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

        # Verify filtered results
        self.assertEqual(ray_dataset.count(), 2, "Should have 2 rows after filtering")
        df = ray_dataset.to_pandas()
        self.assertEqual(set(df['category'].tolist()), {'A'}, "All rows should have category='A'")
        self.assertEqual(set(df['id'].tolist()), {1, 3}, "Should have IDs 1 and 3")

    def test_ray_data_with_projection(self):
        """Test Ray Data read with column projection."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
            ('email', pa.string()),
            ('age', pa.int32()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_projection', schema, False)
        table = self.catalog.get_table('default.test_ray_projection')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3],
            'name': ['Alice', 'Bob', 'Charlie'],
            'email': ['alice@example.com', 'bob@example.com', 'charlie@example.com'],
            'age': [25, 30, 35],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read with projection
        read_builder = table.new_read_builder()
        read_builder = read_builder.with_projection(['id', 'name'])

        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

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

        # Verify projection
        self.assertEqual(ray_dataset.count(), 3, "Should have 3 rows")
        df = ray_dataset.to_pandas()
        self.assertEqual(set(df.columns), {'id', 'name'}, "Should only have id and name columns")
        self.assertFalse('email' in df.columns, "Should not have email column")
        self.assertFalse('age' in df.columns, "Should not have age column")

    def test_ray_data_map_operation(self):
        """Test Ray Data map operations after reading from PyPaimon."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_map', schema, False)
        table = self.catalog.get_table('default.test_ray_map')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'value': [10, 20, 30, 40, 50],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read using Ray Data
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

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

        # Apply map operation (double the value)
        def double_value(row):
            row['value'] = row['value'] * 2
            return row

        mapped_dataset = ray_dataset.map(double_value)

        # Verify mapped results
        df = mapped_dataset.to_pandas()
        # Sort by id to ensure order-independent comparison
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['value']), [20, 40, 60, 80, 100], "Values should be doubled")

    def test_ray_data_filter_operation(self):
        """Test Ray Data filter operations after reading from PyPaimon."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('score', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_filter', schema, False)
        table = self.catalog.get_table('default.test_ray_filter')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'score': [60, 70, 80, 90, 100],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read using Ray Data
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

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

        # Apply filter operation (score >= 80)
        filtered_dataset = ray_dataset.filter(lambda row: row['score'] >= 80)

        # Verify filtered results
        df = filtered_dataset.to_pandas()
        self.assertEqual(len(df), 3, "Should have 3 rows with score >= 80")
        self.assertEqual(set(df['id'].tolist()), {3, 4, 5}, "Should have IDs 3, 4, 5")
        self.assertEqual(set(df['score'].tolist()), {80, 90, 100}, "Should have scores 80, 90, 100")

    def test_ray_data_distributed_vs_simple(self):
        """Test that both distributed and simple reading modes work correctly."""
        # Create schema
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_modes', schema, False)
        table = self.catalog.get_table('default.test_ray_modes')

        # Write test data
        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3],
            'value': [10, 20, 30],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Read using distributed mode
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

        ray_dataset_distributed = table_read.to_ray(splits, override_num_blocks=2)
        ray_dataset_simple = table_read.to_ray(splits, override_num_blocks=1)

        # Both should produce the same results
        self.assertEqual(ray_dataset_distributed.count(), 3, "Distributed mode should have 3 rows")
        self.assertEqual(ray_dataset_simple.count(), 3, "Simple mode should have 3 rows")

        df_distributed = ray_dataset_distributed.to_pandas()
        df_simple = ray_dataset_simple.to_pandas()

        # Sort both dataframes by id to ensure order-independent comparison
        df_distributed_sorted = df_distributed.sort_values(by='id').reset_index(drop=True)
        df_simple_sorted = df_simple.sort_values(by='id').reset_index(drop=True)

        self.assertEqual(list(df_distributed_sorted['id']), list(df_simple_sorted['id']), "IDs should match")
        self.assertEqual(list(df_distributed_sorted['value']), list(df_simple_sorted['value']), "Values should match")

    def test_ray_data_primary_key_basic(self):
        """Test Ray Data read from PrimaryKey table."""
        pa_schema = pa.schema([
            pa.field('id', pa.int32(), nullable=False),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema, primary_keys=['id'], options={'bucket': '2'})
        self.catalog.create_table('default.test_ray_pk_basic', schema, False)
        table = self.catalog.get_table('default.test_ray_pk_basic')

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

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

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

        ray_dataset = table_read.to_ray(splits, override_num_blocks=1)
        self.assertIsNotNone(ray_dataset, "Ray dataset should not be None")
        self.assertEqual(ray_dataset.count(), 5, "Should have 5 rows")

        df = ray_dataset.to_pandas()
        self.assertEqual(len(df), 5, "DataFrame should have 5 rows")
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['id']), [1, 2, 3, 4, 5], "ID column should match")
        self.assertEqual(
            list(df_sorted['name']),
            ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
            "Name column should match"
        )

    def test_ray_data_primary_key_update(self):
        """Test Ray Data read from PrimaryKey table with updates (upsert behavior)."""
        pa_schema = pa.schema([
            pa.field('id', pa.int32(), nullable=False),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema, primary_keys=['id'], options={'bucket': '2'})
        self.catalog.create_table('default.test_ray_pk_update', schema, False)
        table = self.catalog.get_table('default.test_ray_pk_update')

        initial_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()
        writer = write_builder.new_write()
        writer.write_arrow(initial_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        updated_data = pa.Table.from_pydict({
            'id': [1, 2, 4],  # id=1,2 updated, id=4 new
            'name': ['Alice-Updated', 'Bob-Updated', 'David'],
            'value': [150, 250, 400],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(updated_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

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

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

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

        df = ray_dataset.to_pandas()
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['id']), [1, 2, 3, 4], "ID column should match")
        self.assertEqual(
            list(df_sorted['name']),
            ['Alice-Updated', 'Bob-Updated', 'Charlie', 'David'],
            "Name column should reflect updates"
        )
        self.assertEqual(list(df_sorted['value']), [150, 250, 300, 400], "Value column should reflect updates")

    def test_ray_data_primary_key_with_predicate(self):
        """Test Ray Data read from PrimaryKey table with predicate filtering."""
        pa_schema = pa.schema([
            pa.field('id', pa.int32(), nullable=False),
            ('category', pa.string()),
            ('amount', pa.int64()),
            pa.field('dt', pa.string(), nullable=False),
        ])

        schema = Schema.from_pyarrow_schema(
            pa_schema,
            primary_keys=['id', 'dt'],
            partition_keys=['dt'],
            options={'bucket': '2'}
        )
        self.catalog.create_table('default.test_ray_pk_predicate', schema, False)
        table = self.catalog.get_table('default.test_ray_pk_predicate')

        test_data = pa.Table.from_pydict({
            'id': [1, 2, 3, 4, 5],
            'category': ['A', 'B', 'A', 'C', 'B'],
            'amount': [100, 200, 150, 300, 250],
            'dt': ['2024-01-01', '2024-01-01', '2024-01-02', '2024-01-02', '2024-01-03'],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(test_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        read_builder = table.new_read_builder()
        predicate_builder = read_builder.new_predicate_builder()
        predicate = predicate_builder.equal('category', 'A')
        read_builder = read_builder.with_filter(predicate)

        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

        ray_dataset = table_read.to_ray(splits, override_num_blocks=1)

        # Verify filtered results
        self.assertEqual(ray_dataset.count(), 2, "Should have 2 rows after filtering")
        df = ray_dataset.to_pandas()
        self.assertEqual(set(df['category'].tolist()), {'A'}, "All rows should have category='A'")
        self.assertEqual(set(df['id'].tolist()), {1, 3}, "Should have IDs 1 and 3")

        # Read with predicate on partition
        read_builder = table.new_read_builder()
        predicate_builder = read_builder.new_predicate_builder()
        predicate = predicate_builder.equal('dt', '2024-01-01')
        read_builder = read_builder.with_filter(predicate)

        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

        ray_dataset = table_read.to_ray(splits, override_num_blocks=1)

        # Verify filtered results by partition
        self.assertEqual(ray_dataset.count(), 2, "Should have 2 rows in partition 2024-01-01")
        df = ray_dataset.to_pandas()
        self.assertEqual(set(df['dt'].tolist()), {'2024-01-01'}, "All rows should be in partition 2024-01-01")

    def test_ray_data_primary_key_multiple_splits_same_bucket(self):
        """Test Ray Data read from PrimaryKey table with small target_split_size."""

        pa_schema = pa.schema([
            pa.field('id', pa.int32(), nullable=False),
            ('name', pa.string()),
            ('value', pa.int64()),
        ])

        schema = Schema.from_pyarrow_schema(
            pa_schema,
            primary_keys=['id'],
            options={
                'bucket': '2',
                CoreOptions.SOURCE_SPLIT_TARGET_SIZE.key(): '1b'
            }
        )
        self.catalog.create_table('default.test_ray_pk_multi_split', schema, False)
        table = self.catalog.get_table('default.test_ray_pk_multi_split')

        initial_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()
        writer = write_builder.new_write()
        writer.write_arrow(initial_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        updated_data = pa.Table.from_pydict({
            'id': [1, 2, 4],
            'name': ['Alice-Updated', 'Bob-Updated', 'David'],
            'value': [150, 250, 400],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(updated_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

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

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

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

        df = ray_dataset.to_pandas()
        df_sorted = df.sort_values(by='id').reset_index(drop=True)
        self.assertEqual(list(df_sorted['id']), [1, 2, 3, 4], "ID column should match")
        self.assertEqual(
            list(df_sorted['name']),
            ['Alice-Updated', 'Bob-Updated', 'Charlie', 'David'],
            "Name column should reflect updates"
        )
        self.assertEqual(list(df_sorted['value']), [150, 250, 300, 400], "Value column should reflect updates")

    def test_ray_data_invalid_parallelism(self):
        pa_schema = pa.schema([
            ('id', pa.int32()),
            ('name', pa.string()),
        ])

        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_invalid_parallelism', schema, False)
        table = self.catalog.get_table('default.test_ray_invalid_parallelism')

        # Write some data
        data = pa.Table.from_pydict({
            'id': [1, 2, 3],
            'name': ['Alice', 'Bob', 'Charlie'],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

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

        with self.assertRaises(ValueError) as context:
            table_read.to_ray(splits, override_num_blocks=0)
        self.assertIn("override_num_blocks must be at least 1", str(context.exception))

        with self.assertRaises(ValueError) as context:
            table_read.to_ray(splits, override_num_blocks=-1)
        self.assertIn("override_num_blocks must be at least 1", str(context.exception))

        with self.assertRaises(ValueError) as context:
            table_read.to_ray(splits, override_num_blocks=-10)
        self.assertIn("override_num_blocks must be at least 1", str(context.exception))

    def test_dict_return_loses_large_binary_type(self):
        # Original data with large_binary
        original = pa.table({
            'data': pa.array([b'hello', b'world'], type=pa.large_binary())
        })
        self.assertTrue(
            pa_types.is_large_binary(original.schema.field('data').type),
            "Original should be large_binary"
        )

        # Simulate map_batches returning dict: convert to Python list then rebuild
        d = {'data': original['data'].to_pylist()}
        rebuilt = pa.Table.from_pydict(d)
        self.assertTrue(
            pa_types.is_binary(rebuilt.schema.field('data').type),
            f"Rebuilt from dict should be binary (PyArrow default), but got {rebuilt.schema.field('data').type}"
        )
        self.assertFalse(
            pa_types.is_large_binary(rebuilt.schema.field('data').type),
            "large_binary type should be lost after dict roundtrip"
        )

    def test_ray_data_read_and_write_with_blob(self):
        import time
        pa_schema = pa.schema([
            ('id', pa.int64()),
            ('name', pa.string()),
            ('data', pa.large_binary()),  # Table uses large_binary for blob
        ])

        schema = Schema.from_pyarrow_schema(
            pa_schema,
            options={
                'row-tracking.enabled': 'true',
                'data-evolution.enabled': 'true',
                'blob-field': 'data',
            }
        )

        table_name = f'default.test_ray_read_write_blob_{int(time.time() * 1000000)}'
        self.catalog.create_table(table_name, schema, False)
        table = self.catalog.get_table(table_name)

        # Step 1: Write data to Paimon table using write_arrow (large_binary type)
        initial_data = pa.Table.from_pydict({
            'id': [1, 2, 3],
            'name': ['Alice', 'Bob', 'Charlie'],
            'data': [b'blob_data_1', b'blob_data_2', b'blob_data_3'],
        }, schema=pa_schema)

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(initial_data)
        commit_messages = writer.prepare_commit()
        commit = write_builder.new_commit()
        commit.commit(commit_messages)
        writer.close()

        # Step 2: Read from Paimon table using to_ray()
        read_builder = table.new_read_builder()
        table_read = read_builder.new_read()
        table_scan = read_builder.new_scan()
        splits = table_scan.plan().splits()

        ray_dataset = table_read.to_ray(splits)

        # Verify Ray blocks preserve large_binary type from Paimon
        for batch in ray_dataset.iter_batches(batch_size=10, batch_format="pyarrow"):
            ray_data_field = batch.schema.field('data')
            self.assertTrue(
                pa_types.is_large_binary(ray_data_field.type),
                f"Ray block should preserve large_binary() from Paimon, but got {ray_data_field.type}"
            )
            break

        # Verify Paimon table schema is large_binary (BLOB)
        table_pa_schema = PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
        self.assertTrue(
            pa_types.is_large_binary(table_pa_schema.field('data').type),
            "Paimon table should have large_binary() for BLOB field"
        )

        # Step 3: Simulate user pipeline: map_batches returns Python dict,
        def process_blob(batch):
            return {
                'id': batch['id'].to_pylist(),
                'name': batch['name'].to_pylist(),
                'data': batch['data'].to_pylist(),  # Python bytes -> binary
            }

        mapped_dataset = ray_dataset.map_batches(process_blob, batch_format="pyarrow")

        # Verify map_batches caused type downgrade: large_binary -> binary
        for batch in mapped_dataset.iter_batches(batch_size=10, batch_format="pyarrow"):
            mapped_data_field = batch.schema.field('data')
            self.assertTrue(
                pa_types.is_binary(mapped_data_field.type),
                f"After map_batches returning dict, data should be binary(), but got {mapped_data_field.type}"
            )
            break

        # Step 4: Write mapped dataset back via write_ray().
        write_builder2 = table.new_batch_write_builder()
        writer2 = write_builder2.new_write()

        writer2.write_ray(
            mapped_dataset,
            overwrite=False,
            concurrency=1
        )
        writer2.close()

        # Step 5: Verify the data was written correctly
        read_builder2 = table.new_read_builder()
        table_read2 = read_builder2.new_read()
        result = table_read2.to_arrow(read_builder2.new_scan().plan().splits())

        self.assertEqual(result.num_rows, 6, "Table should have 6 rows after roundtrip")

        result_df = result.to_pandas()
        result_df_sorted = result_df.sort_values(by='id').reset_index(drop=True)

        self.assertEqual(list(result_df_sorted['id']), [1, 1, 2, 2, 3, 3], "ID column should match")
        self.assertEqual(
            list(result_df_sorted['name']),
            ['Alice', 'Alice', 'Bob', 'Bob', 'Charlie', 'Charlie'],
            "Name column should match"
        )

        written_data_values = [bytes(d) if d is not None else None for d in result_df_sorted['data']]
        self.assertEqual(
            written_data_values,
            [b'blob_data_1', b'blob_data_1', b'blob_data_2', b'blob_data_2', b'blob_data_3', b'blob_data_3'],
            "Blob data column should match"
        )

    def test_to_ray_with_nested_projection(self):
        """to_ray() respects a nested-leaf projection.

        Sibling of the read_paimon() nested-projection test: this exercises
        the PreResolvedSplitProvider entry point (TableRead.to_ray), which
        must also forward nested_name_paths to the worker TableRead. Without
        it the worker treats the flattened leaf name as a missing top-level
        column and reads the projected leaf as NULL.
        """
        inner = pa.struct([('a', pa.int64()), ('b', pa.string())])
        pa_schema = pa.schema([('id', pa.int64()), ('payload', inner)])
        schema = Schema.from_pyarrow_schema(pa_schema)
        self.catalog.create_table('default.test_ray_nested_proj', schema, False)
        table = self.catalog.get_table('default.test_ray_nested_proj')

        write_builder = table.new_batch_write_builder()
        writer = write_builder.new_write()
        writer.write_arrow(pa.Table.from_pylist(
            [{'id': 1, 'payload': {'a': 10, 'b': 'x'}},
             {'id': 2, 'payload': {'a': 20, 'b': 'y'}}],
            schema=pa_schema))
        commit = write_builder.new_commit()
        commit.commit(writer.prepare_commit())
        writer.close()

        read_builder = table.new_read_builder().with_projection(['id', 'payload.a'])
        table_read = read_builder.new_read()
        splits = read_builder.new_scan().plan().splits()

        ray_dataset = table_read.to_ray(splits, override_num_blocks=1)
        rows = {r['id']: r for r in ray_dataset.take_all()}
        self.assertEqual(set(rows.keys()), {1, 2})
        self.assertEqual(rows[1]['payload_a'], 10)
        self.assertEqual(rows[2]['payload_a'], 20)


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