See [CaseExpr::case_when_with_expr].
(
&self,
batch: &RecordBatch,
return_type: &DataType,
)
| 745 | |
| 746 | /// See [CaseExpr::case_when_with_expr]. |
| 747 | fn case_when_with_expr( |
| 748 | &self, |
| 749 | batch: &RecordBatch, |
| 750 | return_type: &DataType, |
| 751 | ) -> Result<ColumnarValue> { |
| 752 | let mut result_builder = ResultBuilder::new(return_type, batch.num_rows()); |
| 753 | |
| 754 | // `remainder_rows` contains the indices of the rows that need to be evaluated |
| 755 | let mut remainder_rows: ArrayRef = |
| 756 | Arc::new(UInt32Array::from_iter_values(0..batch.num_rows() as u32)); |
| 757 | // `remainder_batch` contains the rows themselves that need to be evaluated |
| 758 | let mut remainder_batch = Cow::Borrowed(batch); |
| 759 | |
| 760 | // evaluate the base expression |
| 761 | let mut base_values = self |
| 762 | .expr |
| 763 | .as_ref() |
| 764 | .unwrap() |
| 765 | .evaluate(batch)? |
| 766 | .into_array(batch.num_rows())?; |
| 767 | |
| 768 | // Fill in a result value already for rows where the base expression value is null |
| 769 | // Since each when expression is tested against the base expression using the equality |
| 770 | // operator, null base values can never match any when expression. `x = NULL` is falsy, |
| 771 | // for all possible values of `x`. |
| 772 | let base_null_count = base_values.logical_null_count(); |
| 773 | if base_null_count > 0 { |
| 774 | // Use `is_not_null` since this is a cheap clone of the null buffer from 'base_value'. |
| 775 | // We already checked there are nulls, so we can be sure a new buffer will not be |
| 776 | // created. |
| 777 | let base_not_nulls = is_not_null(base_values.as_ref())?; |
| 778 | let base_all_null = base_null_count == remainder_batch.num_rows(); |
| 779 | |
| 780 | // If there is an else expression, use that as the default value for the null rows |
| 781 | // Otherwise the default `null` value from the result builder will be used. |
| 782 | if let Some(e) = &self.else_expr { |
| 783 | let expr = try_cast(Arc::clone(e), &batch.schema(), return_type.clone())?; |
| 784 | |
| 785 | if base_all_null { |
| 786 | // All base values were null, so no need to filter |
| 787 | let nulls_value = expr.evaluate(&remainder_batch)?; |
| 788 | result_builder.add_branch_result(&remainder_rows, nulls_value)?; |
| 789 | } else { |
| 790 | // Filter out the null rows and evaluate the else expression for those |
| 791 | let nulls_filter = create_filter(¬(&base_not_nulls)?, true); |
| 792 | let nulls_batch = |
| 793 | filter_record_batch(&remainder_batch, &nulls_filter)?; |
| 794 | let nulls_rows = filter_array(&remainder_rows, &nulls_filter)?; |
| 795 | let nulls_value = expr.evaluate(&nulls_batch)?; |
| 796 | result_builder.add_branch_result(&nulls_rows, nulls_value)?; |
| 797 | } |
| 798 | } |
| 799 | |
| 800 | // All base values are null, so we can return early |
| 801 | if base_all_null { |
| 802 | return result_builder.finish(); |
| 803 | } |
| 804 |
no test coverage detected