| # 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() |
