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

Author: skalwaghe-56

src/builder/analyzer.rs

@@ -255,14 +255,87 @@ fn try_merge_collector_schemas(
     schema1: &CollectorSchema,
     schema2: &CollectorSchema,
 ) -> Result<CollectorSchema> {
-    let fields = try_merge_fields_schemas(&schema1.fields, &schema2.fields)?;
+    let schema1_fields = &schema1.fields;
+    let schema2_fields = &schema2.fields;
+
+    // Create a map from field name to index in schema1
+    let field_map: HashMap<FieldName, usize> = schema1_fields
+        .iter()
+        .enumerate()
+        .map(|(i, f)| (f.name.clone(), i))
+        .collect();
+
+    let mut output_fields = Vec::new();
+    let mut next_field_id_1 = 0;
+    let mut next_field_id_2 = 0;
+
+    for (idx, field) in schema2_fields.iter().enumerate() {
+        if let Some(&idx1) = field_map.get(&field.name) {
+            if idx1 < next_field_id_1 {
+                api_bail!(
+                    "Order mismatch for field '{}': appears at position {} in first schema but expected after '{}'",
+                    field.name,
+                    idx1,
+                    schema1_fields[next_field_id_1 - 1].name
+                );
+            }
+            // Add intervening fields from schema1
+            for i in next_field_id_1..idx1 {
+                output_fields.push(schema1_fields[i].clone());
+            }
+            // Add intervening fields from schema2
+            for i in next_field_id_2..idx {
+                output_fields.push(schema2_fields[i].clone());
+            }
+            // Merge the field
+            let merged_type =
+                try_make_common_value_type(&schema1_fields[idx1].value_type, &field.value_type)?;
+            output_fields.push(FieldSchema {
+                name: field.name.clone(),
+                value_type: merged_type,
+                description: None,
+            });
+            next_field_id_1 = idx1 + 1;
+            next_field_id_2 = idx + 1;
+            // Fields not in schema1 and not UUID are added at the end
+        }
+    }
+
+    // Add remaining fields from schema1
+    for i in next_field_id_1..schema1_fields.len() {
+        output_fields.push(schema1_fields[i].clone());
+    }
+
+    // Add remaining fields from schema2
+    for i in next_field_id_2..schema2_fields.len() {
+        output_fields.push(schema2_fields[i].clone());
+    }
+
+    // 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 position of the auto_uuid field in the merged output
+                output_fields.iter().position(|f| &f.name == name1)
+            } else {
+                return api_bail!(
+                    "Different auto_uuid fields are not allowed: '{}' vs '{}'",
+                    name1,
+                    name2
+                );
+            }
+        }
+        (Some(_), None) | (None, Some(_)) => {
+            return api_bail!("Inconsistent auto_uuid field presence is not allowed");
+        }
+        (None, None) => None,
+    };
+
     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
-        },
+        fields: output_fields,
+        auto_uuid_field_idx,
     })
 }
 
@@ -803,16 +876,47 @@ 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: Vec<FieldName> =
+                    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()
+                };
+
+                // Pre-compute field index mappings for efficient evaluation
+                let field_name_to_index: HashMap<&FieldName, usize> = collector_schema
+                    .fields
+                    .iter()
+                    .enumerate()
+                    .map(|(i, f)| (&f.name, i))
+                    .collect();
+                let field_index_mapping: Vec<usize> = input_field_names
+                    .iter()
+                    .map(|field_name| {
+                        field_name_to_index
+                            .get(field_name)
+                            .copied()
+                            .unwrap_or(usize::MAX)
+                    })
+                    .collect();
+
                 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,
+                    field_index_mapping,
                     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,7 +91,12 @@ 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,
+    /// Pre-computed mapping from input field index to collector field index.
+    /// For missing fields, the value is usize::MAX.
+    pub field_index_mapping: Vec<usize>,
     /// Fingerprinter of the collector's schema. Used to decide when to reuse auto-generated UUIDs.
     pub fingerprinter: Fingerprinter,
 }

src/execution/evaluator.rs

@@ -515,6 +515,43 @@ async fn evaluate_op_scope(
                 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::<Result<Vec<_>>>()?;
+
+                // 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()];
+
+                // Use pre-computed field index mappings for O(1) field placement
+                for (i, &collector_field_idx) in op.field_index_mapping.iter().enumerate() {
+                    if collector_field_idx != usize::MAX {
+                        field_values[collector_field_idx] = input_values[i].clone();
+                    }
+                }
+
+                // Handle auto_uuid_field (assumed to be at position 0 for efficiency)
+                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]