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

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

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


##########
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,
 }
 
-fn register_aggregate(ctx: &mut SessionContext) {
-    let return_type = Arc::new(FirstSelector::output_datatype());
-    let state_type = Arc::new(FirstSelector::state_datatypes());
+impl TestCounters {
+    fn new() -> Self {
+        Default::default()
+    }
+
+    /// Has `update_batch` been called?
+    fn update_batch(&self) -> bool {
+        self.update_batch.load(Ordering::SeqCst)
+    }
+
+    /// Has `update_batch` been called?

Review Comment:
   ```suggestion
       /// Has `retract_batch` been called?
   ```



##########
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?

Review Comment:
   ```suggestion
       /// was retract_batch called?
   ```



-- 
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