feat!: partition data to payloads (#201)

* feat!: partition data to payloads

* feat!: unit test

* feat!: unit test

Author: gangliao

.codecov.yml

@@ -11,5 +11,6 @@ codecov:
 
 ignore:
   - "src/legacy"
+  - "src/function/src/aws/lambda.rs"
   - "src/driver/src/build"
   - "**/*.md"

src/driver/src/deploy/mod.rs

@@ -74,7 +74,7 @@ impl ExecutionEnvironment {
         for (_, ctx) in flow.ctx.iter() {
             let _: Vec<_> = lambda::function_name(&ctx)
                 .iter()
-                .map(move |name| async move {
+                .map(|name| async move {
                     client
                         .create_function(CreateFunctionRequest {
                             code: lambda::function_code(),

src/function/Cargo.toml

@@ -14,10 +14,14 @@ snmalloc = [ "snmalloc-rs" ]
 arrow = { git = "https://github.com/DSLAM-UMD/arrow", branch = "scq", features = [ "simd" ] }
 aws_lambda_events = "0.4"
 datafusion = { git = "https://github.com/DSLAM-UMD/arrow", branch = "scq", features = [ "simd" ] }
+futures = "0.3.12"
 lambda_runtime = { git = "https://github.com/awslabs/aws-lambda-rust-runtime/", branch = "master" }
+log = "0.4.14"
 rayon = "1.5"
-
 runtime = { path = "../runtime" }
+rusoto_core = "0.46.0"
+
+rusoto_lambda = "0.46.0"
 serde_json = "1.0"
 snmalloc-rs = { version = "0.2", optional = true, features = [ "cache-friendly" ] }
 tokio = { version = "1.2", features = [ "macros", "io-util", "sync", "rt-multi-thread" ] }

src/function/src/aws/lambda.rs

@@ -14,11 +14,16 @@
 
 //! The generic lambda function for sub-plan execution on AWS Lambda.
 
+use arrow::record_batch::RecordBatch;
 use aws_lambda_events::event::kafka::KafkaEvent;
 use aws_lambda_events::event::kinesis::KinesisEvent;
 use datafusion::physical_plan::Partitioning;
+use futures::executor::block_on;
 use lambda_runtime::{handler_fn, Context};
+use log::warn;
 use runtime::prelude::*;
+use rusoto_core::Region;
+use rusoto_lambda::{InvokeAsyncRequest, Lambda, LambdaClient};
 use serde_json::Value;
 use std::cell::Cell;
 use std::sync::Once;
@@ -78,6 +83,42 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+/// Invoke functions in the next stage of the data flow.
+fn invoke_async_functions(ctx: &ExecutionContext, batches: &mut Vec<RecordBatch>) -> Result<()> {
+    // retrieve the next lambda function names
+    let next_func = LambdaExecutor::next_function(&ctx)?;
+
+    // create uuid builder to assign id to each payload
+    let mut uuid_builder = UuidBuilder::new(&ctx.name, batches.len());
+
+    let client = &LambdaClient::new(Region::default());
+    let nums = batches.len();
+    (0..nums).for_each(|_| {
+        let uuid = uuid_builder.next();
+        // call the lambda function asynchronously until it succeeds.
+        loop {
+            let request = InvokeAsyncRequest {
+                function_name: next_func.clone(),
+                invoke_args:   Payload::to_bytes(&[batches.pop().unwrap()], uuid.clone()),
+            };
+
+            if let Ok(reponse) = block_on(client.invoke_async(request)) {
+                if let Some(code) = reponse.status {
+                    // A success response (202 Accepted) indicates that the request
+                    // is queued for invocation.
+                    if code == 202 {
+                        break;
+                    } else {
+                        warn!("Unknown invoke error: {}, retry ... ", code);
+                    }
+                }
+            }
+        }
+    });
+
+    Ok(())
+}
+
 async fn source_handler(ctx: &mut ExecutionContext, event: Value) -> Result<Value> {
     let batch = match &ctx.datasource {
         DataSource::KinesisEvent(_) => {
@@ -142,19 +183,28 @@ async fn source_handler(ctx: &mut ExecutionContext, event: Value) -> Result<Valu
             LambdaExecutor::event_sink(vec![batches]).await
         }
         ExecutionStrategy::Distributed => {
-            unimplemented!();
+            let mut batches = LambdaExecutor::coalesce_batches(
+                vec![batch],
+                globals["lambda"]["payload_batch_size"]
+                    .parse::<usize>()
+                    .unwrap(),
+            )
+            .await?;
+            assert_eq!(1, batches.len());
+
+            invoke_async_functions(&ctx, &mut batches[0])?;
+            Ok(serde_json::to_value(&ctx.name)?)
         }
     }
 }
 
 async fn payload_handler(ctx: &mut ExecutionContext, event: Value) -> Result<Value> {
-    let (batch, uuid) = Payload::to_batch(event);
-    let schema = batch.schema();
+    let (batches, uuid) = Payload::to_batch(event);
 
-    ctx.feed_one_source(&vec![vec![batch]]);
+    ctx.feed_one_source(&vec![batches]);
     let batches = ctx.execute().await?;
 
-    Ok(Payload::from(&batches[0], schema, uuid))
+    Ok(Payload::to_value(&batches, uuid))
 }
 
 async fn handler(event: Value, _: Context) -> Result<Value> {
@@ -259,12 +309,12 @@ mod tests {
             let res = handler(event, Context::default()).await?;
 
             // check the result of function execution
-            let (batch, _) = Payload::to_batch(res);
+            let (batches, _) = Payload::to_batch(res);
 
             if i == 0 {
                 println!(
                     "{}",
-                    arrow::util::pretty::pretty_format_batches(&[batch]).unwrap(),
+                    arrow::util::pretty::pretty_format_batches(&batches).unwrap(),
                 );
             }
         }

src/runtime/Cargo.toml

@@ -14,11 +14,13 @@ async-trait = "0.1.42"
 aws_lambda_events = "0.4"
 base64 = "0.13.0"
 blake2 = "0.9"
+bytes = "1.0.1"
 datafusion = { git = "https://github.com/DSLAM-UMD/arrow", branch = "scq", features = [ "simd" ] }
 futures = "0.3.12"
 json = "0.12.4"
 lazy_static = "1.4"
 lz4 = "1.23.1"
+rand = "0.8.3"
 rayon = "1.5"
 rusoto_core = "0.46.0"
 rusoto_kafka = "0.46.0"

src/runtime/src/config/squirtle.toml

@@ -9,9 +9,12 @@ production = false
 
 name = "execution_context"
 
-# default target batch size
+# default target batch size (16 KB)
 target_batch_size = 16384
 
+# default raw record batch size in the payload (512 KB)
+payload_batch_size = 524288
+
 # multi-thread parallelism inside the cloud function
 parallelism = 8
 

src/runtime/src/executor/mod.rs

@@ -25,8 +25,9 @@
 //! distributed dataflow model.
 
 use crate::config::GLOBALS as globals;
+use crate::context::CloudFunction;
 use crate::context::ExecutionContext;
-use crate::error::Result;
+use crate::error::{Result, SquirtleError};
 use crate::payload::{Payload, Uuid};
 use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
@@ -36,6 +37,7 @@ use datafusion::physical_plan::repartition::RepartitionExec;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
 use futures::stream::StreamExt;
 use plan::*;
+use rand::Rng;
 use rayon::prelude::*;
 use serde_json::Value;
 use std::sync::Arc;
@@ -121,11 +123,7 @@ pub trait Executor {
         assert_eq!(1, output_partitions.len());
         assert_eq!(1, output_partitions[0].len());
 
-        Ok(Payload::from(
-            &output_partitions[0][0],
-            output_partitions[0][0].schema(),
-            Uuid::default(),
-        ))
+        Ok(Payload::to_value(&output_partitions[0], Uuid::default()))
     }
 }
 
@@ -178,14 +176,43 @@ impl LambdaExecutor {
             ExecutionStrategy::Distributed
         }
     }
+
+    /// Returns the next cloud function names for invocation.
+    pub fn next_function(ctx: &ExecutionContext) -> Result<String> {
+        let mut lambdas = match &ctx.next {
+            CloudFunction::None => vec![],
+            CloudFunction::Chorus((name, num)) => {
+                (0..*num).map(|i| format!("{}-{}", name, i)).collect()
+            }
+            CloudFunction::Solo(name) => vec![name.to_owned()],
+        };
+
+        if lambdas.is_empty() {
+            return Err(SquirtleError::Internal(
+                "No distributed execution plan".to_owned(),
+            ));
+        }
+
+        let mut function_name = lambdas[0].clone();
+        if lambdas.len() > 1 {
+            // mapping to the same lambda function name through hashing technology.
+            let mut rng = rand::thread_rng();
+            function_name = lambdas.remove(rng.gen_range(0..lambdas.len()));
+        }
+
+        Ok(function_name)
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::datasource::{kinesis, DataSource};
     use crate::error::SquirtleError;
     use arrow::array::UInt32Array;
     use arrow::datatypes::{DataType, Field, Schema};
+    use aws_lambda_events::event::kinesis::KinesisEvent;
+    use datafusion::datasource::MemTable;
     use datafusion::physical_plan::expressions::Column;
     use tokio::task::JoinHandle;
 
@@ -335,6 +362,62 @@ mod tests {
         )
         .unwrap()
     }
+
+    #[tokio::test]
+    async fn next_function() -> Result<()> {
+        let input = include_str!("../../../test/data/example-kinesis-event-1.json");
+        let input: KinesisEvent = serde_json::from_str(input).unwrap();
+        let partitions = vec![kinesis::to_batch(input)];
+
+        let mut ctx = datafusion::execution::context::ExecutionContext::new();
+        let provider = MemTable::try_new(partitions[0][0].schema(), partitions.clone())?;
+
+        ctx.register_table("test", Arc::new(provider));
+
+        let sql = "SELECT MAX(c1), MIN(c2), c3 FROM test WHERE c2 < 99 GROUP BY c3";
+        let logical_plan = ctx.create_logical_plan(&sql)?;
+        let logical_plan = ctx.optimize(&logical_plan)?;
+        let physical_plan = ctx.create_physical_plan(&logical_plan)?;
+
+        // Serialize the physical plan and skip its record batches
+        let plan = serde_json::to_string(&physical_plan)?;
+
+        // Deserialize the physical plan that doesn't contain record batches
+        let plan: Arc<dyn ExecutionPlan> = serde_json::from_str(&plan)?;
+
+        // Feed record batches back to the plan
+        let mut ctx = ExecutionContext {
+            plan:       plan.clone(),
+            name:       "test".to_string(),
+            next:       CloudFunction::None,
+            datasource: DataSource::UnknownEvent,
+        };
+        LambdaExecutor::next_function(&ctx).expect_err("No distributed execution plan");
+
+        ctx = ExecutionContext {
+            plan:       plan.clone(),
+            name:       "test".to_string(),
+            next:       CloudFunction::Solo("solo".to_string()),
+            datasource: DataSource::UnknownEvent,
+        };
+        assert_eq!("solo", LambdaExecutor::next_function(&ctx)?);
+
+        ctx = ExecutionContext {
+            plan:       plan.clone(),
+            name:       "test".to_string(),
+            next:       CloudFunction::Chorus(("chorus".to_string(), 24)),
+            datasource: DataSource::UnknownEvent,
+        };
+
+        let lambdas: Vec<String> = (0..100)
+            .map(|_| LambdaExecutor::next_function(&ctx).unwrap())
+            .collect();
+
+        assert_eq!(100, lambdas.len());
+        assert_ne!(lambdas.iter().min(), lambdas.iter().max());
+
+        Ok(())
+    }
 }
 
 pub mod plan;