From a44dd81f46e0f7c239743e5b2a2e3e9affc65d8c Mon Sep 17 00:00:00 2001 From: Georgi Krastev Date: Mon, 22 Jul 2024 11:00:34 +0300 Subject: [PATCH] Add NullState::is_valid and NullState::is_null --- .../aggregate/groups_accumulator/accumulate.rs | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-expr-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/physical-expr-common/src/aggregate/groups_accumulator/accumulate.rs index 3fcd570f514e..cb3666bd0439 100644 --- a/datafusion/physical-expr-common/src/aggregate/groups_accumulator/accumulate.rs +++ b/datafusion/physical-expr-common/src/aggregate/groups_accumulator/accumulate.rs @@ -324,6 +324,18 @@ impl NullState { } } + /// Check if the accumulated value for the group at the given `index` is valid, + /// meaning that there was at least one value passing the filter for this group. + pub fn is_valid(&self, index: usize) -> bool { + self.seen_values.get_bit(index) + } + + /// Check if the accumulated value for the group at the given `index` is `null`, + /// meaning that no values passing the filter were seen yet for this group. + pub fn is_null(&self, index: usize) -> bool { + !self.is_valid(index) + } + /// Creates the a [`NullBuffer`] representing which group_indices /// should have null values (because they never saw any values) /// for the `emit_to` rows. @@ -828,9 +840,12 @@ mod test { // Validate the final buffer (one value per group) let expected_null_buffer = mock.expected_null_buffer(total_num_groups); + for (i, expected) in expected_null_buffer.iter().enumerate() { + assert_eq!(expected, null_state.is_valid(i)); + assert_eq!(!expected, null_state.is_null(i)) + } let null_buffer = null_state.build(EmitTo::All); - assert_eq!(null_buffer, expected_null_buffer); } }