See [CaseExpr::case_when_no_expr].
(
&self,
batch: &RecordBatch,
return_type: &DataType,
)
| 891 | |
| 892 | /// See [CaseExpr::case_when_no_expr]. |
| 893 | fn case_when_no_expr( |
| 894 | &self, |
| 895 | batch: &RecordBatch, |
| 896 | return_type: &DataType, |
| 897 | ) -> Result<ColumnarValue> { |
| 898 | let mut result_builder = ResultBuilder::new(return_type, batch.num_rows()); |
| 899 | |
| 900 | // `remainder_rows` contains the indices of the rows that need to be evaluated |
| 901 | let mut remainder_rows: ArrayRef = |
| 902 | Arc::new(UInt32Array::from_iter(0..batch.num_rows() as u32)); |
| 903 | // `remainder_batch` contains the rows themselves that need to be evaluated |
| 904 | let mut remainder_batch = Cow::Borrowed(batch); |
| 905 | |
| 906 | for i in 0..self.when_then_expr.len() { |
| 907 | // Evaluate the 'when' predicate for the remainder batch |
| 908 | // This results in a boolean array with the same length as the remaining number of rows |
| 909 | let when_predicate = &self.when_then_expr[i].0; |
| 910 | let when_value = when_predicate |
| 911 | .evaluate(&remainder_batch)? |
| 912 | .into_array(remainder_batch.num_rows())?; |
| 913 | let when_value = as_boolean_array(&when_value).map_err(|_| { |
| 914 | internal_datafusion_err!("WHEN expression did not return a BooleanArray") |
| 915 | })?; |
| 916 | |
| 917 | // If the 'when' predicate did not match any rows, continue to the next branch immediately. |
| 918 | // Only counts valid slots that are true (masked-null predicate slots are ignored) |
| 919 | // so no `prep_null_mask_filter` needed here. |
| 920 | if !when_value.has_true() { |
| 921 | continue; |
| 922 | } |
| 923 | |
| 924 | // If the 'when' predicate matched all remaining rows, there is no need to filter |
| 925 | if when_value.null_count() == 0 && !when_value.has_false() { |
| 926 | let then_expression = &self.when_then_expr[i].1; |
| 927 | let then_value = then_expression.evaluate(&remainder_batch)?; |
| 928 | result_builder.add_branch_result(&remainder_rows, then_value)?; |
| 929 | return result_builder.finish(); |
| 930 | } |
| 931 | |
| 932 | // Filter the remainder batch based on the 'when' value |
| 933 | // This results in a batch containing only the rows that need to be evaluated |
| 934 | // for the current branch |
| 935 | // Still no need to call `prep_null_mask_filter` since `create_filter` will already do |
| 936 | // this unconditionally. |
| 937 | let then_filter = create_filter(when_value, true); |
| 938 | let then_batch = filter_record_batch(&remainder_batch, &then_filter)?; |
| 939 | let then_rows = filter_array(&remainder_rows, &then_filter)?; |
| 940 | |
| 941 | let then_expression = &self.when_then_expr[i].1; |
| 942 | let then_value = then_expression.evaluate(&then_batch)?; |
| 943 | result_builder.add_branch_result(&then_rows, then_value)?; |
| 944 | |
| 945 | // If this is the last 'when' branch and there is no 'else' expression, there's no |
| 946 | // point in calculating the remaining rows. |
| 947 | if self.else_expr.is_none() && i == self.when_then_expr.len() - 1 { |
| 948 | return result_builder.finish(); |
| 949 | } |
| 950 |
no test coverage detected