Simplify data model (#16)

* Simplify data model

* Simplify internal data model

Author: ferranbt

crates/strata/src/datagen.rs

@@ -12,7 +12,7 @@ use schema::{DataType, Field, Schema};
 use serde_json::{Map, Value, json};
 
 use crate::dataset::Dataset;
-use crate::record::{Records, stringify_text_columns};
+use crate::record::{Batch, stringify_text_columns};
 
 pub struct Generator {
     schema: Schema,
@@ -42,10 +42,10 @@ impl Generator {
     }
 
     /// The rows in `range`, Arrow-encoded — one page.
-    pub fn rows(&self, range: Range<usize>) -> Result<Records> {
+    pub fn rows(&self, range: Range<usize>) -> Result<Batch> {
         let mut rows: Vec<Value> = range.map(|i| self.row(i)).collect();
         stringify_text_columns(&self.schema, &mut rows);
-        Records::encode(self.schema.clone(), &rows)
+        Batch::encode(&self.schema, &rows)
     }
 
     /// The first `n` rows as a [`Dataset`] (schema + Arrow records).

crates/strata/src/dataset.rs

@@ -7,13 +7,15 @@
 //! over Arrow Flight the same `(schema, rows)` rides natively in a `DoPut`.
 
 use anyhow::{Result, bail};
-use futures::stream::{BoxStream, StreamExt};
+use futures::stream::StreamExt;
 use schema::Schema;
 use serde::Serialize;
 use serde_json::Value;
 
-use crate::page::Cursor;
-use crate::record::Records;
+use crate::{
+    DataStream,
+    record::{Batch, BatchPage},
+};
 
 /// How a sink should apply a written dataset. Rides as metadata on the existing
 /// `put` verb (the reserved `disposition` query param) rather than a new verb, so
@@ -66,11 +68,11 @@ impl Disposition {
 #[derive(Debug, Clone)]
 pub struct Dataset {
     pub schema: Schema,
-    pub records: Records,
+    pub records: Batch,
 }
 
 impl Dataset {
-    pub fn new(schema: Schema, records: Records) -> Self {
+    pub fn new(schema: Schema, records: Batch) -> Self {
         Dataset { schema, records }
     }
 
@@ -79,21 +81,21 @@ impl Dataset {
     /// against it. The typed-input counterpart of a `put`.
     pub fn of<T: serde::Serialize + schema::HasSchema>(rows: &[T]) -> Result<Self> {
         let schema = T::schema();
-        let records = Records::encode(schema.clone(), rows)?;
+        let records = Batch::encode(&schema, rows)?;
         Ok(Dataset::new(schema, records))
     }
 
     /// Interim bridge: decode the Arrow rows to JSON for sinks that aren't yet
     /// Arrow-native. This is the one labeled Arrow→JSON on the write path; Phase B
     /// removes it as each sink binds Arrow columns directly.
     pub fn to_json_rows(&self) -> Result<Vec<Value>> {
-        self.records.to_json_rows()
+        self.records.to_json_rows(&self.schema)
     }
 
     /// Bridge to a single-chunk [`DataStream`], until readers produce batches lazily.
     pub fn into_stream(self) -> DataStream {
-        let chunk = Chunk {
-            records: self.records,
+        let chunk = BatchPage {
+            data: self.records,
             cursor: None,
         };
         DataStream {
@@ -102,57 +104,3 @@ impl Dataset {
         }
     }
 }
-
-/// One batch of a data-plane stream, with an optional resume [`Cursor`] valid
-/// after it (sparse — `Some` only at checkpoint boundaries). Over Flight it maps
-/// to a `FlightData`'s `app_metadata`.
-pub struct Chunk {
-    pub records: Records,
-    pub cursor: Option<Cursor>,
-}
-
-impl Chunk {
-    /// Decode this batch's rows into typed values.
-    pub fn decode<T: serde::de::DeserializeOwned>(&self) -> Result<Vec<T>> {
-        self.records.decode()
-    }
-}
-
-/// The streaming form of a [`Dataset`]: the native `schema` plus a lazy stream of
-/// [`Chunk`]s following it, consumed by a sink without materializing the table.
-pub struct DataStream {
-    pub schema: Schema,
-    pub chunks: BoxStream<'static, Result<Chunk>>,
-}
-
-impl DataStream {
-    /// Pull just the first [`Chunk`] (one page) — the `list_once` case for callers
-    /// that want a bounded read instead of draining. The schema rides alongside.
-    pub async fn first(mut self) -> Result<Option<Chunk>> {
-        self.chunks.next().await.transpose()
-    }
-
-    /// Drain every chunk into one materialized [`Dataset`] — the bridge for sinks
-    /// that don't consume the stream incrementally.
-    pub async fn collect(self) -> Result<Dataset> {
-        let DataStream { schema, mut chunks } = self;
-        let mut batches = Vec::new();
-        let mut arrow_schema = None;
-        while let Some(chunk) = chunks.next().await {
-            let records = chunk?.records;
-            arrow_schema.get_or_insert_with(|| records.schema.clone());
-            batches.extend(records.batches);
-        }
-        let arrow_schema = match arrow_schema {
-            Some(s) => s,
-            None => std::sync::Arc::new(crate::record::strata_schema_to_arrow_schema(&schema)),
-        };
-        Ok(Dataset::new(
-            schema,
-            Records {
-                schema: arrow_schema,
-                batches,
-            },
-        ))
-    }
-}

crates/strata/src/flight/mod.rs

@@ -28,8 +28,10 @@ use tonic::transport::Server;
 use tonic::{Request, Response, Status, Streaming};
 
 use crate::Registry;
-use crate::dataset::{Chunk, DataStream};
-use crate::record::{Records, arrow_schema_to_strata, strata_schema_to_arrow_schema};
+use crate::record::{
+    Batch, BatchPage, DataStream, arrow_schema_to_strata, strata_schema_to_arrow_schema,
+};
+use schema::Schema as StrataSchema;
 use crate::router::{Body, Method};
 
 /// Start the Flight server on `addr`, serving every registered provider.
@@ -115,20 +117,20 @@ fn ticket_target(ticket: &Ticket) -> Result<(String, String), Status> {
 
 type FlightStream<T> = BoxStream<'static, Result<T, Status>>;
 
-/// Encode one [`Chunk`] (a page): its batches as Arrow `FlightData`, then a
-/// metadata-only message whose `app_metadata` is the chunk's cursor (when present)
-/// — the per-batch checkpoint, forwarded opaquely.
-fn encode_chunk(chunk: Chunk) -> FlightStream<FlightData> {
-    let Chunk { records, cursor } = chunk;
-    let schema = records.schema.clone();
-    let batches = futures::stream::iter(
-        records
-            .batches
-            .into_iter()
-            .map(Ok::<RecordBatch, FlightError>),
-    );
+/// Encode one [`BatchPage`]: its columns as Arrow `FlightData`
+fn encode_chunk(chunk: BatchPage, schema: &StrataSchema) -> FlightStream<FlightData> {
+    let BatchPage { data, cursor } = chunk;
+    let batch = match data.to_record_batch(schema) {
+        Ok(batch) => batch,
+        Err(e) => {
+            return futures::stream::once(async move { Err(Status::internal(e.to_string())) })
+                .boxed();
+        }
+    };
+    let arrow_schema = batch.schema();
+    let batches = futures::stream::iter([Ok::<RecordBatch, FlightError>(batch)]);
     let data = FlightDataEncoderBuilder::new()
-        .with_schema(schema)
+        .with_schema(arrow_schema)
         .build(batches)
         .map(|d| d.map_err(|e| Status::internal(e.to_string())));
     let trailer = cursor.map(|cursor| {
@@ -225,10 +227,11 @@ impl FlightService for StrataFlight {
             .read(&path)
             .await
             .map_err(|e| Status::internal(e.to_string()))?;
+        let schema = stream.schema.clone();
         let out = stream
             .chunks
-            .flat_map(|chunk| match chunk {
-                Ok(chunk) => encode_chunk(chunk),
+            .flat_map(move |chunk| match chunk {
+                Ok(chunk) => encode_chunk(chunk, &schema),
                 Err(e) => {
                     futures::stream::once(async move { Err(Status::internal(e.to_string())) })
                         .boxed()
@@ -291,24 +294,19 @@ impl FlightService for StrataFlight {
             .clone();
         let schema = arrow_schema_to_strata(&arrow_schema);
 
-        // One chunk per batch; the tail streams lazily. No cursor on the write path.
-        let rest_schema = arrow_schema.clone();
-        let head_chunk = futures::stream::iter(first_batch.map(move |batch| {
-            Ok::<_, anyhow::Error>(Chunk {
-                records: Records {
-                    schema: arrow_schema,
-                    batches: vec![batch],
-                },
+        // One page per batch; the tail streams lazily. No cursor on the write path.
+        // The first batch was pulled only to make the schema available, so it is put
+        // back on the front here.
+        let head_chunk = futures::stream::iter(first_batch.map(|batch| {
+            Ok::<_, anyhow::Error>(BatchPage {
+                data: Batch::from_record_batch(&batch),
                 cursor: None,
             })
         }));
-        let rest_chunks = decoder.map(move |batch| {
+        let rest_chunks = decoder.map(|batch| {
             let batch = batch.map_err(|e| anyhow::anyhow!(e.to_string()))?;
-            Ok(Chunk {
-                records: Records {
-                    schema: rest_schema.clone(),
-                    batches: vec![batch],
-                },
+            Ok(BatchPage {
+                data: Batch::from_record_batch(&batch),
                 cursor: None,
             })
         });

crates/strata/src/graphql/mod.rs

@@ -110,12 +110,13 @@ async fn tables(registry: &Arc<Registry>, mount: &str) -> Vec<String> {
     let Ok(stream) = provider.read("/tables").await else {
         return Vec::new();
     };
+    let schema = stream.schema.clone();
     let Ok(Some(chunk)) = stream.first().await else {
         return Vec::new();
     };
     chunk
-        .records
-        .to_json_rows()
+        .data
+        .to_json_rows(&schema)
         .unwrap_or_default()
         .iter()
         .filter_map(|row| row.get("name")?.as_str().map(String::from))
@@ -197,8 +198,10 @@ fn nested_field(
                 serde_json::json!({ "cmp": { "field": target_col, "op": "eq", "value": fk } });
             let encoded = urlencoding::encode(&filter.to_string()).into_owned();
             let path = format!("/tables/{target_table}?filter={encoded}&limit=1");
-            let rows = match registry.get(&mount)?.read(&path).await?.first().await? {
-                Some(chunk) => chunk.records.to_json_rows()?,
+            let stream = registry.get(&mount)?.read(&path).await?;
+            let schema = stream.schema.clone();
+            let rows = match stream.first().await? {
+                Some(chunk) => chunk.data.to_json_rows(&schema)?,
                 None => Vec::new(),
             };
             Ok(rows.into_iter().next().map(FieldValue::owned_any))
@@ -249,8 +252,10 @@ fn table_field(
 
             let path = read_path(&table, filter.as_ref(), &fields, limit);
             let provider = registry.get(&mount)?;
-            let rows = match provider.read(&path).await?.first().await? {
-                Some(chunk) => chunk.records.to_json_rows()?,
+            let stream = provider.read(&path).await?;
+            let schema = stream.schema.clone();
+            let rows = match stream.first().await? {
+                Some(chunk) => chunk.data.to_json_rows(&schema)?,
                 None => Vec::new(),
             };
             Ok(Some(FieldValue::list(

crates/strata/src/lib.rs

@@ -24,9 +24,10 @@ use anyhow::Result;
 
 pub use catalog::Catalog;
 pub use config::{Config, ProviderConfig};
-pub use dataset::{Chunk, DataStream, Dataset, Disposition};
+pub use dataset::{Dataset, Disposition};
 pub use page::{Cursor, Page};
 pub use provider::{Provider, ProviderObject, Registry};
+pub use record::DataStream;
 pub use router::{Body, EndpointInfo, Method, Params, Response, Router};
 
 /// Config file consulted by [`registry`] when no explicit path is given.

crates/strata/src/main.rs

@@ -136,13 +136,17 @@ async fn run() -> Result<()> {
             let display = match method {
                 // Data plane: read the stream and show the first chunk (one page)
                 // as `{ items, cursor }`. Drain the whole stream with `--follow`.
-                Verb::List => match provider.read(&path).await?.first().await? {
-                    Some(chunk) => serde_json::json!({
-                        "items": chunk.records.to_json_rows()?,
-                        "cursor": chunk.cursor,
-                    }),
-                    None => serde_json::json!({ "items": [], "cursor": null }),
-                },
+                Verb::List => {
+                    let stream = provider.read(&path).await?;
+                    let schema = stream.schema.clone();
+                    match stream.first().await? {
+                        Some(chunk) => serde_json::json!({
+                            "items": chunk.data.to_json_rows(&schema)?,
+                            "cursor": chunk.cursor,
+                        }),
+                        None => serde_json::json!({ "items": [], "cursor": null }),
+                    }
+                }
                 // Entity plane: `get` takes no body; `create` reads its JSON entity
                 // from stdin (the `meta` channel).
                 Verb::Get | Verb::Create => {
@@ -290,12 +294,13 @@ fn strip_query(path: &str) -> &str {
 /// this just consumes the stream to the tail.
 async fn follow_list(provider: &dyn strata::ProviderObject, path: &str) -> Result<()> {
     let mut stream = provider.read(path).await?;
+    let schema = stream.schema.clone();
     let mut page = 1;
     while let Some(chunk) = stream.chunks.next().await {
         let chunk = chunk?;
         let output = serde_json::json!({
             "page": page,
-            "items": chunk.records.to_json_rows()?,
+            "items": chunk.data.to_json_rows(&schema)?,
             "cursor": chunk.cursor,
         });
         println!("{}", serde_json::to_string_pretty(&output)?);

crates/strata/src/pipe/mod.rs

@@ -9,7 +9,7 @@ use futures::StreamExt;
 use serde_json::Value;
 use strata_types::Pipe;
 
-use crate::dataset::DataStream;
+use crate::record::DataStream;
 use crate::request::{ReadRequest, WriteRequest};
 use crate::router::Body;
 use crate::{Method, Registry};
@@ -48,7 +48,7 @@ pub async fn run_pass<S: PipeStore>(registry: &Registry, store: &S, pipe: &mut P
         let progress = sync::advance(strategy, &chunk);
 
         // Write the rows (skip an empty tail page — nothing to store).
-        let rows = chunk.records.row_count() as u64;
+        let rows = chunk.data.row_count() as u64;
         if rows > 0 {
             let page = DataStream {
                 schema: schema.clone(),
@@ -157,7 +157,7 @@ mod tests {
             let mut chunks = stream.chunks;
             let mut total = 0usize;
             while let Some(chunk) = chunks.next().await {
-                total += chunk?.records.row_count();
+                total += chunk?.data.row_count();
             }
             anyhow::Ok(total)
         };

crates/strata/src/pipe/sync.rs

@@ -17,10 +17,10 @@
 //!   chunk that brought no rows means we've reached the live edge for now, so the
 //!   driver stops pulling and drops the (otherwise infinite) stream.
 
-use crate::dataset::Chunk;
 use crate::page::ListStrategy;
+use crate::record::BatchPage;
 
-/// What the sync driver should do after writing one [`Chunk`].
+/// What the sync driver should do after writing one [`BatchPage`].
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Progress {
     /// More to pull right now: keep consuming the live stream. The token is this
@@ -36,7 +36,7 @@ pub enum Progress {
 ///
 /// `Offset` follows the cursor until it goes `None`; `NextLink`'s cursor never goes
 /// `None`, so it stops when a page comes back empty. See the module docs.
-pub fn advance(strategy: ListStrategy, chunk: &Chunk) -> Progress {
+pub fn advance(strategy: ListStrategy, chunk: &BatchPage) -> Progress {
     let next = chunk.cursor.as_ref().and_then(|c| c.next.clone());
     match strategy {
         // Finite backfill: drain while the cursor advances; `None` is the tail.
@@ -49,7 +49,7 @@ pub fn advance(strategy: ListStrategy, chunk: &Chunk) -> Progress {
         // `None` — marks the live edge. When caught up, resume from the same
         // forward position so the next run picks up whatever has since arrived.
         ListStrategy::NextLink => {
-            if chunk.records.row_count() == 0 {
+            if chunk.data.row_count() == 0 {
                 Progress::CaughtUp(next)
             } else {
                 match next {

crates/strata/src/provider.rs

@@ -13,7 +13,7 @@ use serde_json::{Value, json};
 
 use crate::catalog::Catalog;
 use crate::config::ProviderConfig;
-use crate::dataset::DataStream;
+use crate::record::DataStream;
 use crate::router::{Body, BoxFuture, EndpointInfo, Method, Response, Router, SchemaSource};
 
 /// Implemented by each concrete provider. Knows its state type and how to wire