[python] Add direct LeRobot capture writer

docs/docs/pypaimon/multimodal-api.mdx

@@ -661,6 +661,57 @@ The one-time importer requires a new target table.
 Scalars map to scalar types, vectors to `VECTOR`, higher-rank tensors to nested
 `ARRAY`, and images to `BLOB`. Images keep their compressed bytes.
 
+## Capture LeRobot frames directly into Paimon
+
+`PaimonLeRobotWriter` implements the write-side surface used by LeRobot's
+recording loop without first creating a LeRobot Parquet/image dataset. Pass the
+same user feature mapping that would be passed to `LeRobotDataset.create`.
+The writer adds the standard `timestamp`, `frame_index`, `episode_index`,
+`index`, and `task_index` features itself.
+
+```python
+from pypaimon.multimodal.lerobot import PaimonLeRobotWriter
+
+writer = PaimonLeRobotWriter(
+    conn,
+    "robot_data",
+    fps=30,
+    features=dataset_features,
+    episodes_per_commit=10,
+)
+
+# LeRobot's record_loop only needs writer.fps, writer.features, and
+# writer.add_frame(frame), so the writer can be passed as its dataset argument.
+record_loop(..., fps=30, dataset=writer)
+writer.save_episode()
+
+# Optional durability/visibility boundary before the configured batch fills.
+snapshot_id = writer.flush()
+writer.finalize()
+```
+
+`add_frame` accepts all declared user features plus a non-empty `task`. It
+validates each complete frame before buffering it and derives the standard
+LeRobot indices and timestamp. Raw NumPy, Torch, or PIL image frames are encoded
+as PNG bytes and stored in Paimon `BLOB` columns; no LeRobot data directory or
+MP4 is created. `video` features remain unsupported.
+
+`save_episode` accepts the current episode and writes it to a long-lived Paimon
+batch writer. By default, ten completed episodes share one Paimon commit. It
+returns `None` while the batch is still open and returns the committed snapshot
+ID when the threshold is reached. Set `episodes_per_commit=1` when every
+episode must become visible immediately, or call `flush()` at an operational
+boundary. `finalize()` commits the final partial batch.
+
+Call `clear_episode_buffer()` before `save_episode()` to discard a re-recorded
+episode without advancing frame or episode indices. `finalize()` rejects an
+unfinished episode instead of silently dropping its frames.
+
+The first version creates a new table or uses an existing empty compatible
+table. Resuming into a non-empty table is rejected because recovering global
+frame, episode, and task indices is not implemented yet. A commit exception has
+an unknown result and is not automatically retried.
+
 ## Overwrite
 
 `overwrite` accepts the same input formats as `add` and replaces existing data

paimon-python/pypaimon/multimodal/lerobot/__init__.py

@@ -14,11 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""One-time LeRobot Dataset v3 import into a multimodal Paimon table."""
+"""LeRobot Dataset v3 import and direct Paimon capture."""
 
 from pypaimon.multimodal.lerobot.api import load_from_lerobot
+from pypaimon.multimodal.lerobot.writer import PaimonLeRobotWriter
 
 
 __all__ = [
+    "PaimonLeRobotWriter",
     "load_from_lerobot",
 ]

paimon-python/pypaimon/multimodal/lerobot/loader.py

@@ -450,7 +450,7 @@ def _image_bytes(value, root):
     return _encode_media_frame(value)
 
 
-def _encode_media_frame(value):
+def _encode_media_frame(value, channel_first=None):
     try:
         import numpy as np
         from PIL import Image
@@ -466,7 +466,10 @@ def _encode_media_frame(value):
         if callable(detach):
             value = detach().cpu().numpy()
         array = np.asarray(value)
-        if array.ndim == 3 and array.shape[0] in (1, 3, 4):
+        if channel_first is True or (
+                channel_first is None
+                and array.ndim == 3
+                and array.shape[0] in (1, 3, 4)):
             array = np.transpose(array, (1, 2, 0))
         if np.issubdtype(array.dtype, np.floating):
             array = np.rint(np.clip(array, 0.0, 1.0) * 255.0).astype(np.uint8)

paimon-python/pypaimon/multimodal/lerobot/writer.py

@@ -0,0 +1,298 @@
+# 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-compatible capture writer for multimodal Paimon tables."""
+
+import copy
+from typing import Mapping, Optional
+
+import pyarrow as pa
+
+from pypaimon.multimodal.arrow_utils import strict_arrow_table
+from pypaimon.multimodal.hdf5 import _SnapshotRecorder
+from pypaimon.multimodal.lerobot.loader import (
+    _encode_media_frame,
+    _normalize_value,
+    _safe_array,
+    _value_shape,
+)
+from pypaimon.multimodal.lerobot.schema import (
+    _feature_shape,
+    _schema_from_info,
+    _validate_lerobot_schema,
+)
+from pypaimon.multimodal.table import _target_schema
+
+
+_DEFAULT_FEATURES = {
+    "timestamp": {"dtype": "float32", "shape": (1,), "names": None},
+    "frame_index": {"dtype": "int64", "shape": (1,), "names": None},
+    "episode_index": {"dtype": "int64", "shape": (1,), "names": None},
+    "index": {"dtype": "int64", "shape": (1,), "names": None},
+    "task_index": {"dtype": "int64", "shape": (1,), "names": None},
+}
+_TASK_FEATURE = {"dtype": "string", "shape": (1,), "names": None}
+
+
+class PaimonLeRobotWriter:
+    """Collect LeRobot frames and commit completed episodes to Paimon."""
+
+    def __init__(
+            self,
+            connection,
+            table_name: str,
+            *,
+            fps: int,
+            features: Mapping[str, Mapping[str, object]],
+            episodes_per_commit: int = 10,
+            options: Optional[Mapping[str, object]] = None):
+        if isinstance(fps, bool) or not isinstance(fps, int) or fps <= 0:
+            raise ValueError("fps must be a positive integer.")
+        if isinstance(episodes_per_commit, bool) \
+                or not isinstance(episodes_per_commit, int) \
+                or episodes_per_commit <= 0:
+            raise ValueError("episodes_per_commit must be a positive integer.")
+        if not isinstance(features, Mapping) or not features:
+            raise ValueError("features must be a non-empty mapping.")
+        if "task" in features:
+            raise ValueError("task is managed by PaimonLeRobotWriter.")
+
+        self.fps = fps
+        self.episodes_per_commit = episodes_per_commit
+        self.features = copy.deepcopy(dict(features))
+        self._user_features = copy.deepcopy(dict(features))
+        self.features.update(copy.deepcopy(_DEFAULT_FEATURES))
+        schema_features = dict(self.features)
+        schema_features["task"] = _TASK_FEATURE
+        self._source_schema = _schema_from_info({"features": schema_features})
+        self._table = connection.create_table(
+            table_name,
+            schema=self._source_schema,
+            options=options,
+            ignore_if_exists=True,
+        )
+        self._target_schema = _target_schema(self._table.raw_table)
+        _validate_lerobot_schema(
+            self._source_schema, self._target_schema, table_name)
+        strict_arrow_table(
+            pa.Table.from_batches([], schema=self._source_schema),
+            self._target_schema,
+            table_name,
+            0,
+            "LeRobot",
+        )
+        if self._table.raw_table.snapshot_manager().get_latest_snapshot() \
+                is not None:
+            raise ValueError(
+                "PaimonLeRobotWriter does not support resuming a non-empty "
+                "table yet.")
+
+        self.num_frames = 0
+        self.num_episodes = 0
+        self.pending_episodes = 0
+        self._episode_frames = []
+        self._task_indices = {}
+        self._table_write = None
+        self._table_commit = None
+        self._snapshot_recorder = None
+        self._finalized = False
+        self._failed = False
+
+    def add_frame(self, frame):
+        self._require_open("add_frame")
+        if not isinstance(frame, Mapping):
+            raise ValueError("frame must be a mapping.")
+        expected = set(self._user_features)
+        actual = set(frame) - {"task"}
+        if actual != expected:
+            missing = sorted(expected - actual)
+            extra = sorted(actual - expected)
+            raise ValueError(
+                "LeRobot frame fields do not match features; missing=%s, "
+                "extra=%s." % (missing, extra))
+        task = frame.get("task")
+        if not isinstance(task, str) or not task:
+            raise ValueError("LeRobot frame task must be a non-empty string.")
+
+        values = {"task": task}
+        for name, feature in self._user_features.items():
+            if feature.get("dtype") == "image":
+                value = self._image_bytes(frame[name], feature, name)
+            else:
+                value = _normalize_value(frame[name], feature, name)
+            _safe_array(
+                [value],
+                self._source_schema.field(name),
+                name,
+                str(feature.get("dtype", "")),
+            )
+            values[name] = value
+        self._episode_frames.append(values)
+
+    @staticmethod
+    def _image_bytes(value, feature, name):
+        actual_shape = _value_shape(value)
+        getbands = getattr(value, "getbands", None)
+        image_size = getattr(value, "size", None)
+        if not actual_shape and callable(getbands) \
+                and isinstance(image_size, tuple) and len(image_size) == 2:
+            actual_shape = (image_size[1], image_size[0], len(getbands()))
+        expected_shape = _feature_shape(feature, name)
+        channel_first_shape = ()
+        if len(actual_shape) == 3:
+            channel_first_shape = (
+                actual_shape[1], actual_shape[2], actual_shape[0])
+        if actual_shape and actual_shape != expected_shape \
+                and channel_first_shape != expected_shape:
+            raise ValueError(
+                "LeRobot feature %s expected shape %s, got %s."
+                % (name, expected_shape, actual_shape))
+        return _encode_media_frame(
+            value,
+            channel_first=actual_shape != expected_shape
+            and channel_first_shape == expected_shape,
+        )
+
+    def save_episode(self):
+        self._require_open("save_episode")
+        if not self._episode_frames:
+            raise ValueError("Cannot save an empty LeRobot episode.")
+
+        episode = self._episode_table()
+        try:
+            self._ensure_batch()
+            self._table_write.write_arrow(episode)
+        except BaseException:
+            self._fail_batch(abort=True)
+            raise
+
+        self.num_frames += episode.num_rows
+        self.num_episodes += 1
+        self.pending_episodes += 1
+        self._episode_frames = []
+        if self.pending_episodes >= self.episodes_per_commit:
+            return self.flush()
+        return None
+
+    def clear_episode_buffer(self, delete_images=True):
+        self._require_open("clear_episode_buffer")
+        self._episode_frames = []
+
+    def has_pending_frames(self):
+        return bool(self._episode_frames)
+
+    def flush(self):
+        self._require_open("flush")
+        if self.pending_episodes == 0:
+            return None
+        commit_started = False
+        try:
+            messages = self._table_write.prepare_commit()
+            commit_started = True
+            self._table_commit.commit(messages)
+            snapshot_id = self._snapshot_recorder.snapshot_id
+            if snapshot_id is None:
+                raise RuntimeError(
+                    "LeRobot batch committed without reporting a snapshot id.")
+        except BaseException:
+            self._fail_batch(abort=not commit_started)
+            raise
+        self._close_batch()
+        self.pending_episodes = 0
+        return snapshot_id
+
+    def finalize(self):
+        if self._finalized:
+            return None
+        self._require_open("finalize")
+        if self._episode_frames:
+            raise RuntimeError(
+                "Cannot finalize with unsaved LeRobot frames; call "
+                "save_episode() or clear_episode_buffer() first.")
+        snapshot_id = self.flush()
+        self._finalized = True
+        return snapshot_id
+
+    def _episode_table(self):
+        episode_index = self.num_episodes
+        first_index = self.num_frames
+        size = len(self._episode_frames)
+        tasks = [frame["task"] for frame in self._episode_frames]
+        task_indices = []
+        for task in tasks:
+            if task not in self._task_indices:
+                self._task_indices[task] = len(self._task_indices)
+            task_indices.append(self._task_indices[task])
+
+        generated = {
+            "timestamp": [index / self.fps for index in range(size)],
+            "frame_index": list(range(size)),
+            "episode_index": [episode_index] * size,
+            "index": list(range(first_index, first_index + size)),
+            "task_index": task_indices,
+        }
+        arrays = []
+        for name, feature in self.features.items():
+            values = generated.get(name)
+            if values is None:
+                values = [frame[name] for frame in self._episode_frames]
+            field = self._source_schema.field(name)
+            arrays.append(_safe_array(
+                values, field, name, str(feature.get("dtype", ""))))
+        arrays.append(pa.array(tasks, type=pa.string()))
+        source = pa.Table.from_arrays(arrays, schema=self._source_schema)
+        return strict_arrow_table(
+            source,
+            self._target_schema,
+            self._table.identifier,
+            self.num_episodes,
+            "LeRobot",
+        )
+
+    def _ensure_batch(self):
+        if self._table_write is not None:
+            return
+        builder = self._table.raw_table.new_batch_write_builder()
+        self._table_write = builder.new_write()
+        self._table_commit = builder.new_commit()
+        self._snapshot_recorder = _SnapshotRecorder()
+        self._table_commit.add_commit_callback(self._snapshot_recorder)
+
+    def _fail_batch(self, abort):
+        self._failed = True
+        if abort and self._table_write is not None:
+            self._table_write.abort()
+        self._close_batch()
+
+    def _close_batch(self):
+        try:
+            if self._table_write is not None:
+                self._table_write.close()
+        finally:
+            if self._table_commit is not None:
+                self._table_commit.close()
+        self._table_write = None
+        self._table_commit = None
+        self._snapshot_recorder = None
+
+    def _require_open(self, method):
+        if self._failed:
+            raise RuntimeError(
+                "Cannot call %s() after a Paimon write failure." % method)
+        if self._finalized:
+            raise RuntimeError(
+                "Cannot call %s() after finalize()." % method)