Function generic_set_loop

datafusion/functions-nested/src/set_ops.rs:396–505

Inner loop for set operations, parameterized by const generic to avoid branching inside the hot loop.

(
    l: &GenericListArray<OffsetSize>,
    r: &GenericListArray<OffsetSize>,
    rows_l: &arrow::row::Rows,
    rows_r: &arrow::row::Rows,
    field: Arc<Field>,
    combined_values: &ArrayRef,
    r

/// Inner loop for set operations, parameterized by const generic to
/// avoid branching inside the hot loop.
fn generic_set_loop<OffsetSize: OffsetSizeTrait, const IS_UNION: bool>(
    l: &GenericListArray<OffsetSize>,
    r: &GenericListArray<OffsetSize>,
    rows_l: &arrow::row::Rows,
    rows_r: &arrow::row::Rows,
    field: Arc<Field>,
    combined_values: &ArrayRef,
    r_offset: usize,
) -> Result<ArrayRef> {
    let l_offsets = l.value_offsets();
    let r_offsets = r.value_offsets();
    let l_first = l.offsets()[0].as_usize();
    let r_first = r.offsets()[0].as_usize();

    let mut result_offsets = Vec::with_capacity(l.len() + 1);
    result_offsets.push(OffsetSize::usize_as(0));
    let initial_capacity = if IS_UNION {
        // Union can include all elements from both sides
        rows_l.num_rows()
    } else {
        // Intersect result is bounded by the smaller side
        rows_l.num_rows().min(rows_r.num_rows())
    };

    let mut indices: Vec<usize> = Vec::with_capacity(initial_capacity);

    // Reuse hash sets across iterations
    let mut seen = HashSet::new();
    let mut lookup_set = HashSet::new();
    for i in 0..l.len() {
        let last_offset = *result_offsets.last().unwrap();

        if l.is_null(i) || r.is_null(i) {
            result_offsets.push(last_offset);
            continue;
        }

        let l_start = l_offsets[i].as_usize() - l_first;
        let l_end = l_offsets[i + 1].as_usize() - l_first;
        let r_start = r_offsets[i].as_usize() - r_first;
        let r_end = r_offsets[i + 1].as_usize() - r_first;

        seen.clear();

        if IS_UNION {
            for idx in l_start..l_end {
                let row = rows_l.row(idx);
                if seen.insert(row) {
                    indices.push(idx);
                }
            }
            for idx in r_start..r_end {
                let row = rows_r.row(idx);
                if seen.insert(row) {
                    indices.push(idx + r_offset);
                }
            }
        } else {