duckdb 1.10505.0

use std::sync::Arc;

use arrow::{
    array::{Array, RecordBatch},
    datatypes::DataType,
};

use crate::{
    arrow_interop::{WritableVector, data_chunk_to_arrow, to_duckdb_logical_type, write_arrow_array_to_vector},
    core::DataChunkHandle,
};

use super::{ScalarFunctionSignature, ScalarParams, VScalar};

/// The possible parameters of a scalar function that accepts and returns arrow types
pub enum ArrowScalarParams {
    /// The exact parameters of the scalar function
    Exact(Vec<DataType>),
    /// The variadic parameter of the scalar function
    Variadic(DataType),
}

impl AsRef<[DataType]> for ArrowScalarParams {
    fn as_ref(&self) -> &[DataType] {
        match self {
            Self::Exact(params) => params.as_ref(),
            Self::Variadic(param) => std::slice::from_ref(param),
        }
    }
}

impl From<ArrowScalarParams> for ScalarParams {
    fn from(params: ArrowScalarParams) -> Self {
        match params {
            ArrowScalarParams::Exact(params) => Self::Exact(
                params
                    .into_iter()
                    .map(|v| to_duckdb_logical_type(&v).expect("type should be converted"))
                    .collect(),
            ),
            ArrowScalarParams::Variadic(param) => {
                Self::Variadic(to_duckdb_logical_type(&param).expect("type should be converted"))
            }
        }
    }
}

/// A signature for a scalar function that accepts and returns arrow types
pub struct ArrowFunctionSignature {
    /// The parameters of the scalar function
    pub parameters: Option<ArrowScalarParams>,
    /// The return type of the scalar function
    pub return_type: DataType,
}

impl ArrowFunctionSignature {
    /// Create an exact function signature
    pub fn exact(params: Vec<DataType>, return_type: DataType) -> Self {
        Self {
            parameters: Some(ArrowScalarParams::Exact(params)),
            return_type,
        }
    }

    /// Create a variadic function signature
    pub fn variadic(param: DataType, return_type: DataType) -> Self {
        Self {
            parameters: Some(ArrowScalarParams::Variadic(param)),
            return_type,
        }
    }
}

/// A trait for scalar functions that accept and return arrow types that can be registered with DuckDB
pub trait VArrowScalar: Sized {
    /// State set at registration time. Persists for the lifetime of the catalog entry.
    /// Shared across worker threads and invocations — must not be modified during execution.
    /// Must be `'static` as it is stored in DuckDB and may outlive the current stack frame.
    type State: Default + Sized + Send + Sync + 'static;

    /// The actual function that is called by DuckDB
    fn invoke(state: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>>;

    /// The possible signatures of the scalar function. These will result in DuckDB scalar function overloads.
    /// The invoke method should be able to handle all of these signatures.
    fn signatures() -> Vec<ArrowFunctionSignature>;

    /// Whether the scalar function is volatile.
    ///
    /// Volatile functions are re-evaluated for each row, even if they have no parameters.
    /// This is useful for functions that generate random or unique values, such as random
    /// number generators, UUID generators, or fake data generators.
    ///
    /// By default, DuckDB optimizes zero-argument scalar functions as constants, evaluating
    /// them only once. Returning true from this method prevents this optimization.
    ///
    /// # Default
    /// Returns `false` by default, meaning the function is not volatile.
    fn volatile() -> bool {
        false
    }
}

impl<T> VScalar for T
where
    T: VArrowScalar,
{
    type State = T::State;

    fn invoke(
        state: &Self::State,
        input: &mut DataChunkHandle,
        out: &mut dyn WritableVector,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let array = T::invoke(state, data_chunk_to_arrow(input)?)?;
        write_arrow_array_to_vector(&array, out)
    }

    fn signatures() -> Vec<ScalarFunctionSignature> {
        T::signatures()
            .into_iter()
            .map(|sig| ScalarFunctionSignature {
                parameters: sig.parameters.map(Into::into),
                return_type: to_duckdb_logical_type(&sig.return_type).expect("type should be converted"),
            })
            .collect()
    }
}

#[cfg(test)]
mod test {

    use std::{error::Error, sync::Arc};

    use arrow::{
        array::{Array, AsArray, Int32Array, Int64Array, ListArray, RecordBatch, StringArray},
        datatypes::{
            ArrowPrimitiveType, DataType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType,
            TimestampNanosecondType, TimestampSecondType,
        },
    };

    use crate::{Connection, vscalar::arrow::ArrowFunctionSignature};

    use super::VArrowScalar;

    struct HelloScalarArrow {}

    impl VArrowScalar for HelloScalarArrow {
        type State = ();

        fn invoke(_: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
            let name = input.column(0).as_any().downcast_ref::<StringArray>().unwrap();
            let result = name.iter().map(|v| format!("Hello {}", v.unwrap())).collect::<Vec<_>>();
            Ok(Arc::new(StringArray::from(result)))
        }

        fn signatures() -> Vec<ArrowFunctionSignature> {
            vec![ArrowFunctionSignature::exact(vec![DataType::Utf8], DataType::Utf8)]
        }
    }

    #[derive(Debug)]
    struct MockState {
        info: String,
    }

    impl Default for MockState {
        fn default() -> Self {
            Self {
                info: "some meta".to_string(),
            }
        }
    }

    impl Drop for MockState {
        fn drop(&mut self) {
            println!("dropped meta");
        }
    }

    struct ArrowMultiplyScalar {}

    impl VArrowScalar for ArrowMultiplyScalar {
        type State = MockState;

        fn invoke(_: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
            let a = input
                .column(0)
                .as_any()
                .downcast_ref::<::arrow::array::Float32Array>()
                .unwrap();

            let b = input
                .column(1)
                .as_any()
                .downcast_ref::<::arrow::array::Float32Array>()
                .unwrap();

            let result = a
                .iter()
                .zip(b.iter())
                .map(|(a, b)| a.unwrap() * b.unwrap())
                .collect::<Vec<_>>();
            Ok(Arc::new(::arrow::array::Float32Array::from(result)))
        }

        fn signatures() -> Vec<ArrowFunctionSignature> {
            vec![ArrowFunctionSignature::exact(
                vec![DataType::Float32, DataType::Float32],
                DataType::Float32,
            )]
        }
    }

    // accepts a string or a number and parses to int and multiplies by 2
    struct ArrowOverloaded {}

    impl VArrowScalar for ArrowOverloaded {
        type State = MockState;

        fn invoke(state: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
            assert_eq!("some meta", state.info);

            let a = input.column(0);
            let b = input.column(1);

            let result = match a.data_type() {
                DataType::Utf8 => {
                    let a = a
                        .as_any()
                        .downcast_ref::<::arrow::array::StringArray>()
                        .unwrap()
                        .iter()
                        .map(|v| v.unwrap().parse::<f32>().unwrap())
                        .collect::<Vec<_>>();
                    let b = b
                        .as_any()
                        .downcast_ref::<::arrow::array::Float32Array>()
                        .unwrap()
                        .iter()
                        .map(|v| v.unwrap())
                        .collect::<Vec<_>>();
                    a.iter().zip(b.iter()).map(|(a, b)| a * b).collect::<Vec<_>>()
                }
                DataType::Float32 => {
                    let a = a
                        .as_any()
                        .downcast_ref::<::arrow::array::Float32Array>()
                        .unwrap()
                        .iter()
                        .map(|v| v.unwrap())
                        .collect::<Vec<_>>();
                    let b = b
                        .as_any()
                        .downcast_ref::<::arrow::array::Float32Array>()
                        .unwrap()
                        .iter()
                        .map(|v| v.unwrap())
                        .collect::<Vec<_>>();
                    a.iter().zip(b.iter()).map(|(a, b)| a * b).collect::<Vec<_>>()
                }
                _ => panic!("unsupported type"),
            };

            Ok(Arc::new(::arrow::array::Float32Array::from(result)))
        }

        fn signatures() -> Vec<ArrowFunctionSignature> {
            vec![
                ArrowFunctionSignature::exact(vec![DataType::Utf8, DataType::Float32], DataType::Float32),
                ArrowFunctionSignature::exact(vec![DataType::Float32, DataType::Float32], DataType::Float32),
            ]
        }
    }

    #[test]
    fn test_arrow_scalar() -> Result<(), Box<dyn Error>> {
        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<HelloScalarArrow>("hello")?;

        let batches = conn
            .prepare("select hello('foo') as hello from range(10)")?
            .query_arrow([])?
            .collect::<Vec<_>>();

        for batch in batches.iter() {
            let array = batch.column(0);
            let array = array.as_any().downcast_ref::<::arrow::array::StringArray>().unwrap();
            for i in 0..array.len() {
                assert_eq!(array.value(i), format!("Hello foo"));
            }
        }

        Ok(())
    }

    #[test]
    fn test_arrow_scalar_multiply() -> Result<(), Box<dyn Error>> {
        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<ArrowMultiplyScalar>("multiply_udf")?;

        let batches = conn
            .prepare("select multiply_udf(3.0, 2.0) as mult_result from range(10)")?
            .query_arrow([])?
            .collect::<Vec<_>>();

        for batch in batches.iter() {
            let array = batch.column(0);
            let array = array.as_any().downcast_ref::<::arrow::array::Float32Array>().unwrap();
            for i in 0..array.len() {
                assert_eq!(array.value(i), 6.0);
            }
        }
        Ok(())
    }

    #[test]
    fn test_multiple_signatures_scalar() -> Result<(), Box<dyn Error>> {
        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<ArrowOverloaded>("multi_sig_udf")?;

        let batches = conn
            .prepare("select multi_sig_udf('3', 5) as message from range(2)")?
            .query_arrow([])?
            .collect::<Vec<_>>();

        for batch in batches.iter() {
            let array = batch.column(0);
            let array = array.as_any().downcast_ref::<::arrow::array::Float32Array>().unwrap();
            for i in 0..array.len() {
                assert_eq!(array.value(i), 15.0);
            }
        }

        let batches = conn
            .prepare("select multi_sig_udf(12, 10) as message from range(2)")?
            .query_arrow([])?
            .collect::<Vec<_>>();

        for batch in batches.iter() {
            let array = batch.column(0);
            let array = array.as_any().downcast_ref::<::arrow::array::Float32Array>().unwrap();
            for i in 0..array.len() {
                assert_eq!(array.value(i), 120.0);
            }
        }

        Ok(())
    }

    #[test]
    fn test_split_function() -> Result<(), Box<dyn Error>> {
        struct SplitFunction {}

        impl VArrowScalar for SplitFunction {
            type State = ();

            fn invoke(_: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
                let strings = input.column(0).as_any().downcast_ref::<StringArray>().unwrap();

                let mut builder = arrow::array::ListBuilder::new(arrow::array::StringBuilder::with_capacity(
                    strings.len(),
                    strings.len() * 10,
                ));

                for s in strings.iter() {
                    let s = s.unwrap();
                    for split_value in s.split(' ').collect::<Vec<_>>() {
                        builder.values().append_value(split_value);
                    }
                    builder.append(true);
                }

                Ok(Arc::new(builder.finish()))
            }

            fn signatures() -> Vec<ArrowFunctionSignature> {
                vec![ArrowFunctionSignature::exact(
                    vec![DataType::Utf8],
                    DataType::List(Arc::new(arrow::datatypes::Field::new("item", DataType::Utf8, true))),
                )]
            }
        }

        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<SplitFunction>("split_string")?;

        // Test with single string
        let batches = conn
            .prepare("select split_string('hello world') as result")?
            .query_arrow([])?
            .collect::<Vec<_>>();

        let array = batches[0].column(0);
        let list_array = array.as_any().downcast_ref::<arrow::array::ListArray>().unwrap();
        let values = list_array.value(0);
        let string_values = values.as_any().downcast_ref::<StringArray>().unwrap();

        assert_eq!(string_values.value(0), "hello");
        assert_eq!(string_values.value(1), "world");

        Ok(())
    }

    #[test]
    fn test_arrow_scalar_reads_filtered_list_vectors() -> Result<(), Box<dyn Error>> {
        struct ListFirstValueFunction;

        impl VArrowScalar for ListFirstValueFunction {
            type State = ();

            fn invoke(_: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
                let lists = input.column(0).as_any().downcast_ref::<ListArray>().unwrap();
                let first_values = lists
                    .iter()
                    .map(|value| value.map(|value| value.as_any().downcast_ref::<Int32Array>().unwrap().value(0)))
                    .collect::<Int32Array>();
                Ok(Arc::new(first_values))
            }

            fn signatures() -> Vec<ArrowFunctionSignature> {
                vec![ArrowFunctionSignature::exact(
                    vec![DataType::List(Arc::new(arrow::datatypes::Field::new(
                        "item",
                        DataType::Int32,
                        true,
                    )))],
                    DataType::Int32,
                )]
            }
        }

        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<ListFirstValueFunction>("arrow_list_first_value")?;
        conn.execute_batch(
            "create table list_input as \
             select i::integer as id, \
                    case when i % 7 = 0 then null \
                         else [i::integer, (i + 1)::integer] end as values \
             from range(5000) t(i)",
        )?;

        let first_values = conn
            .prepare(
                "select arrow_list_first_value(values) \
                 from list_input where id % 97 = 0 order by id",
            )?
            .query_map([], |row| row.get::<_, Option<i32>>(0))?
            .collect::<Result<Vec<_>, _>>()?;

        let expected = (0..5000)
            .step_by(97)
            .map(|id| if id % 7 == 0 { None } else { Some(id) })
            .collect::<Vec<_>>();
        assert_eq!(first_values, expected);
        Ok(())
    }

    fn timestamp_value<T>(input: &RecordBatch, column: usize, expected_type: &DataType) -> Result<i64, Box<dyn Error>>
    where
        T: ArrowPrimitiveType<Native = i64>,
    {
        let array = input.column(column);
        if array.data_type() != expected_type {
            return Err(format!(
                "expected timestamp column {column} to have type {expected_type}, got {}",
                array.data_type()
            )
            .into());
        }

        array
            .as_primitive_opt::<T>()
            .map(|timestamps| timestamps.value(0))
            .ok_or_else(|| format!("timestamp column {column} has an unexpected Arrow array implementation").into())
    }

    #[test]
    fn test_arrow_scalar_reads_timestamp_carriers() -> Result<(), Box<dyn Error>> {
        struct TimestampCarrierProbe;

        impl VArrowScalar for TimestampCarrierProbe {
            type State = ();

            fn invoke(_: &Self::State, input: RecordBatch) -> Result<Arc<dyn Array>, Box<dyn std::error::Error>> {
                let timestamp_tz = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
                let actual = [
                    timestamp_value::<TimestampSecondType>(&input, 0, &DataType::Timestamp(TimeUnit::Second, None))?,
                    timestamp_value::<TimestampMillisecondType>(
                        &input,
                        1,
                        &DataType::Timestamp(TimeUnit::Millisecond, None),
                    )?,
                    timestamp_value::<TimestampMicrosecondType>(
                        &input,
                        2,
                        &DataType::Timestamp(TimeUnit::Microsecond, None),
                    )?,
                    timestamp_value::<TimestampNanosecondType>(
                        &input,
                        3,
                        &DataType::Timestamp(TimeUnit::Nanosecond, None),
                    )?,
                    timestamp_value::<TimestampMicrosecondType>(&input, 4, &timestamp_tz)?,
                ];
                let expected = [
                    1_704_067_200,
                    1_704_067_200_123,
                    1_704_067_200_123_456,
                    1_704_067_200_000_000_123,
                    1_704_067_200_123_456,
                ];
                if actual != expected {
                    return Err(format!("expected timestamp carriers {expected:?}, got {actual:?}").into());
                }

                Ok(Arc::new(Int64Array::from_iter_values([actual[3]])))
            }

            fn signatures() -> Vec<ArrowFunctionSignature> {
                vec![ArrowFunctionSignature::exact(
                    vec![
                        DataType::Timestamp(TimeUnit::Second, None),
                        DataType::Timestamp(TimeUnit::Millisecond, None),
                        DataType::Timestamp(TimeUnit::Microsecond, None),
                        DataType::Timestamp(TimeUnit::Nanosecond, None),
                        DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
                    ],
                    DataType::Int64,
                )]
            }
        }

        let conn = Connection::open_in_memory()?;
        conn.register_scalar_function::<TimestampCarrierProbe>("arrow_timestamp_carriers_raw")?;

        // 2024-01-01T00:00:00Z in each carrier's resolution.
        let value = conn.query_row(
            "select arrow_timestamp_carriers_raw(\
             TIMESTAMP_S '2024-01-01 00:00:00', \
             TIMESTAMP_MS '2024-01-01 00:00:00.123', \
             TIMESTAMP '2024-01-01 00:00:00.123456', \
             TIMESTAMP_NS '2024-01-01 00:00:00.000000123', \
             TIMESTAMPTZ '2024-01-01 00:00:00.123456+00')",
            [],
            |row| row.get::<_, i64>(0),
        )?;
        assert_eq!(value, 1_704_067_200_000_000_123);

        Ok(())
    }
}