feldera
diff --git a/‎crates/adapterlib/src/catalog.rs‎
Lines changed: 217 additions & 3 deletions b/‎crates/adapterlib/src/catalog.rs‎
Lines changed: 217 additions & 3 deletions
@@ -14,7 +14,7 @@ use apache_avro::{
 };
 use arrow::record_batch::RecordBatch;
 use dbsp::circuit::NodeId;
-use dbsp::dynamic::{DynData, DynVec};
+use dbsp::dynamic::{ClonableTrait, DynData, DynVec, Factory};
 use dbsp::operator::StagedBuffers;
 use dyn_clone::DynClone;
 use feldera_sqllib::Variant;
@@ -281,6 +281,10 @@ pub trait SerBatchReader: 'static + Send + Sync {
 
     fn snapshot(&self) -> Arc<dyn SerBatchReader>;
 
+    fn keys_factory(&self) -> &'static dyn Factory<DynVec<DynData>>;
+
+    fn data_factory(&self) -> &'static dyn Factory<DynData>;
+
     fn sample_keys(&self, sample_size: usize, sample: &mut DynVec<DynData>);
 
     fn partition_keys(&self, num_partitions: usize, bounds: &mut DynVec<DynData>);
@@ -354,6 +358,216 @@ pub trait SerTrace: SerBatchReader {
     fn as_batch_reader(&self) -> &dyn SerBatchReader;
 }
 
+pub struct SplitCursorBuilder {
+    batch: Arc<dyn SerBatchReader>,
+    start_key: Box<DynData>,
+    end_key: Option<Box<DynData>>,
+    format: RecordFormat,
+}
+
+impl SplitCursorBuilder {
+    /// Create a [`SplitCursorBuilder`] for partition `index` given a batch,
+    /// pre-computed partition `bounds` (as returned by
+    /// [`SerBatchReader::partition_keys`]), and a record `format`.
+    ///
+    /// `bounds` contains `N-1` boundary keys for `N` partitions.
+    /// Partition 0 spans from the start of the batch to `bounds[0]`,
+    /// partition `i` spans from `bounds[i-1]` to `bounds[i]`, and the last
+    /// partition spans from `bounds[N-2]` to the end of the batch.
+    ///
+    /// Returns `None` if the partition is empty (the cursor has no key at the
+    /// start position).
+    pub fn from_bounds(
+        batch: Arc<dyn SerBatchReader>,
+        bounds: &DynVec<DynData>,
+        index: usize,
+        format: RecordFormat,
+    ) -> Option<Self> {
+        let start_bound = if index == 0 {
+            None
+        } else if index <= bounds.len() {
+            Some(bounds.index(index - 1).as_data())
+        } else {
+            None
+        };
+
+        let end_bound = if index < bounds.len() {
+            Some(bounds.index(index).as_data())
+        } else {
+            None
+        };
+
+        let start_key = {
+            let mut cursor = batch.cursor(format.clone()).unwrap();
+
+            // Seek to start. If None, the cursor starts at the beginning.
+            if let Some(start_bound) = start_bound {
+                cursor.seek_key_exact(start_bound);
+            }
+
+            // Clone the actual key the cursor landed on.
+            cursor.get_key().map(|s| {
+                let mut key = batch.data_factory().default_box();
+                s.clone_to(key.as_mut());
+                key
+            })
+        }?;
+
+        let end_key = end_bound.map(|e| {
+            let mut key = batch.data_factory().default_box();
+            e.clone_to(key.as_mut());
+            key
+        });
+
+        Some(SplitCursorBuilder {
+            batch,
+            start_key,
+            end_key,
+            format,
+        })
+    }
+
+    pub fn build<'a>(&'a self) -> SplitCursor<'a> {
+        let mut cursor = self.batch.cursor(self.format.clone()).unwrap();
+        cursor.seek_key_exact(self.start_key.as_data());
+
+        SplitCursor {
+            cursor,
+            start_key: self.start_key.clone(),
+            end_key: self.end_key.clone(),
+        }
+    }
+}
+
+pub struct SplitCursor<'a> {
+    cursor: Box<dyn SerCursor + 'a>,
+    start_key: Box<DynData>,
+    end_key: Option<Box<DynData>>,
+}
+
+impl SplitCursor<'_> {
+    fn finished(&self) -> bool {
+        if let Some(ref end_key) = self.end_key
+            && let Some(current_key) = self.cursor.get_key()
+        {
+            return current_key >= end_key.as_data();
+        }
+
+        false
+    }
+}
+
+impl SerCursor for SplitCursor<'_> {
+    fn key_valid(&self) -> bool {
+        self.cursor.key_valid() && !self.finished()
+    }
+
+    fn val_valid(&self) -> bool {
+        self.cursor.val_valid()
+    }
+
+    fn key(&self) -> &DynData {
+        self.cursor.key()
+    }
+
+    fn get_key(&self) -> Option<&DynData> {
+        if !self.key_valid() {
+            return None;
+        }
+
+        self.cursor.get_key()
+    }
+
+    fn serialize_key(&mut self, dst: &mut Vec<u8>) -> AnyResult<()> {
+        self.cursor.serialize_key(dst)
+    }
+
+    fn key_to_json(&mut self) -> AnyResult<serde_json::Value> {
+        self.cursor.key_to_json()
+    }
+
+    fn serialize_key_fields(
+        &mut self,
+        fields: &HashSet<String>,
+        dst: &mut Vec<u8>,
+    ) -> AnyResult<()> {
+        self.cursor.serialize_key_fields(fields, dst)
+    }
+
+    fn serialize_key_to_arrow(&mut self, dst: &mut ArrayBuilder) -> AnyResult<()> {
+        self.cursor.serialize_key_to_arrow(dst)
+    }
+
+    fn serialize_key_to_arrow_with_metadata(
+        &mut self,
+        metadata: &dyn erased_serde::Serialize,
+        dst: &mut ArrayBuilder,
+    ) -> AnyResult<()> {
+        self.cursor
+            .serialize_key_to_arrow_with_metadata(metadata, dst)
+    }
+
+    fn serialize_val_to_arrow(&mut self, dst: &mut ArrayBuilder) -> AnyResult<()> {
+        self.cursor.serialize_val_to_arrow(dst)
+    }
+
+    fn serialize_val_to_arrow_with_metadata(
+        &mut self,
+        metadata: &dyn erased_serde::Serialize,
+        dst: &mut ArrayBuilder,
+    ) -> AnyResult<()> {
+        self.cursor
+            .serialize_val_to_arrow_with_metadata(metadata, dst)
+    }
+
+    #[cfg(feature = "with-avro")]
+    fn key_to_avro(&mut self, schema: &AvroSchema, refs: &NamesRef<'_>) -> AnyResult<AvroValue> {
+        self.cursor.key_to_avro(schema, refs)
+    }
+
+    fn serialize_key_weight(&mut self, dst: &mut Vec<u8>) -> AnyResult<()> {
+        self.cursor.serialize_key_weight(dst)
+    }
+
+    fn serialize_val(&mut self, dst: &mut Vec<u8>) -> AnyResult<()> {
+        self.cursor.serialize_val(dst)
+    }
+
+    fn val_to_json(&mut self) -> AnyResult<serde_json::Value> {
+        self.cursor.val_to_json()
+    }
+
+    #[cfg(feature = "with-avro")]
+    fn val_to_avro(&mut self, schema: &AvroSchema, refs: &NamesRef<'_>) -> AnyResult<AvroValue> {
+        self.cursor.val_to_avro(schema, refs)
+    }
+
+    fn weight(&mut self) -> i64 {
+        self.cursor.weight()
+    }
+
+    fn step_key(&mut self) {
+        self.cursor.step_key();
+    }
+
+    fn step_val(&mut self) {
+        self.cursor.step_val();
+    }
+
+    fn rewind_keys(&mut self) {
+        self.cursor.rewind_keys();
+        self.cursor.seek_key_exact(self.start_key.as_data());
+    }
+
+    fn rewind_vals(&mut self) {
+        self.cursor.rewind_vals();
+    }
+
+    fn seek_key_exact(&mut self, key: &DynData) -> bool {
+        self.cursor.seek_key_exact(key)
+    }
+}
+
 /// Cursor that allows serializing the contents of a type-erased batch.
 ///
 /// This is a wrapper around the DBSP `Cursor` trait that yields keys and values
@@ -471,11 +685,11 @@ pub trait SerBatchReaderHandle: Send + Sync + DynClone {
     fn num_nonempty_mailboxes(&self) -> usize;
 
     /// Like [`OutputHandle::take_from_worker`](`dbsp::OutputHandle::take_from_worker`),
-    /// but returns output batch as a [`SyncSerBatchReader`] trait object.
+    /// but returns output batch as a [`SerBatchReader`] trait object.
     fn take_from_worker(&self, worker: usize) -> Option<Box<dyn SerBatchReader>>;
 
     /// Like [`OutputHandle::take_from_all`](`dbsp::OutputHandle::take_from_all`),
-    /// but returns output batches as [`SyncSerBatchReader`] trait objects.
+    /// but returns output batches as [`SerBatchReader`] trait objects.
     fn take_from_all(&self) -> Vec<Arc<dyn SerBatchReader>>;
 
     /// Concatenate outputs from all workers into a single batch reader.