less poisoning

cozo-core/benches/pokec.rs

@@ -198,7 +198,7 @@ const UPDATE_QUERIES: [QueryFn; 1] = [single_vertex_update];
 const AGGREGATE_QUERIES: [QueryFn; 4] = [
     aggregation,
     aggregation_filter,
-    aggregation_distinct,
+    aggregation_count,
     min_max,
 ];
 const ANALYTICAL_QUERIES: [QueryFn; 15] = [
@@ -292,16 +292,16 @@ fn aggregation() {
         .unwrap();
 }
 
-fn aggregation_distinct() {
+fn aggregation_count() {
     TEST_DB
-        .run_script("?[count_unique(age)] := *user{age}", Default::default())
+        .run_script("?[count(uid), count(age)] := *user{uid, age}", Default::default())
         .unwrap();
 }
 
 fn aggregation_filter() {
     TEST_DB
         .run_script(
-            "?[age, count(uid)] := *user{uid, age}, try(age >= 18, false)",
+            "?[age, count(age)] := *user{age}, try(age >= 18, false)",
             Default::default(),
         )
         .unwrap();
@@ -537,7 +537,7 @@ fn bench_aggregation(b: &mut Bencher) {
 #[bench]
 fn bench_aggregation_distinct(b: &mut Bencher) {
     initialize(&TEST_DB);
-    b.iter(aggregation_distinct)
+    b.iter(aggregation_count)
 }
 
 #[bench]
@@ -779,7 +779,7 @@ fn throughput(_: &mut Bencher) {
     let aggregation_distinct_time = Instant::now();
 
     (0..count).into_par_iter().for_each(|_| {
-        aggregation_distinct();
+        aggregation_count();
     });
     dbg!((count as f64) / aggregation_distinct_time.elapsed().as_secs_f64());
 
@@ -994,7 +994,7 @@ fn mixed(_: &mut Bencher) {
     let aggregation_distinct_time = Instant::now();
 
     (0..count).into_par_iter().for_each(|_| {
-        wrap(mixed_pct, aggregation_distinct);
+        wrap(mixed_pct, aggregation_count);
     });
     dbg!((count as f64) / aggregation_distinct_time.elapsed().as_secs_f64());
 

cozo-core/src/query/eval.rs

@@ -199,8 +199,8 @@ impl<'a> SessionTx<'a> {
                             store.put(item, 0);
                         }
                         *changed.get_mut(rule_symb).unwrap() = true;
-                        poison.check()?;
                     }
+                    poison.check()?;
                 }
             }
             AggrKind::Normal => {
@@ -217,8 +217,8 @@ impl<'a> SessionTx<'a> {
                         trace!("item for {:?}.{}: {:?} at {}", rule_symb, rule_n, item, 0);
                         store_to_use.normal_aggr_put(&item, &rule.aggr, serial);
                         *changed.get_mut(rule_symb).unwrap() = true;
-                        poison.check()?;
                     }
+                    poison.check()?;
                 }
                 if store_to_use.normal_aggr_scan_and_put(
                     &ruleset[0].aggr,
@@ -311,8 +311,8 @@ impl<'a> SessionTx<'a> {
                             return Ok(true);
                         }
                     }
-                    poison.check()?;
                 }
+                poison.check()?;
             }
         }
         Ok(should_check_limit)

cozo-core/src/runtime/in_mem.rs

@@ -257,8 +257,8 @@ impl InMemRelation {
                         op.set(val)?;
                     }
                 }
-                poison.check()?;
             }
+            poison.check()?;
             for (i, aggr) in aggrs.iter().enumerate() {
                 if let Some((aggr_op, _aggr_args)) = aggr {
                     let op = aggr_op.normal_op.as_ref().unwrap();