feat: collector automatically merge and align multiple collect() called with different schema

Author: skalwaghe-56

src/builder/analyzer.rs

@@ -255,14 +255,54 @@ fn try_merge_collector_schemas(
     schema1: &CollectorSchema,
     schema2: &CollectorSchema,
 ) -> Result<CollectorSchema> {
-    let fields = try_merge_fields_schemas(&schema1.fields, &schema2.fields)?;
+    // Union all fields from both schemas
+    let mut field_map: HashMap<FieldName, EnrichedValueType> = HashMap::new();
+
+    // Add fields from schema1
+    for field in &schema1.fields {
+        field_map.insert(field.name.clone(), field.value_type.clone());
+    }
+
+    // Merge fields from schema2
+    for field in &schema2.fields {
+        if let Some(existing_type) = field_map.get(&field.name) {
+            // Try to merge types if they differ
+            let merged_type = try_make_common_value_type(existing_type, &field.value_type)?;
+            field_map.insert(field.name.clone(), merged_type);
+        } else {
+            field_map.insert(field.name.clone(), field.value_type.clone());
+        }
+    }
+
+    // Sort fields by name for consistent ordering
+    let mut fields: Vec<FieldSchema> = field_map
+        .into_iter()
+        .map(|(name, value_type)| FieldSchema {
+            name,
+            value_type,
+            description: None,
+        })
+        .collect();
+    fields.sort_by(|a, b| a.name.cmp(&b.name));
+
+    // Handle auto_uuid_field_idx
+    let auto_uuid_field_idx = match (schema1.auto_uuid_field_idx, schema2.auto_uuid_field_idx) {
+        (Some(idx1), Some(idx2)) => {
+            let name1 = &schema1.fields[idx1].name;
+            let name2 = &schema2.fields[idx2].name;
+            if name1 == name2 {
+                // Find the new index of the auto_uuid field
+                fields.iter().position(|f| &f.name == name1)
+            } else {
+                None // Different auto_uuid fields, disable
+            }
+        }
+        _ => None, // If either doesn't have it, or both don't, disable
+    };
+
     Ok(CollectorSchema {
         fields,
-        auto_uuid_field_idx: if schema1.auto_uuid_field_idx == schema2.auto_uuid_field_idx {
-            schema1.auto_uuid_field_idx
-        } else {
-            None
-        },
+        auto_uuid_field_idx,
     })
 }
 
@@ -803,16 +843,29 @@ impl AnalyzerContext {
                 let (struct_mapping, fields_schema) = analyze_struct_mapping(&op.input, op_scope)?;
                 let has_auto_uuid_field = op.auto_uuid_field.is_some();
                 let fingerprinter = Fingerprinter::default().with(&fields_schema)?;
+                let input_field_names = fields_schema.iter().map(|f| f.name.clone()).collect();
+                let collector_ref = add_collector(
+                    &op.scope_name,
+                    op.collector_name.clone(),
+                    CollectorSchema::from_fields(fields_schema, op.auto_uuid_field.clone()),
+                    op_scope,
+                )?;
+                // Get the merged collector schema after adding
+                let collector_schema: Arc<CollectorSchema> = {
+                    let scope = find_scope(&op.scope_name, op_scope)?.1;
+                    let states = scope.states.lock().unwrap();
+                    let collector = states.collectors
+                        .get(&op.collector_name)
+                        .unwrap();
+                    collector.schema.clone()
+                };
                 let collect_op = AnalyzedReactiveOp::Collect(AnalyzedCollectOp {
                     name: reactive_op.name.clone(),
                     has_auto_uuid_field,
                     input: struct_mapping,
-                    collector_ref: add_collector(
-                        &op.scope_name,
-                        op.collector_name.clone(),
-                        CollectorSchema::from_fields(fields_schema, op.auto_uuid_field.clone()),
-                        op_scope,
-                    )?,
+                    input_field_names,
+                    collector_schema,
+                    collector_ref,
                     fingerprinter,
                 });
                 async move { Ok(collect_op) }.boxed()

src/builder/plan.rs

@@ -1,4 +1,5 @@
 use crate::base::schema::FieldSchema;
+use crate::base::spec::FieldName;
 use crate::prelude::*;
 
 use crate::ops::interface::*;
@@ -90,6 +91,8 @@ pub struct AnalyzedCollectOp {
     pub name: String,
     pub has_auto_uuid_field: bool,
     pub input: AnalyzedStructMapping,
+    pub input_field_names: Vec<FieldName>,
+    pub collector_schema: Arc<schema::CollectorSchema>,
     pub collector_ref: AnalyzedCollectorReference,
     /// Fingerprinter of the collector's schema. Used to decide when to reuse auto-generated UUIDs.
     pub fingerprinter: Fingerprinter,

src/execution/evaluator.rs

@@ -483,26 +483,36 @@ async fn evaluate_op_scope(
             }
 
             AnalyzedReactiveOp::Collect(op) => {
-                let mut field_values = Vec::with_capacity(
-                    op.input.fields.len() + if op.has_auto_uuid_field { 1 } else { 0 },
-                );
-                let field_values_iter = assemble_input_values(&op.input.fields, scoped_entries);
-                if op.has_auto_uuid_field {
-                    field_values.push(value::Value::Null);
-                    field_values.extend(field_values_iter);
-                    let uuid = memory.next_uuid(
-                        op.fingerprinter
-                            .clone()
-                            .with(&field_values[1..])?
-                            .into_fingerprint(),
-                    )?;
-                    field_values[0] = value::Value::Basic(value::BasicValue::Uuid(uuid));
-                } else {
-                    field_values.extend(field_values_iter);
-                };
                 let collector_entry = scoped_entries
                     .headn(op.collector_ref.scope_up_level as usize)
                     .ok_or_else(|| anyhow::anyhow!("Collector level out of bound"))?;
+
+                // Assemble input values
+                let input_values: Vec<value::Value> = assemble_input_values(&op.input.fields, scoped_entries).collect();
+
+                // Create field_values vector for all fields in the merged schema
+                let mut field_values: Vec<value::Value> = vec![value::Value::Null; op.collector_schema.fields.len()];
+
+                // Map input fields to their positions in the merged schema
+                for (i, field_name) in op.input_field_names.iter().enumerate() {
+                    if let Some(pos) = op.collector_schema.fields.iter().position(|f| &f.name == field_name) {
+                        field_values[pos] = input_values[i].clone();
+                    }
+                }
+
+                // Handle auto_uuid_field
+                if op.has_auto_uuid_field {
+                    if let Some(uuid_idx) = op.collector_schema.auto_uuid_field_idx {
+                        let uuid = memory.next_uuid(
+                            op.fingerprinter
+                                .clone()
+                                .with(&field_values.iter().enumerate().filter(|(i, _)| *i != uuid_idx).map(|(_, v)| v).collect::<Vec<_>>())?
+                                .into_fingerprint(),
+                        )?;
+                        field_values[uuid_idx] = value::Value::Basic(value::BasicValue::Uuid(uuid));
+                    }
+                }
+
                 {
                     let mut collected_records = collector_entry.collected_values
                         [op.collector_ref.local.collector_idx as usize]