[GitHub] [arrow-datafusion] stuartcarnie commented on a diff in pull request #6669: Minor: Add tests for User Defined Aggregate functions

["stuartcarnie (via GitHub)" <gi...@apache.org>]

stuartcarnie commented on code in PR #6669:
URL: https://github.com/apache/arrow-datafusion/pull/6669#discussion_r1230075514


##########
datafusion/core/tests/user_defined_aggregates.rs:
##########
@@ -82,56 +150,151 @@ async fn execute(ctx: &SessionContext, sql: &str) -> Vec<RecordBatch> {
 ///  3.0  | 1970-01-01T00:00:00.000003
 ///  2.0  | 1970-01-01T00:00:00.000002
 ///  1.0  | 1970-01-01T00:00:00.000004
+///  5.0  | 1970-01-01T00:00:00.000005
+///  5.0  | 1970-01-01T00:00:00.000005
 /// ```
-fn udaf_struct_context() -> SessionContext {
-    let value: Float64Array = vec![3.0, 2.0, 1.0].into_iter().map(Some).collect();
-    let time = TimestampNanosecondArray::from(vec![3000, 2000, 4000]);
+struct TestContext {
+    ctx: SessionContext,
+    counters: Arc<TestCounters>,
+}
+
+impl TestContext {
+    fn new() -> Self {
+        let counters = Arc::new(TestCounters::new());
+
+        let value = Float64Array::from(vec![3.0, 2.0, 1.0, 5.0, 5.0]);
+        let time = TimestampNanosecondArray::from(vec![3000, 2000, 4000, 5000, 5000]);
+
+        let batch = RecordBatch::try_from_iter(vec![
+            ("value", Arc::new(value) as _),
+            ("time", Arc::new(time) as _),
+        ])
+        .unwrap();
 
-    let batch = RecordBatch::try_from_iter(vec![
-        ("value", Arc::new(value) as _),
-        ("time", Arc::new(time) as _),
-    ])
-    .unwrap();
+        let mut ctx = SessionContext::new();
 
-    let mut ctx = SessionContext::new();
-    ctx.register_batch("t", batch).unwrap();
+        ctx.register_batch("t", batch).unwrap();
 
-    // Tell datafusion about the "first" function
-    register_aggregate(&mut ctx);
+        // Tell DataFusion about the "first" function
+        FirstSelector::register(&mut ctx);
+        // Tell DataFusion about the "time_sum" function
+        TimeSum::register(&mut ctx, Arc::clone(&counters));
 
-    ctx
+        Self { ctx, counters }
+    }
+}
+
+#[derive(Debug, Default)]
+struct TestCounters {
+    /// was update_batch called?
+    update_batch: AtomicBool,
+    /// was retract batch called?
+    retract_batch: AtomicBool,

Review Comment:
   Nice way to test expectations 👌🏻



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org