Function async_udf

datafusion-examples/examples/udf/async_udf.rs:45–113

In this example we register `AskLLM` as an asynchronous user defined function and invoke it via the DataFrame API and SQL

()

/// In this example we register `AskLLM` as an asynchronous user defined function
/// and invoke it via the DataFrame API and SQL
pub async fn async_udf() -> Result<()> {
    // Use a hard coded parallelism level of 4 so the explain plan
    // is consistent across machines.
    let config = SessionConfig::new().with_target_partitions(4);
    let ctx =
        SessionContext::from(SessionStateBuilder::new().with_config(config).build());

    // Similarly to regular UDFs, you create an AsyncScalarUDF by implementing
    // `AsyncScalarUDFImpl` and creating an instance of `AsyncScalarUDF`.
    let async_equal = AskLLM::new();
    let udf = AsyncScalarUDF::new(Arc::new(async_equal));

    // Async UDFs are registered with the SessionContext, using the same
    // `register_udf` method as regular UDFs.
    ctx.register_udf(udf.into_scalar_udf());

    // Create a table named 'animal' with some sample data
    ctx.register_batch("animal", animal()?)?;

    // You can use the async UDF as normal in SQL queries
    //
    // Note: Async UDFs can currently be used in the select list and filter conditions.
    let results = ctx
        .sql("select * from animal a where ask_llm(a.name, 'Is this animal furry?')")
        .await?
        .collect()
        .await?;

    assert_batches_eq!(
        [
            "+----+------+",
            "| id | name |",
            "+----+------+",
            "| 1  | cat  |",
            "| 2  | dog  |",
            "+----+------+",
        ],
        &results
    );

    // While the interface is the same for both normal and async UDFs, you can
    // use `EXPLAIN` output to see that the async UDF uses a special
    // `AsyncFuncExec` node in the physical plan:
    let results = ctx
        .sql("explain select * from animal a where ask_llm(a.name, 'Is this animal furry?')")
        .await?
        .collect()
        .await?;

    assert_batches_eq!(
        [
            "+---------------+------------------------------------------------------------------------------------------------------------------------------+",
            "| plan_type     | plan                                                                                                                         |",
            "+---------------+------------------------------------------------------------------------------------------------------------------------------+",
            "| logical_plan  | SubqueryAlias: a                                                                                                             |",
            "|               |   Filter: ask_llm(CAST(animal.name AS Utf8View), Utf8View(\"Is this animal furry?\"))                                          |",
            "|               |     TableScan: animal projection=[id, name]                                                                                  |",
            "| physical_plan | FilterExec: __async_fn_0@2, projection=[id@0, name@1]                                                                        |",