feat: Async back-pressure on adding to a source (#746)

This closes #743 --------- Co-authored-by: Jordan Frazier <[email protected]>
kaskada-ai · Sep 12, 2023 · bf294dd · bf294dd
1 parent 83d4b59
commit bf294dd
Show file tree

Hide file tree

Showing 72 changed files with 728 additions and 390 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -23,6 +23,7 @@ arrow-ord = { version = "43.0.0" }
 arrow-schema = { version = "43.0.0", features = ["serde"] }
 arrow-select = { version = "43.0.0" }
 arrow-string = { version = "43.0.0" }
+async-broadcast = "0.5.1"
 async-once-cell = "0.5.3"
 async-stream = "0.3.4"
 async-trait = "0.1.68"

diff --git a/crates/sparrow-merge/Cargo.toml b/crates/sparrow-merge/Cargo.toml
@@ -19,6 +19,7 @@ arrow-array.workspace = true
 arrow-csv = { workspace = true, optional = true }
 arrow-schema.workspace = true
 arrow-select.workspace = true
+async-broadcast.workspace = true
 async-stream.workspace = true
 bit-set.workspace = true
 derive_more.workspace = true

diff --git a/crates/sparrow-merge/src/in_memory_batches.rs b/crates/sparrow-merge/src/in_memory_batches.rs
@@ -2,7 +2,7 @@ use std::sync::RwLock;
 
 use arrow_array::RecordBatch;
 use arrow_schema::SchemaRef;
-use error_stack::{IntoReport, IntoReportCompat, ResultExt};
+use error_stack::{IntoReportCompat, ResultExt};
 use futures::Stream;
 
 use crate::old::homogeneous_merge;
@@ -20,12 +20,15 @@ impl error_stack::Context for Error {}
 /// Struct for managing in-memory batches.
 #[derive(Debug)]
 pub struct InMemoryBatches {
-    retained: bool,
+    /// Whether rows added will be available for interactive queries.
+    /// If False, rows will be discarded after being sent to any active
+    /// materializations.
+    queryable: bool,
     current: RwLock<Current>,
-    updates: tokio::sync::broadcast::Sender<(usize, RecordBatch)>,
+    sender: async_broadcast::Sender<(usize, RecordBatch)>,
     /// A subscriber that is never used -- it exists only to keep the sender
     /// alive.
-    _subscriber: tokio::sync::broadcast::Receiver<(usize, RecordBatch)>,
+    _receiver: async_broadcast::InactiveReceiver<(usize, RecordBatch)>,
 }
 
 #[derive(Debug)]
@@ -61,38 +64,43 @@ impl Current {
 }
 
 impl InMemoryBatches {
-    pub fn new(retained: bool, schema: SchemaRef) -> Self {
-        let (updates, _subscriber) = tokio::sync::broadcast::channel(10);
+    pub fn new(queryable: bool, schema: SchemaRef) -> Self {
+        let (mut sender, receiver) = async_broadcast::broadcast(10);
+
+        // Don't wait for a receiver. If no-one receives, `send` will fail.
+        sender.set_await_active(false);
+
         let current = RwLock::new(Current::new(schema.clone()));
         Self {
-            retained,
+            queryable,
             current,
-            updates,
-            _subscriber,
+            sender,
+            _receiver: receiver.deactivate(),
         }
     }
 
     /// Add a batch, merging it into the in-memory version.
     ///
     /// Publishes the new batch to the subscribers.
-    pub fn add_batch(&self, batch: RecordBatch) -> error_stack::Result<(), Error> {
+    pub async fn add_batch(&self, batch: RecordBatch) -> error_stack::Result<(), Error> {
         if batch.num_rows() == 0 {
             return Ok(());
         }
 
         let new_version = {
             let mut write = self.current.write().map_err(|_| Error::Add)?;
-            if self.retained {
+            if self.queryable {
                 write.add_batch(&batch)?;
             }
             write.version += 1;
             write.version
         };
 
-        self.updates
-            .send((new_version, batch))
-            .into_report()
-            .change_context(Error::Add)?;
+        let send_result = self.sender.broadcast((new_version, batch)).await;
+        if send_result.is_err() {
+            assert!(!self.sender.is_closed());
+            tracing::info!("No-one subscribed for new batch");
+        }
         Ok(())
     }
 
@@ -107,7 +115,7 @@ impl InMemoryBatches {
             let read = self.current.read().unwrap();
             (read.version, read.batch.clone())
         };
-        let mut recv = self.updates.subscribe();
+        let mut recv = self.sender.new_receiver();
 
         async_stream::try_stream! {
             tracing::info!("Starting subscriber with version {version}");
@@ -126,11 +134,11 @@ impl InMemoryBatches {
                             tracing::warn!("Ignoring old version {recv_version}");
                         }
                     }
-                    Err(tokio::sync::broadcast::error::RecvError::Closed) => {
+                    Err(async_broadcast::RecvError::Closed) => {
                         tracing::info!("Sender closed.");
                         break;
                     },
-                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
+                    Err(async_broadcast::RecvError::Overflowed(_)) => {
                         Err(Error::ReceiverLagged)?;
                     }
                 }

diff --git a/crates/sparrow-runtime/src/key_hash_inverse.rs b/crates/sparrow-runtime/src/key_hash_inverse.rs
@@ -324,27 +324,6 @@ impl ThreadSafeKeyHashInverse {
         }
     }
 
-    pub fn blocking_add(
-        &self,
-        keys: &dyn Array,
-        key_hashes: &UInt64Array,
-    ) -> error_stack::Result<(), Error> {
-        error_stack::ensure!(
-            keys.len() == key_hashes.len(),
-            Error::MismatchedLengths {
-                keys: keys.len(),
-                key_hashes: key_hashes.len()
-            }
-        );
-        let has_new_keys = self.key_map.blocking_read().has_new_keys(key_hashes);
-
-        if has_new_keys {
-            self.key_map.blocking_write().add(keys, key_hashes)
-        } else {
-            Ok(())
-        }
-    }
-
     /// Stores the KeyHashInverse to the compute store.
     ///
     /// This method is thread-safe and acquires the read-lock.

diff --git a/crates/sparrow-runtime/src/prepare/preparer.rs b/crates/sparrow-runtime/src/prepare/preparer.rs
@@ -1,3 +1,4 @@
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use arrow::array::{ArrayRef, UInt64Array};
@@ -31,7 +32,7 @@ pub struct Preparer {
     prepared_schema: SchemaRef,
     time_column_name: String,
     subsort_column_name: Option<String>,
-    next_subsort: u64,
+    next_subsort: AtomicU64,
     key_column_name: String,
     time_multiplier: Option<i64>,
 }
@@ -51,7 +52,7 @@ impl Preparer {
             prepared_schema,
             time_column_name,
             subsort_column_name,
-            next_subsort: prepare_hash,
+            next_subsort: prepare_hash.into(),
             key_column_name,
             time_multiplier,
         })
@@ -66,10 +67,7 @@ impl Preparer {
     /// - This computes and adds the key columns.
     /// - This sorts the batch by time, subsort and key hash.
     /// - This adds or casts columns as needed.
-    ///
-    /// Self is mutated as necessary to ensure the `subsort` column is increasing, if
-    /// it is added.
-    pub fn prepare_batch(&mut self, batch: RecordBatch) -> error_stack::Result<RecordBatch, Error> {
+    pub fn prepare_batch(&self, batch: RecordBatch) -> error_stack::Result<RecordBatch, Error> {
         let time = get_required_column(&batch, &self.time_column_name)?;
         let time = cast_to_timestamp(time, self.time_multiplier)?;
 
@@ -80,8 +78,10 @@ impl Preparer {
                 .into_report()
                 .change_context_lazy(|| Error::ConvertSubsort(subsort.data_type().clone()))?
         } else {
-            let subsort: UInt64Array = (self.next_subsort..).take(num_rows).collect();
-            self.next_subsort += num_rows as u64;
+            let subsort_start = self
+                .next_subsort
+                .fetch_add(num_rows as u64, Ordering::SeqCst);
+            let subsort: UInt64Array = (subsort_start..).take(num_rows).collect();
             Arc::new(subsort)
         };
 

diff --git a/crates/sparrow-session/src/session.rs b/crates/sparrow-session/src/session.rs
@@ -87,7 +87,7 @@ impl Session {
         name: &str,
         schema: SchemaRef,
         time_column_name: &str,
-        retained: bool,
+        queryable: bool,
         subsort_column_name: Option<&str>,
         key_column_name: &str,
         grouping_name: Option<&str>,
@@ -150,7 +150,7 @@ impl Session {
             key_hash_inverse,
             key_column,
             expr,
-            retained,
+            queryable,
             time_unit,
         )
     }

diff --git a/crates/sparrow-session/src/table.rs b/crates/sparrow-session/src/table.rs
@@ -26,7 +26,7 @@ impl Table {
         key_hash_inverse: Arc<ThreadSafeKeyHashInverse>,
         key_column: usize,
         expr: Expr,
-        retained: bool,
+        queryable: bool,
         time_unit: Option<&str>,
     ) -> error_stack::Result<Self, Error> {
         let prepared_fields: Fields = KEY_FIELDS
@@ -38,7 +38,7 @@ impl Table {
         let prepare_hash = 0;
 
         assert!(table_info.in_memory.is_none());
-        let in_memory_batches = Arc::new(InMemoryBatches::new(retained, prepared_schema.clone()));
+        let in_memory_batches = Arc::new(InMemoryBatches::new(queryable, prepared_schema.clone()));
         table_info.in_memory = Some(in_memory_batches.clone());
 
         let preparer = Preparer::new(
@@ -66,7 +66,7 @@ impl Table {
         self.preparer.schema()
     }
 
-    pub fn add_data(&mut self, batch: RecordBatch) -> error_stack::Result<(), Error> {
+    pub async fn add_data(&self, batch: RecordBatch) -> error_stack::Result<(), Error> {
         let prepared = self
             .preparer
             .prepare_batch(batch)
@@ -75,11 +75,13 @@ impl Table {
         let key_hashes = prepared.column(2).as_primitive();
         let keys = prepared.column(self.key_column);
         self.key_hash_inverse
-            .blocking_add(keys.as_ref(), key_hashes)
+            .add(keys.as_ref(), key_hashes)
+            .await
             .change_context(Error::Prepare)?;
 
         self.in_memory_batches
             .add_batch(prepared)
+            .await
             .change_context(Error::Prepare)?;
         Ok(())
     }

diff --git a/examples/event-api/server.py b/examples/event-api/server.py
@@ -13,15 +13,14 @@ async def main():
 
     # Initialize event source with schema from historical data.
     events = kd.sources.PyDict(
-        rows = [],
-        schema = pa.schema([
+        schema=pa.schema([
             pa.field("ts", pa.float64()),
             pa.field("user", pa.string()),
             pa.field("request_id", pa.string()),
         ]),
-        time_column = "ts",
-        key_column = "user",
-        time_unit = "s",
+        time_column="ts",
+        key_column="user",
+        time_unit="s",
         retained=False,
     )
 

diff --git a/examples/slackbot/Notebook.ipynb b/examples/slackbot/Notebook.ipynb
@@ -119,10 +119,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import asyncio\n",
     "import pandas\n",
     "import sparrow_pi.sources as sources\n",
     "\n",
-    "messages = kt.sources.Parquet(\"./messages.parquet\", time = \"ts\", entity = \"channel\")\n",
+    "messages = await kt.sources.Parquet.create(\"./messages.parquet\", time = \"ts\", entity = \"channel\")\n",
     "messages = messages.with_key(kt.record({  # !!!\n",
     "        \"channel\": messages.col(\"channel\"),\n",
     "        \"thread\": messages.col(\"thread_ts\"),\n",

diff --git a/python/Cargo.lock b/python/Cargo.lock
diff --git a/python/docs/source/examples/time_centric.ipynb b/python/docs/source/examples/time_centric.ipynb
@@ -125,6 +125,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import asyncio\n",
+    "\n",
     "# For demo simplicity, instead of a CSV file, we read and then parse data from a\n",
     "# CSV string. Kaskadaa\n",
     "event_data_string = \"\"\"\n",
@@ -151,7 +153,7 @@
     "    ev_00020,2022-01-01 22:20:00,user_002,view_item,0\n",
     "\"\"\"\n",
     "\n",
-    "events = kd.sources.CsvString(\n",
+    "events = await kd.sources.CsvString.create(\n",
     "    event_data_string, time_column=\"event_at\", key_column=\"entity_id\"\n",
     ")\n",
     "\n",

diff --git a/python/docs/source/guide/entities.md b/python/docs/source/guide/entities.md
@@ -67,6 +67,7 @@ This is helpful since the _feature vector_ for an entity will depend only on eve
 ---
 tags: [hide-input]
 ---
+import asyncio
 import kaskada as kd
 kd.init_session()
 data = "\n".join(
@@ -79,7 +80,7 @@ data = "\n".join(
         "1996-12-23T16:40:01,A,12",
     ]
 )
-multi_entity = kd.sources.CsvString(data, time_column="time", key_column="key")
+multi_entity = await kd.sources.CsvString.create(data, time_column="time", key_column="key")
 
 kd.plot.render(
     kd.plot.Plot(multi_entity.col("m"), name="m"),

diff --git a/python/docs/source/guide/quickstart.md b/python/docs/source/guide/quickstart.md
@@ -28,6 +28,7 @@ The following Python code imports the Kaskada library, creates a session, and lo
 It then runs a query to produce a Pandas DataFrame.
 
 ```{code-cell}
+import asyncio
 import kaskada as kd
 kd.init_session()
 content = "\n".join(
@@ -41,6 +42,6 @@ content = "\n".join(
         "1996-12-19T16:40:02,A,,",
     ]
 )
-source = kd.sources.CsvString(content, time_column="time", key_column="key")
+source = await kd.sources.CsvString.create(content, time_column="time", key_column="key")
 source.select("m", "n").extend({"sum_m": source.col("m").sum() }).to_pandas()
 ```