paimon-python/pypaimon/multimodal/connection.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.

from typing import Dict, Iterable, Optional

import pyarrow as pa

from pypaimon import CatalogFactory, Schema as PaimonSchema
from pypaimon.catalog.catalog import Catalog
from pypaimon.catalog.catalog_exception import (
    DatabaseNotExistException,
    TableAlreadyExistException,
    TableNotExistException,
)
from pypaimon.multimodal.table import MultimodalTable, _to_arrow_table

_DEFAULT_OPTIONS = {
    "row-tracking.enabled": "true",
    "data-evolution.enabled": "true",
    "blob-as-descriptor": "true",
    "global-index.search-mode": "full",
}

_DEFAULT_DATABASE = Catalog.DEFAULT_DATABASE


def connect(
        *,
        database: str = _DEFAULT_DATABASE,
        options: Optional[Dict[str, str]] = None):
    """Connect to a Paimon catalog through the multimodal facade."""
    return MultimodalConnection(
        _resolve_catalog(options),
        database=database,
    )


class MultimodalConnection:
    """High-level entry point backed by a Paimon catalog and database."""

    def __init__(self, catalog, database: str = _DEFAULT_DATABASE):
        if not database:
            raise ValueError("database is required.")
        self.catalog = catalog
        self.database = database

    def create_table(
            self,
            name: str,
            data=None,
            schema=None,
            options: Optional[Dict[str, str]] = None,
            partitioned: Optional[Iterable[str]] = None,
            ignore_if_exists: bool = False):
        """Create a multimodal table and optionally add initial data."""
        identifier = self._identifier(name)
        already_exists = _table_exists(self.catalog, identifier)
        if already_exists and ignore_if_exists:
            try:
                return self.get_table(name)
            except (DatabaseNotExistException, TableNotExistException):
                pass
        try:
            paimon_schema = _to_paimon_schema(
                schema, data, options, partitioned)
            _validate_multimodal_schema(paimon_schema, identifier)
        except ValueError:
            if ignore_if_exists:
                try:
                    return self.get_table(name)
                except (DatabaseNotExistException, TableNotExistException):
                    pass
            raise

        self._create_database_for(identifier)
        created = False
        try:
            self.catalog.create_table(
                identifier, paimon_schema, False)
            created = True
        except TableAlreadyExistException:
            if not ignore_if_exists:
                raise

        table = self.get_table(name)
        if data is not None and created:
            table.add(data)
        return table

    def get_table(self, name: str):
        identifier = self._identifier(name)
        raw_table = self.catalog.get_table(identifier)
        _validate_multimodal_table(raw_table, identifier)
        return MultimodalTable(
            self.catalog,
            identifier,
            raw_table,
        )

    def load_from_hdf5(
            self,
            table_name: str,
            paths,
            *,
            transform,
            source_options=None):
        """Load HDF5 transforms into an existing table as one append commit.

        Repeating a call appends the rows again. A commit exception has an
        unknown result and is not safe to retry without checking table state.
        Source filesystem options are isolated from the target warehouse.
        """
        from pypaimon.multimodal.hdf5 import load_from_hdf5
        return load_from_hdf5(
            self.get_table(table_name),
            paths,
            transform=transform,
            source_options=source_options,
        )

    def load_from_lerobot(
            self,
            table_name: str,
            source,
            *,
            batch_size: int = 1024,
            options=None,
            source_options=None,
            tag_name=None,
            parallelism=1,
            dataset_name=None,
            extra_dataset_fields=None,
            schema_evolution=True) -> None:
        """Import LeRobot v3 sources using table_name as the group prefix.

        Creates <table_name>_frames, <table_name>_episodes, and
        <table_name>_datasets.
        """
        from pypaimon.multimodal.lerobot import load_from_lerobot
        load_from_lerobot(
            self,
            table_name,
            source,
            batch_size=batch_size,
            options=options,
            source_options=source_options,
            tag_name=tag_name,
            parallelism=parallelism,
            dataset_name=dataset_name,
            extra_dataset_fields=extra_dataset_fields,
            schema_evolution=schema_evolution,
        )

    def append_to_lerobot(
            self,
            table_name: str,
            source,
            *,
            batch_size: int = 1024,
            source_options=None,
            tag_name=None,
            parallelism=1,
            dataset_name=None,
            extra_dataset_fields=None,
            schema_evolution=True) -> None:
        """Append LeRobot v3 datasets to a source-local three-table group."""
        from pypaimon.multimodal.lerobot import load_from_lerobot
        load_from_lerobot(
            self, table_name, source, batch_size=batch_size,
            source_options=source_options, tag_name=tag_name,
            parallelism=parallelism, dataset_name=dataset_name,
            extra_dataset_fields=extra_dataset_fields,
            schema_evolution=schema_evolution,
            append=True,
        )

    def recover_lerobot_append(self, table_name: str):
        """Recover a three-table group after an interrupted append."""
        from pypaimon.multimodal.lerobot.aggregate import _recover_incomplete_append
        return _recover_incomplete_append(self, table_name)

    def create_lerobot_tag(self, table_name: str, tag_name: str):
        """Pin a LeRobot group by prefix or frame table name.

        Pause group writes until this returns. Use the tag only after success
        and retain it on every component for the lifetime of a training run.
        """
        from pypaimon.multimodal.lerobot.metadata import create_lerobot_tag
        return create_lerobot_tag(self, table_name, tag_name)

    def load_from_rosbag(
            self,
            table_name: str,
            paths,
            *,
            transform,
            default_typestore=None,
            typestore_factory=None,
            source_options=None,
            staging=None,
            allow_storage_fragment: bool = False):
        """Validate and append ROS1/ROS2 transforms in one commit."""
        from pypaimon.multimodal.rosbag import load_from_rosbag
        return load_from_rosbag(
            self.get_table(table_name),
            paths,
            transform=transform,
            default_typestore=default_typestore,
            typestore_factory=typestore_factory,
            source_options=source_options,
            staging=staging,
            allow_storage_fragment=allow_storage_fragment,
        )

    def drop_table(self, name: str, ignore_if_not_exists: bool = False):
        self.catalog.drop_table(
            self._identifier(name),
            ignore_if_not_exists=ignore_if_not_exists,
        )

    def _identifier(self, name: str) -> str:
        return _resolve_identifier(name, self.database)

    def _create_database_for(self, identifier: str):
        database_name = _database_name(identifier)
        if database_name is not None:
            self.catalog.create_database(database_name, ignore_if_exists=True)


def _resolve_catalog(options):
    resolved_options = {}
    if options:
        resolved_options.update(options)
    if not resolved_options:
        raise ValueError("options is required.")
    return CatalogFactory.create(resolved_options)


def _table_exists(catalog, identifier: str) -> bool:
    try:
        catalog.get_table(identifier)
        return True
    except (DatabaseNotExistException, TableNotExistException):
        return False


def _validate_multimodal_table(table, identifier: str):
    _validate_multimodal_schema(table.table_schema, identifier)


def _validate_multimodal_schema(table_schema, identifier: str):
    options = table_schema.options
    if str(options.get("data-evolution.enabled", "false")).lower() != "true":
        raise ValueError(
            "Table %s is not a data-evolution table; "
            "data-evolution.enabled must be true." % identifier)
    if table_schema.primary_keys:
        raise ValueError(
            "Table %s has primary keys %s; multimodal tables must not have "
            "primary keys." % (identifier, table_schema.primary_keys))


def _resolve_identifier(name: str, database: str) -> str:
    if not name:
        raise ValueError("table name is required.")
    if "." in name:
        return name
    return "%s.%s" % (database, name)


def _to_paimon_schema(schema, data, options, partitioned):
    if schema is None:
        if data is None:
            raise ValueError("schema or data is required.")
        pa_schema, inferred_options = _infer_pyarrow_schema(data)
    else:
        pa_schema = _to_pyarrow_schema(schema)
        inferred_options = {}
    return PaimonSchema.from_pyarrow_schema(
        pa_schema,
        partition_keys=_normalize_partitioned(partitioned),
        options=_merge_options(inferred_options, options),
    )


def _merge_options(*option_groups):
    table_options = dict(_DEFAULT_OPTIONS)
    for options in option_groups:
        if options:
            table_options.update({str(k): str(v) for k, v in options.items()})
    return table_options


def _normalize_partitioned(partitioned):
    if partitioned is None:
        return []
    if isinstance(partitioned, str):
        return [partitioned]
    return list(partitioned)


def _to_pyarrow_schema(schema):
    if isinstance(schema, pa.Schema):
        return schema
    raise ValueError("schema must be a pyarrow.Schema.")


def _infer_pyarrow_schema(data):
    table = _to_arrow_table(data)
    if table.num_columns == 0:
        raise ValueError("Cannot infer schema from empty data.")
    return table.schema, {}


def _database_name(identifier: str):
    parts = identifier.split(".")
    if len(parts) == 2 and parts[0]:
        return parts[0]
    return None