feat(rollout): add ordered trajectory reads

crates/lance-context-api/src/lib.rs

@@ -137,6 +137,35 @@ pub trait RolloutStoreApi {
         filters: Option<Value>,
     ) -> impl Future<Output = ContextResult<Vec<RolloutRecordDto>>> + Send;
 
+    fn get_trajectory(
+        &self,
+        rollout_id: &str,
+    ) -> impl Future<Output = ContextResult<Vec<RolloutRecordDto>>> + Send
+    where
+        Self: Sync,
+    {
+        async move {
+            if rollout_id.is_empty() {
+                return Err(ContextError::InvalidRequest(
+                    "rollout_id must not be empty".to_string(),
+                ));
+            }
+            let mut records = self
+                .list_filtered(
+                    None,
+                    None,
+                    Some(serde_json::json!({"rollout_id": rollout_id})),
+                )
+                .await?;
+            records.sort_by(|left, right| {
+                left.sequence_order
+                    .cmp(&right.sequence_order)
+                    .then_with(|| left.id.cmp(&right.id))
+            });
+            Ok(records)
+        }
+    }
+
     fn get(&self, id: &str)
         -> impl Future<Output = ContextResult<Option<RolloutRecordDto>>> + Send;
 

crates/lance-context-core/src/api_impl.rs

@@ -416,6 +416,13 @@ impl RolloutStoreApi for RolloutStore {
         Ok(records.into_iter().map(rollout_record_to_dto).collect())
     }
 
+    async fn get_trajectory(&self, rollout_id: &str) -> ContextResult<Vec<RolloutRecordDto>> {
+        let records = RolloutStore::get_trajectory(self, rollout_id)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(records.into_iter().map(rollout_record_to_dto).collect())
+    }
+
     async fn get(&self, id: &str) -> ContextResult<Option<RolloutRecordDto>> {
         let record = RolloutStore::get_by_id(self, id)
             .await

crates/lance-context-core/src/rollout_store.rs

@@ -965,6 +965,27 @@ impl RolloutStore {
         Ok(results)
     }
 
+    /// Return one complete trajectory in deterministic message order.
+    pub async fn get_trajectory(&self, rollout_id: &str) -> LanceResult<Vec<RolloutRecord>> {
+        if rollout_id.is_empty() {
+            return Err(LanceError::from(ArrowError::InvalidArgumentError(
+                "rollout_id must not be empty".to_string(),
+            )));
+        }
+
+        let filters = RolloutFilters {
+            rollout_id: Some(rollout_id.to_string()),
+            ..Default::default()
+        };
+        let mut records = self.list_with_filters(None, None, Some(&filters)).await?;
+        records.sort_by(|left, right| {
+            left.sequence_order
+                .cmp(&right.sequence_order)
+                .then_with(|| left.id.cmp(&right.id))
+        });
+        Ok(records)
+    }
+
     /// Filter and page rollout rows in the LSM execution plan.
     ///
     /// Reads one row beyond the requested page to report `has_more`, avoiding
@@ -2264,6 +2285,54 @@ mod tests {
         });
     }
 
+    #[test]
+    fn trajectory_rows_are_filtered_and_sorted_across_fragments() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+        runtime.block_on(async {
+            let mut store = RolloutStore::open_with_options(
+                &uri,
+                RolloutStoreOptions {
+                    shard_id: Some("trajectory-test".to_string()),
+                    merge_after_generations: Some(1),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+            let mut later = assistant_record("row-a");
+            later.rollout_id = "target".to_string();
+            later.sequence_order = 2;
+            let mut unrelated = assistant_record("other");
+            unrelated.rollout_id = "other".to_string();
+            store.add(&[later, unrelated]).await.unwrap();
+
+            let mut earlier = assistant_record("row-b");
+            earlier.rollout_id = "target".to_string();
+            earlier.sequence_order = 0;
+            store.add(&[earlier]).await.unwrap();
+
+            assert_eq!(store.observe().await.unwrap().fragment_count, 2);
+            let rows = store.get_trajectory("target").await.unwrap();
+            let ids: Vec<&str> = rows.iter().map(|row| row.id.as_str()).collect();
+            assert_eq!(ids, vec!["row-b", "row-a"]);
+        });
+    }
+
+    #[test]
+    fn trajectory_read_rejects_an_empty_rollout_id() {
+        let dir = TempDir::new().unwrap();
+        let uri = dir.path().to_string_lossy().to_string();
+        let runtime = tokio::runtime::Runtime::new().unwrap();
+        runtime.block_on(async {
+            let store = RolloutStore::open(&uri).await.unwrap();
+            let err = store.get_trajectory("").await.unwrap_err();
+            assert!(err.to_string().contains("rollout_id must not be empty"));
+        });
+    }
+
     /// Read the number of un-merged flushed generations recorded for a store's
     /// own write shard. Used by merge tests to assert the manifest drains.
     async fn flushed_generation_count(store: &RolloutStore) -> usize {

python/python/lance_context/api.py

@@ -2594,6 +2594,11 @@ def list(
         raw = self._sync.list(limit, offset, _json_dumps(filters, "filters"))
         return [_rollout_record_from_json(r) for r in json.loads(raw)]
 
+    def get_trajectory(self, rollout_id: str) -> list[dict[str, Any]]:
+        """Return one complete trajectory ordered by sequence and row id."""
+        raw = self._sync.get_trajectory(rollout_id)
+        return [_rollout_record_from_json(r) for r in json.loads(raw)]
+
     def get(self, id: str) -> dict[str, Any] | None:
         """Fetch a single rollout row by id, or ``None``."""
         raw = self._sync.get(id)
@@ -2681,6 +2686,14 @@ async def list(
         )
         return [_rollout_record_from_json(r) for r in json.loads(raw)]
 
+    async def get_trajectory(self, rollout_id: str) -> list[dict[str, Any]]:
+        """Return one complete trajectory ordered by sequence and row id."""
+        loop = asyncio.get_running_loop()
+        raw = await loop.run_in_executor(
+            None, lambda: self._sync.get_trajectory(rollout_id)
+        )
+        return [_rollout_record_from_json(r) for r in json.loads(raw)]
+
     async def get(self, id: str) -> dict[str, Any] | None:
         """Fetch a single rollout row by id, or ``None``."""
         loop = asyncio.get_running_loop()

python/src/lib.rs

@@ -2498,6 +2498,14 @@ impl RolloutStore {
         rollout_records_to_json(&records)
     }
 
+    /// Return one complete trajectory in deterministic message order.
+    fn get_trajectory(&self, py: Python<'_>, rollout_id: &str) -> PyResult<String> {
+        let records = py
+            .allow_threads(|| self.runtime.block_on(self.store.get_trajectory(rollout_id)))
+            .map_err(to_py_err)?;
+        rollout_records_to_json(&records)
+    }
+
     /// Fetch a single rollout row by id, or `None`. Returns a JSON object
     /// string (artifact bytes projected out).
     fn get(&self, py: Python<'_>, id: &str) -> PyResult<Option<String>> {

python/tests/test_rollout.py

@@ -179,3 +179,17 @@ def test_list_rejects_invalid_filter_fields(store_uri):
     store = RolloutStore.open(store_uri)
     with pytest.raises(RuntimeError, match="unsupported rollout filter"):
         store.list(filters={"reward": 1.0})
+
+
+def test_get_trajectory_orders_rows(store_uri):
+    store = RolloutStore.open(store_uri)
+    store.add(
+        [
+            {"id": "row-a", "rollout_id": "target", "sequence_order": 2},
+            {"id": "other", "rollout_id": "other", "sequence_order": 0},
+            {"id": "row-b", "rollout_id": "target", "sequence_order": 0},
+        ]
+    )
+
+    rows = store.get_trajectory("target")
+    assert [row["id"] for row in rows] == ["row-b", "row-a"]

python/tests/test_rollout_remote.py

@@ -159,3 +159,20 @@ async def run():
         assert [row["id"] for row in rows] == ["row-7"]
 
     asyncio.run(run())
+
+
+def test_remote_get_trajectory_orders_rows(server):
+    async def run():
+        store = await AsyncRolloutStore.connect_or_create(server, "rl-trajectory")
+        await store.add(
+            [
+                {"id": "row-a", "rollout_id": "target", "sequence_order": 2},
+                {"id": "other", "rollout_id": "other", "sequence_order": 0},
+                {"id": "row-b", "rollout_id": "target", "sequence_order": 0},
+            ]
+        )
+
+        rows = await store.get_trajectory("target")
+        assert [row["id"] for row in rows] == ["row-b", "row-a"]
+
+    asyncio.run(run())