arrow-udf 0.10.0

//! DuckDB adapters used by code generated by [`crate::function`].

use std::sync::Arc;

use arrow_array::{Array, ArrayRef, ListArray, RecordBatch};
use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow_schema::Field;
use duckdb::core::{DataChunkHandle, LogicalTypeId};
use duckdb::ffi::duckdb_list_entry;

use crate::ScalarFunction;

enum InputColumn {
    Flat(ArrayRef),
    List {
        entries: Vec<duckdb_list_entry>,
        valid: Vec<bool>,
        values: ArrayRef,
    },
}

impl InputColumn {
    fn empty(&self) -> ArrayRef {
        match self {
            Self::Flat(array) => array.slice(0, 0),
            Self::List { values, .. } => Arc::new(ListArray::new(
                Arc::new(Field::new("item", values.data_type().clone(), true)),
                OffsetBuffer::new(ScalarBuffer::from(vec![0])),
                values.slice(0, 0),
                None,
            )),
        }
    }

    fn row(&self, row: usize) -> Result<ArrayRef, Box<dyn std::error::Error>> {
        match self {
            Self::Flat(array) => Ok(array.slice(row, 1)),
            Self::List {
                entries,
                valid,
                values,
            } => {
                let entry = entries[row];
                let length = i32::try_from(entry.length)?;
                let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, length]));
                let nulls = (!valid[row]).then(|| NullBuffer::from(vec![false]));
                Ok(Arc::new(ListArray::new(
                    Arc::new(Field::new("item", values.data_type().clone(), true)),
                    offsets,
                    values.slice(entry.offset as usize, entry.length as usize),
                    nulls,
                )))
            }
        }
    }
}

/// Invoke an Arrow scalar function on a DuckDB chunk containing LIST columns.
///
/// DuckDB LIST entries can reference a shared child range (notably for constant
/// lists), while Arrow list offsets must be monotonically increasing. This
/// adapter evaluates one row at a time with zero-copy child slices, avoiding a
/// potentially large materialization of repeated list values. Any number of
/// top-level LIST arguments with scalar children is supported.
///
/// # Safety
///
/// `input` must contain valid DuckDB vectors for the duration of this call.
pub unsafe fn invoke_scalar_with_lists(
    function: ScalarFunction,
    input: &mut DataChunkHandle,
) -> Result<ArrayRef, Box<dyn std::error::Error>> {
    let len = input.len();
    let mut columns = Vec::with_capacity(input.num_columns());

    for index in 0..input.num_columns() {
        let mut flat = input.flat_vector(index);
        if flat.logical_type().id() != LogicalTypeId::List {
            columns.push(InputColumn::Flat(
                duckdb::vtab::arrow::flat_vector_to_arrow_array(&mut flat, len)?,
            ));
            continue;
        }

        let entries = flat.as_slice_with_len::<duckdb_list_entry>(len).to_vec();
        let valid = (0..len).map(|row| !flat.row_is_null(row as u64)).collect();
        let list = input.list_vector(index);
        let mut child = list.child(list.len());
        if child.logical_type().id() == LogicalTypeId::List {
            return Err("nested DuckDB LIST arguments are not supported".into());
        }
        let values = duckdb::vtab::arrow::flat_vector_to_arrow_array(&mut child, list.len())?;
        columns.push(InputColumn::List {
            entries,
            valid,
            values,
        });
    }

    let mut outputs = Vec::with_capacity(len);
    for row in 0..len {
        let row_columns = columns
            .iter()
            .enumerate()
            .map(|(index, column)| Ok((index.to_string(), column.row(row)?)))
            .collect::<Result<Vec<_>, Box<dyn std::error::Error>>>()?;
        let batch = RecordBatch::try_from_iter(row_columns)?;
        let result = function(&batch)?;
        outputs.push(result.column(0).clone());
    }

    if outputs.is_empty() {
        let empty_columns = columns
            .iter()
            .enumerate()
            .map(|(index, column)| (index.to_string(), column.empty()));
        let empty = RecordBatch::try_from_iter(empty_columns)?;
        return Ok(function(&empty)?.column(0).clone());
    }
    let output_refs = outputs
        .iter()
        .map(|array| array.as_ref())
        .collect::<Vec<_>>();
    Ok(arrow_select::concat::concat(&output_refs)?)
}