Adapt versioned dataset reads into evaluation case sources

src/api/btql.rs

@@ -33,6 +33,9 @@ pub struct DatasetRecord {
 #[derive(Debug, Clone, Serialize)]
 pub struct BTQLQuery {
     pub query: BTQLQueryInner,
+    /// Optional immutable dataset snapshot version.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub version: Option<String>,
 }
 
 /// Inner BTQL query structure.

src/dataset.rs

@@ -595,6 +595,7 @@ struct DatasetCursor {
     scope: CredentialScope,
     dataset_id: String,
     batch_size: usize,
+    version: Option<String>,
     cursor: Option<String>,
     buffer: VecDeque<DatasetRecord>,
     done: bool,
@@ -634,6 +635,7 @@ impl DatasetIterator {
                 scope,
                 dataset_id,
                 batch_size,
+                version: None,
                 cursor: None,
                 buffer: VecDeque::new(),
                 done: false,
@@ -690,6 +692,8 @@ impl DatasetIterator {
     }
 }
 
+pub(crate) mod read_session;
+
 impl DatasetCursor {
     /// Get the next record, fetching a new batch if needed.
     async fn next(&mut self) -> Result<Option<DatasetRecord>> {
@@ -707,6 +711,7 @@ impl DatasetCursor {
                     limit: Some(self.batch_size),
                     cursor: previous_cursor.clone(),
                 },
+                version: self.version.clone(),
             };
             // A failed request leaves the cursor and EOF state untouched.
             let response = self

src/dataset/read_session.rs

@@ -0,0 +1,143 @@
+//! One admitted, demand-driven dataset cursor on its resource's fixed binding.
+use super::*;
+use std::future::IntoFuture;
+use tokio::sync::{mpsc, oneshot};
+
+#[derive(Debug, Clone)]
+pub(crate) struct ReadError(Arc<BraintrustError>);
+impl std::fmt::Display for ReadError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+impl std::error::Error for ReadError {
+    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
+        Some(self.0.as_ref())
+    }
+}
+impl From<BraintrustError> for ReadError {
+    fn from(error: BraintrustError) -> Self {
+        Self(Arc::new(error))
+    }
+}
+impl From<crate::ExecutionError> for ReadError {
+    fn from(error: crate::ExecutionError) -> Self {
+        BraintrustError::from(error).into()
+    }
+}
+impl From<crate::EnqueueError> for ReadError {
+    fn from(error: crate::EnqueueError) -> Self {
+        BraintrustError::from(error).into()
+    }
+}
+
+type Reply = oneshot::Sender<std::result::Result<Option<DatasetRecord>, ReadError>>;
+type Driver = futures::future::Shared<BoxFuture<'static, std::result::Result<(), ReadError>>>;
+
+pub(crate) struct ReadSession {
+    requests: Option<mpsc::Sender<Reply>>,
+    ready: Option<oneshot::Receiver<DatasetMetadata>>,
+    metadata: Option<DatasetMetadata>,
+    driver: Driver,
+    done: bool,
+}
+impl ReadSession {
+    pub fn start(
+        dataset: Dataset,
+        batch_size: Option<usize>,
+        version: Option<String>,
+    ) -> std::result::Result<Self, ReadError> {
+        let (requests, mut receiver) = mpsc::channel::<Reply>(1);
+        let (ready, metadata) = oneshot::channel();
+        let execution = dataset.execution.clone();
+        let completion = execution.response(async move {
+            let metadata = dataset.ensure_registered().await?.clone();
+            let mut cursor = DatasetCursor {
+                fetcher: dataset.api.clone(),
+                scope: dataset.scope.clone(),
+                dataset_id: metadata.dataset_id.clone(),
+                batch_size: batch_size.unwrap_or(1000).max(1),
+                version,
+                cursor: None,
+                buffer: VecDeque::new(),
+                done: false,
+            };
+            if ready.send(metadata).is_err() {
+                return Ok(());
+            }
+            while let Some(reply) = receiver.recv().await {
+                let result = cursor.next().await.map_err(ReadError::from);
+                let terminal = !matches!(result, Ok(Some(_)));
+                let failure = result.as_ref().err().cloned();
+                if reply.send(result).is_err() {
+                    break;
+                }
+                if let Some(error) = failure {
+                    return Err(error);
+                }
+                if terminal {
+                    break;
+                }
+            }
+            Ok(())
+        });
+        if let Err(error) = &*completion {
+            return Err(error.clone());
+        }
+        Ok(Self {
+            requests: Some(requests),
+            ready: Some(metadata),
+            metadata: None,
+            driver: completion.into_future().boxed().shared(),
+            done: false,
+        })
+    }
+
+    pub async fn prepare(&mut self) -> std::result::Result<&DatasetMetadata, ReadError> {
+        if let Some(ready) = self.ready.take() {
+            self.metadata = Some(self.receive(ready).await?);
+        }
+        self.metadata.as_ref().ok_or_else(stopped)
+    }
+
+    pub async fn next(&mut self) -> std::result::Result<Option<DatasetRecord>, ReadError> {
+        self.prepare().await?;
+        if self.done {
+            return Ok(None);
+        }
+        let (reply, received) = oneshot::channel();
+        if self
+            .requests
+            .as_ref()
+            .ok_or_else(stopped)?
+            .send(reply)
+            .await
+            .is_err()
+        {
+            return Err(self.driver.clone().await.err().unwrap_or_else(stopped));
+        }
+        let result = self.receive(received).await?;
+        let result = result?;
+        if result.is_none() {
+            self.done = true;
+            self.requests = None;
+        }
+        Ok(result)
+    }
+
+    async fn receive<T>(
+        &self,
+        receiver: oneshot::Receiver<T>,
+    ) -> std::result::Result<T, ReadError> {
+        use futures::future::{select, Either};
+        // Prefer a reply already produced before a terminal driver result.
+        match select(receiver, self.driver.clone()).await {
+            Either::Left((Ok(value), _)) => Ok(value),
+            Either::Left((Err(_), driver)) => Err(driver.await.err().unwrap_or_else(stopped)),
+            Either::Right((result, _)) => Err(result.err().unwrap_or_else(stopped)),
+        }
+    }
+}
+fn stopped() -> ReadError {
+    BraintrustError::Background("dataset read session stopped before replying".into()).into()
+}

src/eval/executor.rs

@@ -87,12 +87,19 @@ impl EvalExecutor {
         binding: Option<&crate::internal::execution::ExecutionBinding>,
     ) -> std::result::Result<EvalSummary<I, O>, EvalRunError<I, O>> {
         reporter.started(&plan.name);
+        let prepared = plan.source.prepare().await;
         let dataset = plan.source.provenance();
         let mut results = Vec::new();
         let mut failures = Vec::new();
+        if let Err(source) = prepared {
+            failures.push(EvalRunFailure::Source {
+                case_index: 0,
+                source,
+            });
+        }
         let mut case_count = 0;
         let mut pending = FuturesUnordered::new();
-        loop {
+        while failures.is_empty() {
             // Keep one source read alive while cases run. Sources may await
             // case progress and need not support cancellation/restarting next().
             let next_case = {

src/eval/mod.rs

@@ -73,7 +73,8 @@ pub use run::{run, EvalRequest};
 pub use runner::{EvalOpts, Evaluator};
 pub use scorer::{scorer, ExactMatch, LevenshteinScorer, NumericDiffScorer, Scorer, SyncFnScorer};
 pub use source::{
-    CaseSource, DatasetVersionRef, IterCaseSource, SourceError, StreamCaseSource, VecCaseSource,
+    CaseSource, DatasetCaseSource, DatasetVersionRef, IterCaseSource, SourceError,
+    StreamCaseSource, VecCaseSource,
 };
 pub use task::{task, SimpleFnTask, SyncFnTask, Task};
 pub use types::{

src/eval/source.rs

@@ -34,6 +34,11 @@ pub struct DatasetVersionRef {
 /// without requiring Sync, cloning user values, or serializing local cases.
 #[async_trait]
 pub trait CaseSource<I, O>: Send {
+    /// Resolve source identity before provenance is captured or cases are read.
+    async fn prepare(&mut self) -> Result<(), SourceError> {
+        Ok(())
+    }
+
     async fn next(&mut self) -> Result<Option<Case<I, O>>, SourceError>;
     fn provenance(&self) -> Option<DatasetVersionRef> {
         None
@@ -103,3 +108,94 @@ impl<I: Send + 'static, O: Send + 'static> From<Vec<Case<I, O>>> for Box<dyn Cas
         Box::new(VecCaseSource::new(cases))
     }
 }
+
+/// Demand-driven typed cases from an SDK dataset on its bound runtime.
+///
+/// Construction admits one bounded read session immediately. Consume the source
+/// or drop it before waiting for its client's shutdown; a live unread source is
+/// accepted work. Dropping the source closes its session. Pages are fetched only
+/// as cases are requested, with no nested response-slot reservation.
+///
+/// `version` pins every page to that snapshot. With `None`, the source reads live
+/// data and reports no snapshot version. Input/output deserialization is required
+/// only at this source boundary; values need not implement Clone or Serialize.
+///
+/// ```no_run
+/// use braintrust_sdk_rust::{Dataset, eval::{CaseSource, DatasetCaseSource}};
+/// # async fn read(dataset: Dataset) -> anyhow::Result<()> {
+/// let mut cases = DatasetCaseSource::<String, String>::new(
+///     dataset, Some("snapshot-version".into()), Some(100),
+/// )?;
+/// cases.prepare().await?;
+/// while let Some(case) = cases.next().await? { println!("{}", case.input); }
+/// # Ok(())
+/// # }
+/// ```
+pub struct DatasetCaseSource<I, O> {
+    session: crate::dataset::read_session::ReadSession,
+    provenance: Option<DatasetVersionRef>,
+    version: Option<String>,
+    marker: std::marker::PhantomData<fn() -> (I, O)>,
+}
+impl<I, O> DatasetCaseSource<I, O> {
+    pub fn new(
+        dataset: crate::Dataset,
+        version: Option<String>,
+        batch_size: Option<usize>,
+    ) -> Result<Self, SourceError> {
+        Ok(Self {
+            session: crate::dataset::read_session::ReadSession::start(
+                dataset,
+                batch_size,
+                version.clone(),
+            )?,
+            provenance: None,
+            version,
+            marker: std::marker::PhantomData,
+        })
+    }
+}
+#[async_trait]
+impl<I: serde::de::DeserializeOwned + Send, O: serde::de::DeserializeOwned + Send> CaseSource<I, O>
+    for DatasetCaseSource<I, O>
+{
+    async fn prepare(&mut self) -> Result<(), SourceError> {
+        let metadata = self.session.prepare().await?;
+        if self.provenance.is_none() {
+            self.provenance = Some(DatasetVersionRef {
+                dataset_id: metadata.dataset_id.clone(),
+                version: self.version.clone(),
+            });
+        }
+        Ok(())
+    }
+    fn provenance(&self) -> Option<DatasetVersionRef> {
+        self.provenance.clone()
+    }
+    async fn next(&mut self) -> Result<Option<Case<I, O>>, SourceError> {
+        self.prepare().await?;
+        let Some(record) = self.session.next().await? else {
+            return Ok(None);
+        };
+        let input = record.input.unwrap_or(serde_json::Value::Null);
+        let input = serde_json::from_value(input).map_err(|error| {
+            anyhow::Error::new(error).context(format!("dataset record {} input", record.id))
+        })?;
+        let expected = record
+            .expected
+            .map(serde_json::from_value)
+            .transpose()
+            .map_err(|error| {
+                anyhow::Error::new(error).context(format!("dataset record {} expected", record.id))
+            })?;
+        Ok(Some(Case {
+            input,
+            expected,
+            metadata: record.metadata,
+            tags: record.tags,
+            id: Some(record.id),
+            xact_id: record.xact_id,
+            created: record.created,
+        }))
+    }
+}

src/internal/upload.rs

@@ -313,7 +313,11 @@ mod tests {
                     .any(|key| key.as_str() == "x-bt-org-name"));
                 assert_eq!(request.body, data);
             }
-            let grant: Value = serde_json::from_slice(&requests[0].body).unwrap();
+            let grant_request = requests
+                .iter()
+                .find(|request| request.url.path() == "/logs3/overflow")
+                .expect("overflow grant request");
+            let grant: Value = serde_json::from_slice(&grant_request.body).unwrap();
             assert_eq!(grant["size_bytes"], data.len());
             assert_eq!(grant["rows"][0]["object_ids"]["project_id"], "project");
         }

tests/api_client.rs

@@ -147,6 +147,7 @@ async fn btql_readback_preserves_cursor_and_request_errors() {
                 limit: Some(5),
                 cursor: Some("next".into()),
             },
+            version: None,
         })
         .await
         .unwrap();
@@ -169,7 +170,8 @@ async fn btql_readback_preserves_cursor_and_request_errors() {
                     from: "dataset('id')".into(),
                     limit: None,
                     cursor: None,
-                }
+                },
+                version: None,
             })
             .await,
         Err(api::Error::Http { status: 500, .. })