## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/async_scalar_function.rs

    Line
    Count
    Source
    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
    use crate::ScalarFunctionExpr;
    19
    use arrow::array::{make_array, MutableArrayData, RecordBatch};
    20
    use arrow::datatypes::{DataType, Field, FieldRef, Schema};
    21
    use datafusion_common::config::ConfigOptions;
    22
    use datafusion_common::Result;
    23
    use datafusion_common::{internal_err, not_impl_err};
    24
    use datafusion_expr::async_udf::AsyncScalarUDF;
    25
    use datafusion_expr::ScalarFunctionArgs;
    26
    use datafusion_expr_common::columnar_value::ColumnarValue;
    27
    use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
    28
    use std::any::Any;
    29
    use std::fmt::Display;
    30
    use std::hash::{Hash, Hasher};
    31
    use std::sync::Arc;
    32
    33
    /// Wrapper around a scalar function that can be evaluated asynchronously
    34
    #[derive(Debug, Clone, Eq)]
    35
    pub struct AsyncFuncExpr {
    36
        /// The name of the output column this function will generate
    37
        pub name: String,
    38
        /// The actual function (always `ScalarFunctionExpr`)
    39
        pub func: Arc<dyn PhysicalExpr>,
    40
        /// The field that this function will return
    41
        return_field: FieldRef,
    42
    }
    43
    44
    impl Display for AsyncFuncExpr {
    45
    0
        fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
    46
    0
            write!(f, "async_expr(name={}, expr={})", self.name, self.func)
    47
    0
        }
    48
    }
    49
    50
    impl PartialEq for AsyncFuncExpr {
    51
    0
        fn eq(&self, other: &Self) -> bool {
    52
    0
            self.name == other.name && self.func == Arc::clone(&other.func)
    53
    0
        }
    54
    }
    55
    56
    impl Hash for AsyncFuncExpr {
    57
    0
        fn hash<H: Hasher>(&self, state: &mut H) {
    58
    0
            self.name.hash(state);
    59
    0
            self.func.as_ref().hash(state);
    60
    0
        }
    61
    }
    62
    63
    impl AsyncFuncExpr {
    64
        /// create a new AsyncFuncExpr
    65
    0
        pub fn try_new(
    66
    0
            name: impl Into<String>,
    67
    0
            func: Arc<dyn PhysicalExpr>,
    68
    0
            schema: &Schema,
    69
    0
        ) -> Result<Self> {
    70
    0
            let Some(_) = func.as_any().downcast_ref::<ScalarFunctionExpr>() else {
    71
    0
                return internal_err!(
    72
    0
                    "unexpected function type, expected ScalarFunctionExpr, got: {:?}",
    73
                    func
    74
                );
    75
            };
    76
    77
    0
            let return_field = func.return_field(schema)?;
    78
    0
            Ok(Self {
    79
    0
                name: name.into(),
    80
    0
                func,
    81
    0
                return_field,
    82
    0
            })
    83
    0
        }
    84
    85
        /// return the name of the output column
    86
    0
        pub fn name(&self) -> &str {
    87
    0
            &self.name
    88
    0
        }
    89
    90
        /// Return the output field generated by evaluating this function
    91
    0
        pub fn field(&self, input_schema: &Schema) -> Result<Field> {
    92
    0
            Ok(Field::new(
    93
    0
                &self.name,
    94
    0
                self.func.data_type(input_schema)?,
    95
    0
                self.func.nullable(input_schema)?,
    96
            ))
    97
    0
        }
    98
    99
        /// Return the ideal batch size for this function
    100
    0
        pub fn ideal_batch_size(&self) -> Result<Option<usize>> {
    101
    0
            if let Some(expr) = self.func.as_any().downcast_ref::<ScalarFunctionExpr>() {
    102
    0
                if let Some(udf) =
    103
    0
                    expr.fun().inner().as_any().downcast_ref::<AsyncScalarUDF>()
    104
                {
    105
    0
                    return Ok(udf.ideal_batch_size());
    106
    0
                }
    107
    0
            }
    108
    0
            not_impl_err!("Can't get ideal_batch_size from {:?}", self.func)
    109
    0
        }
    110
    111
        /// This (async) function is called for each record batch to evaluate the LLM expressions
    112
        ///
    113
        /// The output is the output of evaluating the async expression and the input record batch
    114
    0
        pub async fn invoke_with_args(
    115
    0
            &self,
    116
    0
            batch: &RecordBatch,
    117
    0
            config_options: Arc<ConfigOptions>,
    118
    0
        ) -> Result<ColumnarValue> {
    119
    0
            let Some(scalar_function_expr) =
    120
    0
                self.func.as_any().downcast_ref::<ScalarFunctionExpr>()
    121
            else {
    122
    0
                return internal_err!(
    123
    0
                    "unexpected function type, expected ScalarFunctionExpr, got: {:?}",
    124
                    self.func
    125
                );
    126
            };
    127
    128
    0
            let Some(async_udf) = scalar_function_expr
    129
    0
                .fun()
    130
    0
                .inner()
    131
    0
                .as_any()
    132
    0
                .downcast_ref::<AsyncScalarUDF>()
    133
            else {
    134
    0
                return not_impl_err!(
    135
    0
                    "Don't know how to evaluate async function: {:?}",
    136
                    scalar_function_expr
    137
                );
    138
            };
    139
    140
    0
            let arg_fields = scalar_function_expr
    141
    0
                .args()
    142
    0
                .iter()
    143
    0
                .map(|e| e.return_field(batch.schema_ref()))
    144
    0
                .collect::<Result<Vec<_>>>()?;
    145
    146
    0
            let mut result_batches = vec![];
    147
    0
            if let Some(ideal_batch_size) = self.ideal_batch_size()? {
    148
    0
                let mut remainder = batch.clone();
    149
    0
                while remainder.num_rows() > 0 {
    150
    0
                    let size = if ideal_batch_size > remainder.num_rows() {
    151
    0
                        remainder.num_rows()
    152
                    } else {
    153
    0
                        ideal_batch_size
    154
                    };
    155
    156
    0
                    let current_batch = remainder.slice(0, size); // get next 10 rows
    157
    0
                    remainder = remainder.slice(size, remainder.num_rows() - size);
    158
    0
                    let args = scalar_function_expr
    159
    0
                        .args()
    160
    0
                        .iter()
    161
    0
                        .map(|e| e.evaluate(&current_batch))
    162
    0
                        .collect::<Result<Vec<_>>>()?;
    163
    0
                    result_batches.push(
    164
    0
                        async_udf
    165
    0
                            .invoke_async_with_args(ScalarFunctionArgs {
    166
    0
                                args,
    167
    0
                                arg_fields: arg_fields.clone(),
    168
    0
                                number_rows: current_batch.num_rows(),
    169
    0
                                return_field: Arc::clone(&self.return_field),
    170
    0
                                config_options: Arc::clone(&config_options),
    171
    0
                            })
    172
    0
                            .await?,
    173
                    );
    174
                }
    175
            } else {
    176
    0
                let args = scalar_function_expr
    177
    0
                    .args()
    178
    0
                    .iter()
    179
    0
                    .map(|e| e.evaluate(batch))
    180
    0
                    .collect::<Result<Vec<_>>>()?;
    181
    182
    0
                result_batches.push(
    183
    0
                    async_udf
    184
    0
                        .invoke_async_with_args(ScalarFunctionArgs {
    185
    0
                            args: args.to_vec(),
    186
    0
                            arg_fields,
    187
    0
                            number_rows: batch.num_rows(),
    188
    0
                            return_field: Arc::clone(&self.return_field),
    189
    0
                            config_options: Arc::clone(&config_options),
    190
    0
                        })
    191
    0
                        .await?,
    192
                );
    193
            }
    194
    195
    0
            let datas = ColumnarValue::values_to_arrays(&result_batches)?
    196
    0
                .iter()
    197
    0
                .map(|b| b.to_data())
    198
    0
                .collect::<Vec<_>>();
    199
    0
            let total_len = datas.iter().map(|d| d.len()).sum();
    200
    0
            let mut mutable = MutableArrayData::new(datas.iter().collect(), false, total_len);
    201
    0
            datas.iter().enumerate().for_each(|(i, data)| {
    202
    0
                mutable.extend(i, 0, data.len());
    203
    0
            });
    204
    0
            let array_ref = make_array(mutable.freeze());
    205
    0
            Ok(ColumnarValue::Array(array_ref))
    206
    0
        }
    207
    }
    208
    209
    impl PhysicalExpr for AsyncFuncExpr {
    210
    0
        fn as_any(&self) -> &dyn Any {
    211
    0
            self
    212
    0
        }
    213
    214
    0
        fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
    215
    0
            self.func.data_type(input_schema)
    216
    0
        }
    217
    218
    0
        fn nullable(&self, input_schema: &Schema) -> Result<bool> {
    219
    0
            self.func.nullable(input_schema)
    220
    0
        }
    221
    222
    0
        fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
    223
            // TODO: implement this for scalar value input
    224
    0
            not_impl_err!("AsyncFuncExpr.evaluate")
    225
    0
        }
    226
    227
    0
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    228
    0
            self.func.children()
    229
    0
        }
    230
    231
    0
        fn with_new_children(
    232
    0
            self: Arc<Self>,
    233
    0
            children: Vec<Arc<dyn PhysicalExpr>>,
    234
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    235
    0
            let new_func = Arc::clone(&self.func).with_new_children(children)?;
    236
    0
            Ok(Arc::new(AsyncFuncExpr {
    237
    0
                name: self.name.clone(),
    238
    0
                func: new_func,
    239
    0
                return_field: Arc::clone(&self.return_field),
    240
    0
            }))
    241
    0
        }
    242
    243
    0
        fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    244
    0
            write!(f, "{}", self.func)
    245
    0
        }
    246
    }