Repository: rcoh/angle-grinder (language: rust)

Issue:
fields after count causes panic
```sudo lsof -i -P | grep '(LISTEN)' | agrind '* | parse ":* (" as port | count by port | fields port'```

leads to a panic. Probably a bug in the pre-aggregate adapter.


Discussion / hints:

Fix patch:
diff --git a/src/lib.rs b/src/lib.rs
index 9345ff2..c14c2b6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -243,7 +243,10 @@ pub mod pipeline {
 
             // Drop tx when causes the thread to exit.
             drop(tx);
-            t.join().unwrap();
+            match t.join() {
+                Ok(_) => (),
+                Err(e) => println!("Error: {:?}", e),
+            }
         }
 
         /// Process a record using the pre-agg operators.  The output of the last operator will be
diff --git a/src/operator.rs b/src/operator.rs
index 33bf5f0..c55f369 100644
--- a/src/operator.rs
+++ b/src/operator.rs
@@ -130,7 +130,6 @@ impl AggregateOperator for PreAggAdapter {
             Row::Record(_) => panic!("PreAgg adaptor should only be used after aggregates"),
             Row::Aggregate(agg) => {
                 let mut op = self.op_builder.build();
-                let columns = agg.columns;
                 let mut processed_records: Vec<data::VMap> = {
                     let records = agg
                         .data
@@ -144,17 +143,27 @@ impl AggregateOperator for PreAggAdapter {
                     records.collect()
                 };
                 processed_records.extend(op.drain().map(|rec| rec.data));
-                let new_keys: Vec<String> = {
+                let resulting_columns: Vec<String> = {
                     processed_records
                         .iter()
                         .flat_map(|vmap| vmap.keys())
-                        .filter(|col| !columns.contains(col))
                         .unique()
                         .cloned()
                 }
                 .collect();
-                let mut columns = columns;
-                columns.extend(new_keys);
+                let output_column_set: HashSet<String> =
+                    HashSet::from_iter(resulting_columns.iter().cloned());
+                let input_column_set = HashSet::from_iter(agg.columns.iter().cloned());
+                let new_columns: Vec<String> = output_column_set
+                    .difference(&input_column_set)
+                    .cloned()
+                    .collect();
+                let mut columns = agg.columns;
+                columns.extend(new_columns);
+                let columns: Vec<String> = columns
+                    .into_iter()
+                    .filter(|col| output_column_set.contains(col))
+                    .collect();
                 self.state = Aggregate {
                     data: processed_records,
                     columns,


Test patch:
diff --git a/tests/integration.rs b/tests/integration.rs
index c84b1d7..db565a7 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -101,6 +101,11 @@ mod integration {
         structured_test(include_str!("structured_tests/total_agg.toml"));
     }
 
+    #[test]
+    fn fields_after_agg_bug() {
+        structured_test(include_str!("structured_tests/fields_after_agg.toml"));
+    }
+
     #[test]
     fn limit() {
         structured_test(include_str!("structured_tests/limit.toml"));
diff --git a/tests/structured_tests/fields_after_agg.toml b/tests/structured_tests/fields_after_agg.toml
new file mode 100644
index 0000000..051f4db
--- /dev/null
+++ b/tests/structured_tests/fields_after_agg.toml
@@ -0,0 +1,10 @@
+query = ' * | parse ":* (" as port | count by port | fields port'
+input = ": 45 ("
+output = """
+port
+------------
+45
+"""
+error = """
+"""
+succeeds = true