Function try_array_positions_scalar

datafusion/functions-nested/src/position.rs:433–474

Attempts the scalar-needle fast path for `array_positions`.

(args: &[ColumnarValue])

/// Attempts the scalar-needle fast path for `array_positions`.
fn try_array_positions_scalar(args: &[ColumnarValue]) -> Result<Option<ColumnarValue>> {
    let [haystack_arg, needle_arg] = take_function_args("array_positions", args)?;

    let scalar_needle = match needle_arg {
        ColumnarValue::Scalar(s) => s,
        ColumnarValue::Array(_) => return Ok(None),
    };

    // `not_distinct` doesn't support nested types (List, Struct, etc.),
    // so fall back to the per-row path for those.
    if scalar_needle.data_type().is_nested() {
        return Ok(None);
    }

    let (num_rows, all_inputs_scalar) = match haystack_arg {
        ColumnarValue::Array(a) => (a.len(), false),
        ColumnarValue::Scalar(_) => (1, true),
    };

    let needle = scalar_needle.to_array_of_size(1)?;
    let haystack = haystack_arg.to_array(num_rows)?;

    let result = match haystack.data_type() {
        List(_) => {
            let list = as_list_array(&haystack)?;
            array_positions_scalar::<i32>(list, &needle)
        }
        LargeList(_) => {
            let list = as_large_list_array(&haystack)?;
            array_positions_scalar::<i64>(list, &needle)
        }
        t => exec_err!("array_positions does not support type '{t}'"),
    }?;

    if all_inputs_scalar {
        Ok(Some(ColumnarValue::Scalar(ScalarValue::try_from_array(
            &result, 0,
        )?)))
    } else {
        Ok(Some(ColumnarValue::Array(result)))
    }
}

fn array_positions_inner(args: &[ArrayRef]) -> Result<ArrayRef> {
    let [haystack, needle] = take_function_args("array_positions", args)?;