paimon-python/pypaimon/tests/rest/rest_object_table_test.py

# 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 pyarrow as pa

from pypaimon import PaimonVirtualFileSystem, Schema
from pypaimon.table.object import ObjectTable
from pypaimon.tests.rest.rest_base_test import RESTBaseTest


class RESTObjectTableTest(RESTBaseTest):

    def setUp(self):
        super().setUp()
        pvfs_options = {
            'uri': self.options['uri'],
            'warehouse': self.options['warehouse'],
            'dlf.region': self.options['dlf.region'],
            'token.provider': self.options['token.provider'],
            'token': self.options['token'],
        }
        self.pvfs = PaimonVirtualFileSystem(pvfs_options)

    def _create_object_table(self, table_name, extra_options=None):
        """Helper to create an object table with the given name.

        ObjectTable has a fixed schema and does not support custom fields.
        Only options (including type=object-table) are needed.
        """
        options = {"type": "object-table"}
        if extra_options:
            options.update(extra_options)
        schema = Schema(options=options)
        self.rest_catalog.drop_table(table_name, True)
        self.rest_catalog.create_table(table_name, schema, False)
        return self.rest_catalog.get_table(table_name)

    def _pvfs_table_path(self, table_name, sub_path=None):
        """Build a pvfs:// path for the given table and optional sub-path."""
        warehouse = self.options['warehouse']
        base = "pvfs://{}/{}".format(warehouse, table_name.replace('.', '/'))
        if sub_path:
            return "{}/{}".format(base, sub_path)
        return base

    def _write_file_via_pvfs(self, table_name, filename, content):
        """Write a file into the table's location using pvfs."""
        path = self._pvfs_table_path(table_name, filename)
        # Ensure parent directory exists when filename contains subdirectories
        if '/' in filename:
            parent_dir = self._pvfs_table_path(
                table_name, filename.rsplit('/', 1)[0]
            )
            self.pvfs.makedirs(parent_dir, exist_ok=True)
        with self.pvfs.open(path, 'wb') as f:
            f.write(content)

    def test_get_object_table(self):
        table_name = "default.object_table_basic"
        table = self._create_object_table(table_name)

        self.assertIsInstance(table, ObjectTable)
        self.assertEqual(table.name(), "object_table_basic")
        self.assertEqual(table.full_name(), table_name)
        self.assertEqual(table.partition_keys, [])
        self.assertEqual(table.primary_keys, [])
        self.assertEqual(table.options().get("type"), "object-table")

    def test_object_table_read_files(self):
        table_name = "default.object_table_read"
        table = self._create_object_table(table_name)

        # Write some test files into the table's location via pvfs
        test_files = {
            "file_a.txt": b"hello world",
            "file_b.dat": b"some binary data here",
        }
        for filename, content in test_files.items():
            self._write_file_via_pvfs(table_name, filename, content)

        # Read the object table
        read_builder = table.new_read_builder()
        scan = read_builder.new_scan()
        plan = scan.plan()
        splits = plan.splits()

        self.assertEqual(len(splits), 1)

        table_read = read_builder.new_read()
        result = table_read.to_arrow(splits)

        self.assertIsInstance(result, pa.Table)
        # Should contain the schema, snapshot and manifest files plus our test files
        # Filter to only our test files
        result_names = result.column("name").to_pylist()
        for filename in test_files:
            self.assertIn(filename, result_names)

        # Verify file sizes for our test files
        result_dict = {}
        for i in range(result.num_rows):
            name = result.column("name")[i].as_py()
            length = result.column("length")[i].as_py()
            result_dict[name] = length

        for filename, content in test_files.items():
            self.assertEqual(result_dict[filename], len(content))

        # Verify schema columns
        self.assertIn("path", result.column_names)
        self.assertIn("name", result.column_names)
        self.assertIn("length", result.column_names)
        self.assertIn("mtime", result.column_names)
        self.assertIn("atime", result.column_names)
        self.assertIn("owner", result.column_names)

    def test_object_table_read_with_subdirectories(self):
        table_name = "default.object_table_subdir"
        table = self._create_object_table(table_name)

        # Write files via pvfs, including a nested subdirectory file
        self._write_file_via_pvfs(table_name, "root_file.txt", b"root content")
        self._write_file_via_pvfs(table_name, "subdir/nested_file.txt", b"nested content")

        read_builder = table.new_read_builder()
        splits = read_builder.new_scan().plan().splits()
        result = read_builder.new_read().to_arrow(splits)

        result_names = result.column("name").to_pylist()
        self.assertIn("root_file.txt", result_names)
        self.assertIn("nested_file.txt", result_names)

        # Verify relative paths contain subdirectory
        result_paths = result.column("path").to_pylist()
        nested_paths = [p for p in result_paths if "nested_file.txt" in p]
        self.assertTrue(len(nested_paths) > 0)
        self.assertIn("subdir/", nested_paths[0])

    def test_object_table_with_projection(self):
        table_name = "default.object_table_projection"
        table = self._create_object_table(table_name)

        self._write_file_via_pvfs(table_name, "proj_test.txt", b"test data")

        # Test projection with two columns
        read_builder = table.new_read_builder()
        read_builder.with_projection(["name", "length"])
        splits = read_builder.new_scan().plan().splits()
        result = read_builder.new_read().to_arrow(splits)

        self.assertEqual(result.column_names, ["name", "length"])
        self.assertNotIn("path", result.column_names)
        self.assertNotIn("mtime", result.column_names)
        self.assertNotIn("atime", result.column_names)
        self.assertNotIn("owner", result.column_names)

        # Verify data content for our test file
        result_names = result.column("name").to_pylist()
        self.assertIn("proj_test.txt", result_names)
        idx = result_names.index("proj_test.txt")
        self.assertEqual(result.column("length")[idx].as_py(), len(b"test data"))

    def test_object_table_with_limit(self):
        table_name = "default.object_table_limit"
        table = self._create_object_table(table_name)

        for i in range(5):
            self._write_file_via_pvfs(
                table_name, "file_{}.txt".format(i), "content {}".format(i).encode()
            )

        read_builder = table.new_read_builder()
        read_builder.with_limit(2)
        splits = read_builder.new_scan().plan().splits()
        result = read_builder.new_read().to_arrow(splits)

        self.assertEqual(result.num_rows, 2)

    def test_object_table_options_and_copy(self):
        table_name = "default.object_table_options"
        table = self._create_object_table(
            table_name, extra_options={"custom.key": "custom_value"}
        )

        self.assertEqual(table.options().get("custom.key"), "custom_value")
        self.assertEqual(table.options().get("type"), "object-table")

        # Test copy with dynamic options
        copied = table.copy({"new.key": "new_value"})
        self.assertIsInstance(copied, ObjectTable)
        self.assertEqual(copied.options().get("custom.key"), "custom_value")
        self.assertEqual(copied.options().get("new.key"), "new_value")

        # Original should not be modified
        self.assertIsNone(table.options().get("new.key"))

        # Test copy with override
        overridden = table.copy({"custom.key": "overridden"})
        self.assertEqual(overridden.options().get("custom.key"), "overridden")
        self.assertEqual(table.options().get("custom.key"), "custom_value")

    def test_object_table_unsupported_write(self):
        table_name = "default.object_table_no_write"
        table = self._create_object_table(table_name)

        with self.assertRaises(NotImplementedError):
            table.new_batch_write_builder()
        with self.assertRaises(NotImplementedError):
            table.new_stream_write_builder()

    def test_object_table_unsupported_drop_partitions(self):
        table_name = "default.object_table_no_drop_partitions"
        self._create_object_table(table_name)

        with self.assertRaisesRegex(
            ValueError,
            "drop_partitions is not supported for table type 'ObjectTable'",
        ):
            self.rest_catalog.drop_partitions(
                table_name,
                [{"dt": "20250101"}],
            )

    def test_object_table_to_pandas(self):
        table_name = "default.object_table_pandas"
        table = self._create_object_table(table_name)

        self._write_file_via_pvfs(table_name, "pandas_test.txt", b"pandas data")

        read_builder = table.new_read_builder()
        splits = read_builder.new_scan().plan().splits()
        result_df = read_builder.new_read().to_pandas(splits)

        self.assertIn("name", result_df.columns)
        self.assertIn("pandas_test.txt", result_df["name"].values)