Function array_position_scalar

datafusion/functions-nested/src/position.rs:226–300

Fast path for `array_position` when the needle is scalar. Performs a single bulk `not_distinct` comparison of the needle against the entire flat values buffer, then walks the result bitmap using offsets to find per-row first-match positions.

(
    haystack: &GenericListArray<O>,
    needle: &ArrayRef,
    arr_from: &[i64], // 0-indexed
)

/// entire flat values buffer, then walks the result bitmap using offsets to
/// find per-row first-match positions.
fn array_position_scalar<O: OffsetSizeTrait>(
    haystack: &GenericListArray<O>,
    needle: &ArrayRef,
    arr_from: &[i64], // 0-indexed
) -> Result<ArrayRef> {
    crate::utils::check_datatypes("array_position", &[haystack.values(), needle])?;

    if haystack.len() == 0 {
        return Ok(Arc::new(UInt64Array::new_null(0)));
    }

    let needle_datum = Scalar::new(Arc::clone(needle));
    let validity = haystack.nulls();

    // Only convert the visible portion of the values array. For sliced
    // ListArrays, values() returns the full underlying array but only
    // elements between the first and last offset are referenced.
    let offsets = haystack.offsets();
    let first_offset = offsets[0].as_usize();
    let last_offset = offsets[haystack.len()].as_usize();
    let visible_values = haystack
        .values()
        .slice(first_offset, last_offset - first_offset);

    // `not_distinct` treats NULL=NULL as true, matching the semantics of
    // `array_position`.
    let eq_array = arrow_ord::cmp::not_distinct(&visible_values, &needle_datum)?;
    let eq_bits = eq_array.values();

    let mut result: Vec<Option<u64>> = Vec::with_capacity(haystack.len());
    let mut matches = eq_bits.set_indices().peekable();

    // Match positions are relative to visible_values (0-based), so
    // subtract first_offset from each offset when comparing.
    for i in 0..haystack.len() {
        let start = offsets[i].as_usize() - first_offset;
        let end = offsets[i + 1].as_usize() - first_offset;

        if validity.is_some_and(|v| v.is_null(i)) {
            // Null row -> null output; advance past matches in range
            while matches.peek().is_some_and(|&p| p < end) {
                matches.next();
            }
            result.push(None);
            continue;
        }

        let from = arr_from[i];
        let row_len = end - start;
        if !(from >= 0 && (from as usize) <= row_len) {
            return exec_err!("start_from out of bounds: {}", from + 1);
        }
        let search_start = start + from as usize;

        // Advance past matches before search_start
        while matches.peek().is_some_and(|&p| p < search_start) {
            matches.next();
        }