MCPcopy Create free account
hub / github.com/apache/datafusion / invoke_with_args

Method invoke_with_args

datafusion/functions-nested/src/array_transform.rs:146–217  ·  view source on GitHub ↗
(&self, args: HigherOrderFunctionArgs)

Source from the content-addressed store, hash-verified

144 }
145
146 fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
147 let [list, lambda] = take_function_args(self.name(), &args.args)?;
148 let (ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)) = (list, lambda)
149 else {
150 return plan_err!("{} expects a value followed by a lambda", self.name());
151 };
152
153 let list_array = list.to_array(args.number_rows)?;
154
155 let list_values = match extract_list_values(&list_array, args.return_type())? {
156 ListValuesResult::EarlyReturn(v) => return Ok(v),
157 ListValuesResult::Values(v) => v,
158 };
159
160 // by passing closures, lambda.evaluate can evaluate only those actually needed
161 let values_param = || Ok(Arc::clone(&list_values));
162
163 // call the transforming lambda
164 let transformed_values = lambda
165 .evaluate(&[&values_param], |arrays| {
166 // if any column got captured, we need to adjust it to the values arrays,
167 // duplicating values of list with multitple values and removing values of empty lists
168 let indices = list_values_row_number(&list_array)?;
169 Ok(take_arrays(arrays, &indices, None)?)
170 })?
171 .into_array(list_values.len())?;
172
173 let field = match args.return_field.data_type() {
174 DataType::List(field) | DataType::LargeList(field) => Arc::clone(field),
175 _ => {
176 return exec_err!(
177 "{} expected ScalarFunctionArgs.return_field to be a list, got {}",
178 self.name(),
179 args.return_field
180 );
181 }
182 };
183
184 let transformed_list = match list_array.data_type() {
185 DataType::List(_) => {
186 let list = list_array.as_list();
187
188 // since we called list_values above which would return sliced values for
189 // a sliced list, we must adjust the offsets here as otherwise they would be invalid
190 let adjusted_offsets = adjust_offsets_for_slice(list);
191
192 Arc::new(ListArray::new(
193 field,
194 adjusted_offsets,
195 transformed_values,
196 list.nulls().cloned(),
197 )) as ArrayRef
198 }
199 DataType::LargeList(_) => {
200 let large_list = list_array.as_list();
201
202 // since we called list_values above which would return sliced values for
203 // a sliced list, we must adjust the offsets here as otherwise they would be invalid

Callers

nothing calls this directly

Calls 14

take_function_argsFunction · 0.85
extract_list_valuesFunction · 0.85
list_values_row_numberFunction · 0.85
adjust_offsets_for_sliceFunction · 0.85
newFunction · 0.85
nameMethod · 0.45
to_arrayMethod · 0.45
return_typeMethod · 0.45
into_arrayMethod · 0.45
evaluateMethod · 0.45
lenMethod · 0.45
data_typeMethod · 0.45

Tested by

no test coverage detected