datafusion_functions_aggregate/
bool_and_or.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Defines physical expressions that can evaluated at runtime during query execution
19
20use std::mem::size_of_val;
21
22use arrow::array::ArrayRef;
23use arrow::array::BooleanArray;
24use arrow::compute::bool_and as compute_bool_and;
25use arrow::compute::bool_or as compute_bool_or;
26use arrow::datatypes::Field;
27use arrow::datatypes::{DataType, FieldRef};
28
29use datafusion_common::internal_err;
30use datafusion_common::{Result, ScalarValue};
31use datafusion_common::{downcast_value, not_impl_err};
32use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
33use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name};
34use datafusion_expr::{
35    Accumulator, AggregateUDFImpl, Documentation, GroupsAccumulator, ReversedUDAF,
36    Signature, Volatility,
37};
38
39use datafusion_functions_aggregate_common::aggregate::groups_accumulator::bool_op::BooleanGroupsAccumulator;
40use datafusion_macros::user_doc;
41
42// returns the new value after bool_and/bool_or with the new values, taking nullability into account
43macro_rules! typed_bool_and_or_batch {
44    ($VALUES:expr, $ARRAYTYPE:ident, $SCALAR:ident, $OP:ident) => {{
45        let array = downcast_value!($VALUES, $ARRAYTYPE);
46        let delta = $OP(array);
47        Ok(ScalarValue::$SCALAR(delta))
48    }};
49}
50
51// bool_and/bool_or the array and returns a ScalarValue of its corresponding type.
52macro_rules! bool_and_or_batch {
53    ($VALUES:expr, $OP:ident) => {{
54        match $VALUES.data_type() {
55            DataType::Boolean => {
56                typed_bool_and_or_batch!($VALUES, BooleanArray, Boolean, $OP)
57            }
58            e => {
59                return internal_err!(
60                    "Bool and/Bool or is not expected to receive the type {e:?}"
61                );
62            }
63        }
64    }};
65}
66
67/// dynamically-typed bool_and(array) -> ScalarValue
68fn bool_and_batch(values: &ArrayRef) -> Result<ScalarValue> {
69    bool_and_or_batch!(values, compute_bool_and)
70}
71
72/// dynamically-typed bool_or(array) -> ScalarValue
73fn bool_or_batch(values: &ArrayRef) -> Result<ScalarValue> {
74    bool_and_or_batch!(values, compute_bool_or)
75}
76
77make_udaf_expr_and_func!(
78    BoolAnd,
79    bool_and,
80    expression,
81    "The values to combine with `AND`",
82    bool_and_udaf
83);
84
85make_udaf_expr_and_func!(
86    BoolOr,
87    bool_or,
88    expression,
89    "The values to combine with `OR`",
90    bool_or_udaf
91);
92
93#[user_doc(
94    doc_section(label = "General Functions"),
95    description = "Returns true if all non-null input values are true, otherwise false.",
96    syntax_example = "bool_and(expression)",
97    sql_example = r#"```sql
98> SELECT bool_and(column_name) FROM table_name;
99+----------------------------+
100| bool_and(column_name)       |
101+----------------------------+
102| true                        |
103+----------------------------+
104```"#,
105    standard_argument(name = "expression", prefix = "The")
106)]
107/// BOOL_AND aggregate expression
108#[derive(Debug, PartialEq, Eq, Hash)]
109pub struct BoolAnd {
110    signature: Signature,
111}
112
113impl BoolAnd {
114    fn new() -> Self {
115        Self {
116            signature: Signature::exact(vec![DataType::Boolean], Volatility::Immutable),
117        }
118    }
119}
120
121impl Default for BoolAnd {
122    fn default() -> Self {
123        Self::new()
124    }
125}
126
127impl AggregateUDFImpl for BoolAnd {
128    fn name(&self) -> &str {
129        "bool_and"
130    }
131
132    fn signature(&self) -> &Signature {
133        &self.signature
134    }
135
136    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
137        Ok(DataType::Boolean)
138    }
139
140    fn accumulator(&self, _: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
141        Ok(Box::<BoolAndAccumulator>::default())
142    }
143
144    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
145        Ok(vec![
146            Field::new(
147                format_state_name(args.name, self.name()),
148                DataType::Boolean,
149                true,
150            )
151            .into(),
152        ])
153    }
154
155    fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
156        true
157    }
158
159    fn create_groups_accumulator(
160        &self,
161        args: AccumulatorArgs,
162    ) -> Result<Box<dyn GroupsAccumulator>> {
163        match args.return_field.data_type() {
164            DataType::Boolean => {
165                Ok(Box::new(BooleanGroupsAccumulator::new(|x, y| x && y, true)))
166            }
167            _ => not_impl_err!(
168                "GroupsAccumulator not supported for {} with {}",
169                args.name,
170                args.return_field.data_type()
171            ),
172        }
173    }
174
175    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
176        AggregateOrderSensitivity::Insensitive
177    }
178
179    fn reverse_expr(&self) -> ReversedUDAF {
180        ReversedUDAF::Identical
181    }
182
183    fn documentation(&self) -> Option<&Documentation> {
184        self.doc()
185    }
186}
187
188#[derive(Debug, Default)]
189struct BoolAndAccumulator {
190    acc: Option<bool>,
191}
192
193impl Accumulator for BoolAndAccumulator {
194    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
195        let values = &values[0];
196        self.acc = match (self.acc, bool_and_batch(values)?) {
197            (None, ScalarValue::Boolean(v)) => v,
198            (Some(v), ScalarValue::Boolean(None)) => Some(v),
199            (Some(a), ScalarValue::Boolean(Some(b))) => Some(a && b),
200            _ => unreachable!(),
201        };
202        Ok(())
203    }
204
205    fn evaluate(&mut self) -> Result<ScalarValue> {
206        Ok(ScalarValue::Boolean(self.acc))
207    }
208
209    fn size(&self) -> usize {
210        size_of_val(self)
211    }
212
213    fn state(&mut self) -> Result<Vec<ScalarValue>> {
214        Ok(vec![ScalarValue::Boolean(self.acc)])
215    }
216
217    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
218        self.update_batch(states)
219    }
220}
221
222#[user_doc(
223    doc_section(label = "General Functions"),
224    description = "Returns true if all non-null input values are true, otherwise false.",
225    syntax_example = "bool_and(expression)",
226    sql_example = r#"```sql
227> SELECT bool_and(column_name) FROM table_name;
228+----------------------------+
229| bool_and(column_name)       |
230+----------------------------+
231| true                        |
232+----------------------------+
233```"#,
234    standard_argument(name = "expression", prefix = "The")
235)]
236/// BOOL_OR aggregate expression
237#[derive(Debug, Clone, PartialEq, Eq, Hash)]
238pub struct BoolOr {
239    signature: Signature,
240}
241
242impl BoolOr {
243    fn new() -> Self {
244        Self {
245            signature: Signature::exact(vec![DataType::Boolean], Volatility::Immutable),
246        }
247    }
248}
249
250impl Default for BoolOr {
251    fn default() -> Self {
252        Self::new()
253    }
254}
255
256impl AggregateUDFImpl for BoolOr {
257    fn name(&self) -> &str {
258        "bool_or"
259    }
260
261    fn signature(&self) -> &Signature {
262        &self.signature
263    }
264
265    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
266        Ok(DataType::Boolean)
267    }
268
269    fn accumulator(&self, _: AccumulatorArgs) -> Result<Box<dyn Accumulator>> {
270        Ok(Box::<BoolOrAccumulator>::default())
271    }
272
273    fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> {
274        Ok(vec![
275            Field::new(
276                format_state_name(args.name, self.name()),
277                DataType::Boolean,
278                true,
279            )
280            .into(),
281        ])
282    }
283
284    fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
285        true
286    }
287
288    fn create_groups_accumulator(
289        &self,
290        args: AccumulatorArgs,
291    ) -> Result<Box<dyn GroupsAccumulator>> {
292        match args.return_field.data_type() {
293            DataType::Boolean => Ok(Box::new(BooleanGroupsAccumulator::new(
294                |x, y| x || y,
295                false,
296            ))),
297            _ => not_impl_err!(
298                "GroupsAccumulator not supported for {} with {}",
299                args.name,
300                args.return_field.data_type()
301            ),
302        }
303    }
304
305    fn order_sensitivity(&self) -> AggregateOrderSensitivity {
306        AggregateOrderSensitivity::Insensitive
307    }
308
309    fn reverse_expr(&self) -> ReversedUDAF {
310        ReversedUDAF::Identical
311    }
312
313    fn documentation(&self) -> Option<&Documentation> {
314        self.doc()
315    }
316}
317
318#[derive(Debug, Default)]
319struct BoolOrAccumulator {
320    acc: Option<bool>,
321}
322
323impl Accumulator for BoolOrAccumulator {
324    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
325        let values = &values[0];
326        self.acc = match (self.acc, bool_or_batch(values)?) {
327            (None, ScalarValue::Boolean(v)) => v,
328            (Some(v), ScalarValue::Boolean(None)) => Some(v),
329            (Some(a), ScalarValue::Boolean(Some(b))) => Some(a || b),
330            _ => unreachable!(),
331        };
332        Ok(())
333    }
334
335    fn evaluate(&mut self) -> Result<ScalarValue> {
336        Ok(ScalarValue::Boolean(self.acc))
337    }
338
339    fn size(&self) -> usize {
340        size_of_val(self)
341    }
342
343    fn state(&mut self) -> Result<Vec<ScalarValue>> {
344        Ok(vec![ScalarValue::Boolean(self.acc)])
345    }
346
347    fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
348        self.update_batch(states)
349    }
350}
datafusion_functions_aggregate/bool_and_or.rs

datafusion_functions_aggregate/
bool_and_or.rs
