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

"""LeRobot source resolution for local, Hub, and FileIO datasets."""

from array import array
import json
import logging
import posixpath
from bisect import bisect_right
from contextlib import closing, contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from urllib.parse import quote, unquote, urlparse, urlunparse

import pyarrow as pa
import pyarrow.fs as pafs
import pyarrow.parquet as pq

from pypaimon.common.options import Options
from pypaimon.common.uri_reader import FileUriReader
from pypaimon.filesystem.pyarrow_file_io import LegacyOssDirectoryListingError
from pypaimon.multimodal.lerobot.schema import _video_feature_names
from pypaimon.multimodal.source_utils import (
    _SourceFileIO,
    _normalize_source_path,
    _qualified_status_path,
)
from pypaimon.multimodal.lerobot.loader import _encode_media_frame


_LOGGER = logging.getLogger(__name__)


@dataclass(frozen=True)
class _LeRobotSource:
    path: str
    root: Optional[Path]
    repo_id: str
    file_io: object = None


@dataclass(frozen=True)
class _RemoteLeRobotMeta:
    info: dict
    episodes: object
    tasks: list


class _RemoteEpisodeIndex:

    def __init__(self, video_fields=()):
        self.starts = array("q")
        self._ends = array("q")
        self._chunk_indices = array("q")
        self._file_indices = array("q")
        self._video_fields = tuple(video_fields)
        self._video_indices = {
            field: (array("q"), array("q"))
            for field in self._video_fields
        }
        self._video_timestamps = {
            field: (array("d"), array("d"))
            for field in self._video_fields
        }

    def append(
            self, begin, end, chunk_index, file_index,
            video_values=None):
        self.starts.append(begin)
        self._ends.append(end)
        self._chunk_indices.append(chunk_index)
        self._file_indices.append(file_index)
        video_values = video_values or {}
        for field in self._video_fields:
            values = video_values[field]
            chunk_indices, file_indices = self._video_indices[field]
            from_timestamps, to_timestamps = self._video_timestamps[field]
            chunk_indices.append(values[0])
            file_indices.append(values[1])
            from_timestamps.append(values[2])
            to_timestamps.append(values[3])

    def __len__(self):
        return len(self.starts)

    def __getitem__(self, index):
        if index < 0:
            index += len(self)
        if index < 0 or index >= len(self):
            raise IndexError(index)
        result = {
            "episode_index": index,
            "dataset_from_index": self.starts[index],
            "dataset_to_index": self._ends[index],
            "length": self._ends[index] - self.starts[index],
            "data/chunk_index": self._chunk_indices[index],
            "data/file_index": self._file_indices[index],
        }
        for field in self._video_fields:
            chunk_indices, file_indices = self._video_indices[field]
            from_timestamps, to_timestamps = self._video_timestamps[field]
            prefix = "videos/%s/" % field
            result[prefix + "chunk_index"] = chunk_indices[index]
            result[prefix + "file_index"] = file_indices[index]
            result[prefix + "from_timestamp"] = from_timestamps[index]
            result[prefix + "to_timestamp"] = to_timestamps[index]
        return result


@contextmanager
def _resolved_source(source, source_options):
    if isinstance(source, Path):
        root = source.expanduser().resolve()
        if not root.is_dir():
            raise FileNotFoundError(
                "LeRobot source directory does not exist: %s" % root)
        yield _local_source(root)
        return
    if not isinstance(source, str) or not source.strip():
        raise ValueError(
            "source must be a local directory or Hugging Face repo_id.")

    value = source.strip()
    candidate = Path(value).expanduser()
    if candidate.is_dir():
        yield _local_source(candidate.resolve())
        return
    if candidate.is_absolute() or value.startswith((".", "~")):
        raise FileNotFoundError(
            "LeRobot source directory does not exist: %s" % candidate)
    if "://" not in value:
        yield _LeRobotSource(path=value, root=None, repo_id=value), None
        return

    source_uri = _normalize_source_path(value).rstrip("/")
    source_file_io = _SourceFileIO(Options(source_options))
    try:
        try:
            status = source_file_io.get_file_status(source_uri)
        except FileNotFoundError as error:
            raise FileNotFoundError(
                "LeRobot source directory does not exist: %s" % source_uri
            ) from error
        if status.type != pafs.FileType.Directory:
            raise ValueError(
                "LeRobot URI source must be a directory: %s" % source_uri)
        info = _read_remote_json(
            source_file_io, _remote_path(source_uri, "meta/info.json"))
        yield (
            _LeRobotSource(
                path=source_uri,
                root=None,
                repo_id="",
                file_io=source_file_io,
            ),
            info,
        )
    finally:
        _close_quietly(source_file_io, "source FileIO")


def _close_quietly(resource, name):
    try:
        resource.close()
    except Exception:
        _LOGGER.warning("Failed to close LeRobot %s.", name, exc_info=True)


def _local_source(root, display_path=None):
    info_path = root / "meta" / "info.json"
    if not info_path.is_file():
        raise ValueError(
            "LeRobot source is missing meta/info.json: %s"
            % (display_path or root))
    try:
        with info_path.open("r", encoding="utf-8") as file:
            info = json.load(file)
    except (OSError, ValueError) as error:
        raise ValueError(
            "Cannot read LeRobot metadata %s: %s"
            % (info_path, error)) from error
    return (
        _LeRobotSource(
            path=str(display_path or root),
            root=root,
            repo_id="local/pypaimon-import",
        ),
        info,
    )


def _import_lerobot_dataset():
    try:
        from lerobot.datasets.lerobot_dataset import LeRobotDataset
    except ImportError as error:
        raise ImportError(
            "load_from_lerobot requires LeRobot; install "
            "'pypaimon[lerobot]'.") from error
    return LeRobotDataset


def _info_dict(info):
    # LeRobot 0.6 exposes DatasetInfo instead of a dict.
    return dict(info.to_dict() if hasattr(info, "to_dict") else info)


def _load_hub_info(source):
    try:
        from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
    except ImportError as error:
        raise ImportError(
            "load_from_lerobot requires LeRobot; install "
            "'pypaimon[lerobot]'.") from error
    try:
        metadata = LeRobotDatasetMetadata(repo_id=source.repo_id)
        # DatasetInfo omits unknown fields; preserve the downloaded source JSON.
        return _local_source(Path(metadata.root), source.path)[1]
    except Exception as error:
        raise ValueError(
            "Cannot open LeRobot Dataset v3 metadata %s: %s"
            % (source.path, error)) from error


def _open_dataset(LeRobotDataset, source, download_videos=False):
    try:
        if source.root is not None:
            return LeRobotDataset(
                repo_id=source.repo_id,
                root=source.root,
                download_videos=download_videos,
            )
        return LeRobotDataset(
            repo_id=source.repo_id,
            download_videos=download_videos,
        )
    except Exception as error:
        raise ValueError(
            "Cannot open LeRobot Dataset v3 source %s: %s"
            % (source.path, error)) from error


def _open_resolved_dataset(
        LeRobotDataset, source, info, download_videos=False):
    if source.file_io is not None:
        return _RemoteLeRobotDataset(source, info)
    return _open_dataset(LeRobotDataset, source, download_videos)


class _RemoteLeRobotDataset:

    _EPISODE_COLUMNS = [
        "episode_index",
        "dataset_from_index",
        "dataset_to_index",
        "data/chunk_index",
        "data/file_index",
    ]

    def __init__(self, source, info):
        self.source = source
        self.root = source.path
        self._file_io = source.file_io
        self._episodes = self._load_episodes(info)
        self._tasks = self._load_tasks(info)
        self.meta = _RemoteLeRobotMeta(info, self._episodes, self._tasks)
        self._episode_starts = self._episodes.starts
        self._data_ranges = self._build_data_ranges(info)
        self._cached_data_path = None
        self._cached_data_table = None

    def __len__(self):
        return int(self.meta.info.get("total_frames", 0))

    def close(self):
        self._cached_data_table = None

    def read_batch(self, begin, end):
        episode = self._episode_for_range(begin, end)
        relative_path = self._data_path(episode, self.meta.info)
        source_path = _remote_source_path(
            self.source.path,
            relative_path,
            "info.data_path",
            self._file_io,
        )
        if source_path != self._cached_data_path:
            table = _read_remote_parquet(self._file_io, source_path)
            expected_begin, expected_end = self._data_ranges[relative_path]
            expected_rows = expected_end - expected_begin
            if table.num_rows != expected_rows:
                raise ValueError(
                    "LeRobot data file %s has %d rows; metadata expects %d."
                    % (source_path, table.num_rows, expected_rows))
            self._cached_data_path = source_path
            self._cached_data_table = table
        file_begin = self._data_ranges[relative_path][0]
        return self._cached_data_table.slice(begin - file_begin, end - begin)

    def image_bytes(self, value):
        if value is None:
            raise ValueError("LeRobot image feature contains a null frame.")
        if isinstance(value, (bytes, bytearray, memoryview)):
            return bytes(value)
        if isinstance(value, dict):
            body = value.get("bytes")
            if body is not None:
                return bytes(body)
            image_path = value.get("path")
            if image_path:
                source_path = _remote_source_path(
                    self.source.path,
                    image_path,
                    "image path",
                    self._file_io,
                )
                return _read_remote_bytes(self._file_io, source_path)
        return _encode_media_frame(value)

    def video_source(self, video_key, episode):
        relative_path = self.meta.info["video_path"].format(
            video_key=video_key,
            chunk_index=int(episode[
                "videos/%s/chunk_index" % video_key]),
            file_index=int(episode[
                "videos/%s/file_index" % video_key]),
        )
        relative_path = _relative_dataset_path(
            relative_path, "info.video_path")
        path = _remote_source_path(
            self.source.path,
            relative_path,
            "info.video_path",
            self._file_io,
        )
        status = self._file_io.get_file_status(path)
        if status.type != pafs.FileType.File \
                or status.size is None or status.size <= 0:
            raise ValueError("LeRobot video file is empty: %s" % path)
        return path, int(status.size)

    @property
    def video_uri_reader_factory(self):
        return _SourceUriReaderFactory(self._file_io)

    def _load_episodes(self, info):
        episode_count = int(info.get("total_episodes", 0))
        if episode_count == 0:
            return []
        directory = _remote_path(self.source.path, "meta/episodes")
        paths = _remote_parquet_files(self._file_io, directory)
        video_fields = _video_feature_names(info)
        selected_columns = list(self._EPISODE_COLUMNS)
        for field in video_fields:
            selected_columns.extend([
                "videos/%s/%s" % (field, suffix)
                for suffix in (
                    "chunk_index",
                    "file_index",
                    "from_timestamp",
                    "to_timestamp",
                )
            ])
        episodes = _RemoteEpisodeIndex(video_fields)
        for path in paths:
            table = _read_remote_parquet(
                self._file_io,
                path,
                columns=selected_columns,
            )
            control_columns = {
                name: table.column(name)
                for name in self._EPISODE_COLUMNS
            }
            for offset in range(table.num_rows):
                episode_index = int(
                    control_columns["episode_index"][offset].as_py())
                if episode_index != len(episodes):
                    raise ValueError(
                        "LeRobot Episode metadata must be ordered by "
                        "episode_index.")
                episodes.append(
                    int(control_columns[
                        "dataset_from_index"][offset].as_py()),
                    int(control_columns[
                        "dataset_to_index"][offset].as_py()),
                    int(control_columns[
                        "data/chunk_index"][offset].as_py()),
                    int(control_columns[
                        "data/file_index"][offset].as_py()),
                    {
                        field: (
                            int(table.column(
                                "videos/%s/chunk_index" % field
                            )[offset].as_py()),
                            int(table.column(
                                "videos/%s/file_index" % field
                            )[offset].as_py()),
                            float(table.column(
                                "videos/%s/from_timestamp" % field
                            )[offset].as_py()),
                            float(table.column(
                                "videos/%s/to_timestamp" % field
                            )[offset].as_py()),
                        )
                        for field in video_fields
                    },
                )
        if len(episodes) != episode_count:
            raise ValueError(
                "LeRobot metadata reports %d Episodes but %d were found."
                % (episode_count, len(episodes)))
        return episodes

    def _load_tasks(self, info):
        task_count = int(info.get("total_tasks", 0))
        if task_count == 0:
            return []
        path = _remote_path(self.source.path, "meta/tasks.parquet")
        table = _read_remote_parquet(self._file_io, path)
        name_column = _pandas_index_column(table.schema, "task")
        rows = table.select(["task_index", name_column]).to_pylist()
        tasks = [None] * task_count
        for row in rows:
            index = int(row["task_index"])
            name = row[name_column]
            if index < 0 or index >= task_count \
                    or not isinstance(name, str) or not name:
                raise ValueError("LeRobot task metadata is invalid: %s" % row)
            tasks[index] = name
        if any(task is None for task in tasks):
            raise ValueError(
                "LeRobot metadata reports %d tasks but %d were found."
                % (task_count, len(rows)))
        return tasks

    def _build_data_ranges(self, info):
        ranges = {}
        for episode in self._episodes:
            path = self._data_path(episode, info)
            begin = int(episode["dataset_from_index"])
            end = int(episode["dataset_to_index"])
            if path in ranges:
                previous_begin, previous_end = ranges[path]
                if begin != previous_end:
                    raise ValueError(
                        "LeRobot data file %s has non-contiguous Episode "
                        "ranges." % path)
                ranges[path] = previous_begin, end
            else:
                ranges[path] = begin, end
        return ranges

    def _episode_for_range(self, begin, end):
        index = bisect_right(self._episode_starts, begin) - 1
        if index < 0:
            raise ValueError("LeRobot frame range starts before Episode 0.")
        episode = self._episodes[index]
        episode_end = int(episode["dataset_to_index"])
        if end > episode_end:
            raise ValueError("LeRobot frame batch crosses an Episode boundary.")
        return episode

    @staticmethod
    def _data_path(episode, info):
        path = info["data_path"].format(
            chunk_index=int(episode["data/chunk_index"]),
            file_index=int(episode["data/file_index"]),
        )
        return _relative_dataset_path(path, "info.data_path")


def _remote_path(root, relative_path):
    return "%s/%s" % (root.rstrip("/"), relative_path.lstrip("/"))


class _SourceUriReaderFactory:

    def __init__(self, file_io):
        self._reader = FileUriReader(file_io)

    def create(self, unused_uri):
        return self._reader


def _pandas_index_column(schema, component):
    encoded = (schema.metadata or {}).get(b"pandas")
    if encoded is None:
        # Arrow-written task tables may omit Pandas index metadata.
        expected = {component + "_index", component}
        if len(schema.names) == 2 and set(schema.names) == expected \
                and (pa.types.is_string(schema.field(component).type)
                     or pa.types.is_large_string(schema.field(component).type)):
            return component
        raise ValueError(
            "LeRobot %s metadata must contain a Pandas index or the "
            "canonical %s and %s columns."
            % (component, component + "_index", component))
    try:
        pandas_metadata = json.loads(encoded.decode("utf-8"))
        index_columns = pandas_metadata["index_columns"]
    except (AttributeError, KeyError, TypeError, ValueError,
            UnicodeDecodeError) as error:
        raise ValueError(
            "LeRobot %s metadata must contain a Pandas index."
            % component
        ) from error
    if not isinstance(index_columns, list) \
            or len(index_columns) != 1 \
            or not isinstance(index_columns[0], str) \
            or index_columns[0] not in schema.names \
            or index_columns[0] == component + "_index":
        raise ValueError(
            "LeRobot %s metadata must contain one text Pandas index."
            % component)
    return index_columns[0]


def _relative_dataset_path(path, name):
    if not isinstance(path, str) or not path:
        raise ValueError("LeRobot %s must be a relative path." % name)
    path = _fully_unquote(path)
    if "\\" in path or urlparse(path).scheme or path.startswith(("/", "~")):
        raise ValueError(
            "LeRobot %s must stay within the source directory: %s"
            % (name, path))
    normalized = posixpath.normpath(path)
    if normalized in (".", "..") or normalized.startswith("../"):
        raise ValueError(
            "LeRobot %s must stay within the source directory: %s"
            % (name, path))
    return normalized


def _remote_source_path(root, path, name, source_file_io=None):
    if not isinstance(path, str) or not path:
        _relative_dataset_path(path, name)
    root_uri = urlparse(root)
    path_uri = urlparse(path)
    root_path = posixpath.normpath(_fully_unquote(root_uri.path) or "/")
    if path_uri.scheme:
        if path_uri.query or path_uri.fragment \
                or path_uri.scheme.lower() != root_uri.scheme.lower() \
                or path_uri.netloc != root_uri.netloc:
            raise ValueError(
                "LeRobot %s must stay within the source directory: %s"
                % (name, path))
        source_path = posixpath.normpath(
            _fully_unquote(path_uri.path) or "/")
    elif root_uri.scheme.lower() == "file" and Path(path).is_absolute():
        source_path = str(Path(path).resolve())
    else:
        relative_path = _relative_dataset_path(path, name)
        source_path = posixpath.normpath(posixpath.join(
            root_path,
            relative_path,
        ))
    if posixpath.commonpath([root_path, source_path]) != root_path:
        raise ValueError(
            "LeRobot %s must stay within the source directory: %s"
            % (name, path))
    source_uri = urlunparse((
        root_uri.scheme,
        root_uri.netloc,
        quote(source_path, safe="/:"),
        "",
        "",
        "",
    ))
    _validate_filesystem_containment(
        source_file_io, root, source_uri, name, path)
    return source_uri


def _fully_unquote(path):
    decoded = unquote(path)
    while decoded != path:
        path = decoded
        decoded = unquote(path)
    return decoded


def _validate_filesystem_containment(
        source_file_io, root, source, name, original_path):
    to_filesystem_path = getattr(
        source_file_io, "to_filesystem_path", None)
    if not callable(to_filesystem_path):
        return
    root_path = _fully_unquote(to_filesystem_path(root))
    source_path = _fully_unquote(to_filesystem_path(source))
    if urlparse(root).scheme.lower() == "file":
        root_path = Path(root_path).resolve()
        source_path = Path(source_path).resolve()
        try:
            source_path.relative_to(root_path)
        except ValueError as error:
            raise ValueError(
                "LeRobot %s must stay within the source directory: %s"
                % (name, original_path)) from error
        return
    else:
        root_path = posixpath.normpath(root_path)
        source_path = posixpath.normpath(source_path)
    relative_path = posixpath.relpath(source_path, root_path)
    if relative_path == ".." or relative_path.startswith("../"):
        raise ValueError(
            "LeRobot %s must stay within the source directory: %s"
            % (name, original_path))


def _validate_info_paths(info):
    for name in ("data_path", "video_path"):
        path = info.get(name)
        if path is not None:
            _relative_dataset_path(path, "info.%s" % name)


def _read_remote_bytes(source_file_io, path):
    stream = source_file_io.new_input_stream(path)
    with closing(stream) as source_stream:
        return source_stream.read()


def _read_remote_json(source_file_io, path):
    try:
        return json.loads(_read_remote_bytes(
            source_file_io, path).decode("utf-8"))
    except (OSError, UnicodeError, ValueError) as error:
        raise ValueError(
            "Cannot read LeRobot metadata %s: %s" % (path, error)) from error


def _read_remote_parquet(source_file_io, path, columns=None):
    stream = source_file_io.new_input_stream(path)
    with closing(stream) as source_stream:
        try:
            return pq.read_table(source_stream, columns=columns)
        except (OSError, ValueError, pa.ArrowException) as error:
            raise ValueError(
                "Cannot read LeRobot Parquet file %s: %s"
                % (path, error)) from error


def _read_remote_parquet_schema(source_file_io, path):
    stream = source_file_io.new_input_stream(path)
    with closing(stream) as source_stream:
        try:
            return pq.read_schema(source_stream)
        except (OSError, ValueError, pa.ArrowException) as error:
            raise ValueError(
                "Cannot read LeRobot Parquet schema %s: %s"
                % (path, error)) from error


def _remote_parquet_files(source_file_io, directory):
    try:
        statuses = source_file_io.list_status(directory)
    except LegacyOssDirectoryListingError as error:
        raise ValueError(
            "LeRobot URI directory listing is unavailable at %s; use "
            "Jindo or upgrade PyArrow." % directory) from error
    paths = []
    for status in statuses:
        path = _qualified_status_path(directory, status)
        if status.type == pafs.FileType.Directory:
            paths.extend(_remote_parquet_files(source_file_io, path))
        elif status.type == pafs.FileType.File and path.endswith(".parquet"):
            paths.append(path)
    return sorted(paths)