Support joining updating data

Change the JoinWithExpiration stream node to operate on updating data.
ArroyoSystems · Nov 28, 2023 · b6c48ee · b6c48ee
1 parent 760fc5f
commit b6c48ee
Show file tree

Hide file tree

Showing 13 changed files with 1,154 additions and 331 deletions.
diff --git a/arroyo-datastream/src/lib.rs b/arroyo-datastream/src/lib.rs
@@ -1789,27 +1789,89 @@ impl Program {
 
                     let in_k = parse_type(&inputs[0].weight().key);
                     let in_t1 = parse_type(&inputs[0].weight().value);
+
+                    let (t1_updating, in_t1_inner) = match extract_container_type("UpdatingData", &in_t1) {
+                        Some(t) => (true, t),
+                        None => (false, in_t1.clone()),
+                    };
+
                     let in_t2 = parse_type(&inputs[1].weight().value);
+
+                    let (t2_updating, in_t2_inner) = match extract_container_type("UpdatingData", &in_t2) {
+                        Some(t) => (true, t),
+                        None => (false, in_t2.clone()),
+                    };
+
                     let left_expiration = duration_to_syn_expr(*left_expiration);
                     let right_expiration = duration_to_syn_expr(*right_expiration);
-                    match join_type {
-                        arroyo_types::JoinType::Inner => quote!{
-                            Box::new(arroyo_worker::operators::join_with_expiration::
+
+                    match (join_type, t1_updating, t2_updating) {
+                        (JoinType::Inner, false, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
                                 inner_join::<#in_k, #in_t1, #in_t2>(#left_expiration, #right_expiration))
                         },
-                        arroyo_types::JoinType::Left => quote!{
-                            Box::new(arroyo_worker::operators::join_with_expiration::
+                        (JoinType::Inner, true, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                inner_join_left_updating::<#in_k, #in_t1_inner, #in_t2>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Inner, false, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                inner_join_right_updating::<#in_k, #in_t1, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Inner, true, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                inner_join_both_updating::<#in_k, #in_t1_inner, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Left, false, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
                                 left_join::<#in_k, #in_t1, #in_t2>(#left_expiration, #right_expiration))
                         },
-                        arroyo_types::JoinType::Right => quote!{
-                            Box::new(arroyo_worker::operators::join_with_expiration::
+                        (JoinType::Left, true, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                left_join_left_updating::<#in_k, #in_t1_inner, #in_t2>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Left, false, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                left_join_right_updating::<#in_k, #in_t1, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Left, true, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                left_join_both_updating::<#in_k, #in_t1_inner, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Right, false, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
                                 right_join::<#in_k, #in_t1, #in_t2>(#left_expiration, #right_expiration))
                         },
-                        arroyo_types::JoinType::Full => quote!{
-                            Box::new(arroyo_worker::operators::join_with_expiration::
+                        (JoinType::Right, true, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                right_join_left_updating::<#in_k, #in_t1_inner, #in_t2>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Right, false, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                right_join_right_updating::<#in_k, #in_t1, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Right, true, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                right_join_both_updating::<#in_k, #in_t1_inner, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Full, false, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
                                 full_join::<#in_k, #in_t1, #in_t2>(#left_expiration, #right_expiration))
                         },
+                        (JoinType::Full, true, false) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                full_join_left_updating::<#in_k, #in_t1_inner, #in_t2>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Full, false, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                full_join_right_updating::<#in_k, #in_t1, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
+                        (JoinType::Full, true, true) => quote! {
+                            Box::new(arroyo_worker::operators::joiners::
+                                full_join_both_updating::<#in_k, #in_t1_inner, #in_t2_inner>(#left_expiration, #right_expiration))
+                        },
                     }
+
                 },
                 Operator::UpdatingOperator { name, expression } => {
                     let expr : syn::Expr = parse_str(expression).expect(expression);

diff --git a/arroyo-sql-testing/golden_outputs/updating_full_join.json b/arroyo-sql-testing/golden_outputs/updating_full_join.json
@@ -0,0 +1,7 @@
+{"before":null,"after":{"left_counter":0,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":null,"after":{"left_counter":2,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":{"left_counter":1,"counter_mod_2":null,"right_count":null},"after":{"left_counter":1,"counter_mod_2":0,"right_count":1},"op":"u"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":1,"right_count":1},"op":"c"}
+{"before":{"left_counter":2,"counter_mod_2":null,"right_count":null},"after":{"left_counter":2,"counter_mod_2":0,"right_count":2},"op":"u"}
+{"before":{"left_counter":1,"counter_mod_2":0,"right_count":1},"after":null,"op":"d"}
diff --git a/arroyo-sql-testing/golden_outputs/updating_inner_join.json b/arroyo-sql-testing/golden_outputs/updating_inner_join.json
@@ -0,0 +1,4 @@
+{"before":null,"after":{"left_counter":1,"counter_mod_2":0,"right_count":1},"op":"c"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":1,"right_count":1},"op":"c"}
+{"before":null,"after":{"left_counter":2,"counter_mod_2":0,"right_count":2},"op":"c"}
+{"before":{"left_counter":1,"counter_mod_2":0,"right_count":1},"after":null,"op":"d"}
diff --git a/arroyo-sql-testing/golden_outputs/updating_left_join.json b/arroyo-sql-testing/golden_outputs/updating_left_join.json
@@ -0,0 +1,7 @@
+{"before":null,"after":{"left_counter":0,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":null,"after":{"left_counter":2,"counter_mod_2":null,"right_count":null},"op":"c"}
+{"before":{"left_counter":1,"counter_mod_2":null,"right_count":null},"after":{"left_counter":1,"counter_mod_2":0,"right_count":1},"op":"u"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":1,"right_count":1},"op":"c"}
+{"before":{"left_counter":2,"counter_mod_2":null,"right_count":null},"after":{"left_counter":2,"counter_mod_2":0,"right_count":2},"op":"u"}
+{"before":{"left_counter":1,"counter_mod_2":0,"right_count":1},"after":null,"op":"d"}
diff --git a/arroyo-sql-testing/golden_outputs/updating_right_join.json b/arroyo-sql-testing/golden_outputs/updating_right_join.json
@@ -0,0 +1,4 @@
+{"before":null,"after":{"left_counter":1,"counter_mod_2":0,"right_count":1},"op":"c"}
+{"before":null,"after":{"left_counter":1,"counter_mod_2":1,"right_count":1},"op":"c"}
+{"before":null,"after":{"left_counter":2,"counter_mod_2":0,"right_count":2},"op":"c"}
+{"before":{"left_counter":1,"counter_mod_2":0,"right_count":1},"after":null,"op":"d"}
diff --git a/arroyo-sql-testing/src/smoke_tests.rs b/arroyo-sql-testing/src/smoke_tests.rs
@@ -869,3 +869,123 @@ FROM (SELECT TUMBLE(interval '1 second'),  counter, count(*) FROM impulse_source
 
 JOIN (SELECT TUMBLE(interval '1 second') as window, counter , count(*) FROM delayed_impulse_source GROUP BY 1,2) b
 ON a.counter = b.counter;"}
+
+correctness_run_codegen! {"updating_left_join", 10,
+"CREATE TABLE impulse (
+  timestamp TIMESTAMP,
+  counter bigint unsigned not null,
+  subtask_index bigint unsigned not null
+) WITH (
+  connector = 'single_file',
+  path = '$input_dir/impulse.json',
+  format = 'json',
+  type = 'source',
+  event_time_field = 'timestamp'
+);
+
+
+CREATE TABLE output (
+  left_counter bigint,
+  counter_mod_2 bigint,
+  right_count bigint
+) WITH (
+  connector = 'single_file',
+  path = '$output_path',
+  format = 'debezium_json',
+  type = 'sink'
+);
+
+INSERT INTO output
+select counter as left_counter, counter_mod_2, right_count from impulse left join
+     (select counter % 2 as counter_mod_2, cast(count(*) as bigint UNSIGNED) as right_count from impulse where counter < 3 group by 1)
+    on counter = right_count where counter < 3;"}
+
+correctness_run_codegen! {"updating_right_join", 10,
+"CREATE TABLE impulse (
+  timestamp TIMESTAMP,
+  counter bigint unsigned not null,
+  subtask_index bigint unsigned not null
+) WITH (
+  connector = 'single_file',
+  path = '$input_dir/impulse.json',
+  format = 'json',
+  type = 'source',
+  event_time_field = 'timestamp'
+);
+
+
+CREATE TABLE output (
+  left_counter bigint,
+  counter_mod_2 bigint,
+  right_count bigint
+) WITH (
+  connector = 'single_file',
+  path = '$output_path',
+  format = 'debezium_json',
+  type = 'sink'
+);
+
+INSERT INTO output
+select counter as left_counter, counter_mod_2, right_count from impulse right join
+     (select counter % 2 as counter_mod_2, cast(count(*) as bigint UNSIGNED) as right_count from impulse where counter < 3 group by 1)
+    on counter = right_count where counter < 3;"}
+
+correctness_run_codegen! {"updating_inner_join", 10,
+"CREATE TABLE impulse (
+  timestamp TIMESTAMP,
+  counter bigint unsigned not null,
+  subtask_index bigint unsigned not null
+) WITH (
+  connector = 'single_file',
+  path = '$input_dir/impulse.json',
+  format = 'json',
+  type = 'source',
+  event_time_field = 'timestamp'
+);
+
+
+CREATE TABLE output (
+  left_counter bigint,
+  counter_mod_2 bigint,
+  right_count bigint
+) WITH (
+  connector = 'single_file',
+  path = '$output_path',
+  format = 'debezium_json',
+  type = 'sink'
+);
+
+INSERT INTO output
+select counter as left_counter, counter_mod_2, right_count from impulse inner join
+     (select counter % 2 as counter_mod_2, cast(count(*) as bigint UNSIGNED) as right_count from impulse where counter < 3 group by 1)
+    on counter = right_count where counter < 3;"}
+
+correctness_run_codegen! {"updating_full_join", 10,
+"CREATE TABLE impulse (
+  timestamp TIMESTAMP,
+  counter bigint unsigned not null,
+  subtask_index bigint unsigned not null
+) WITH (
+  connector = 'single_file',
+  path = '$input_dir/impulse.json',
+  format = 'json',
+  type = 'source',
+  event_time_field = 'timestamp'
+);
+
+
+CREATE TABLE output (
+  left_counter bigint,
+  counter_mod_2 bigint,
+  right_count bigint
+) WITH (
+  connector = 'single_file',
+  path = '$output_path',
+  format = 'debezium_json',
+  type = 'sink'
+);
+
+INSERT INTO output
+select counter as left_counter, counter_mod_2, right_count from impulse full outer join
+     (select counter % 2 as counter_mod_2, cast(count(*) as bigint UNSIGNED) as right_count from impulse where counter < 3 group by 1)
+    on counter = right_count where counter < 3;"}
diff --git a/arroyo-sql/src/code_gen.rs b/arroyo-sql/src/code_gen.rs
@@ -227,9 +227,10 @@ impl JoinPairContext {
         let merge_expr = code_generator.generate(self);
         let left_ident = self.left_ident();
         let right_ident = self.right_ident();
+
         parse_quote!({
-            let #left_ident = &record.value.0;
-            let #right_ident = &record.value.1;
+            let (#left_ident, #right_ident) = &record.value.unwrap_append().clone();
+
             arroyo_types::Record {
                 timestamp: record.timestamp.clone(),
                 key: None,

diff --git a/arroyo-sql/src/pipeline.rs b/arroyo-sql/src/pipeline.rs
@@ -222,6 +222,12 @@ pub struct JoinOperator {
     pub join_type: JoinType,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct InputsUpdating {
+    pub left: bool,
+    pub right: bool,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum JoinType {
     /// Inner Join
@@ -805,9 +811,6 @@ impl<'a> SqlPipelineBuilder<'a> {
     fn insert_join(&mut self, join: &datafusion_expr::logical_plan::Join) -> Result<SqlOperator> {
         let left_input = self.insert_sql_plan(&join.left)?;
         let right_input = self.insert_sql_plan(&join.right)?;
-        if left_input.is_updating() || right_input.is_updating() {
-            bail!("don't support joins with updating inputs");
-        }
         match join.join_constraint {
             JoinConstraint::On => {}
             JoinConstraint::Using => bail!("don't support 'using' in joins"),