Function elt

datafusion/spark/src/function/string/elt.rs:94–142

(args: &[ArrayRef])

}

fn elt(args: &[ArrayRef]) -> Result<ArrayRef, DataFusionError> {
    let n_rows = args[0].len();

    let idx: &PrimitiveArray<Int64Type> =
        args[0].as_primitive_opt::<Int64Type>().ok_or_else(|| {
            DataFusionError::Plan(format!(
                "ELT function: first argument must be Int64 (got {:?})",
                args[0].data_type()
            ))
        })?;

    let num_values = args.len() - 1;
    let mut cols: Vec<Arc<StringArray>> = Vec::with_capacity(num_values);
    for a in args.iter().skip(1) {
        let casted = cast(a, &Utf8)?;
        let sa = as_string_array(&casted)?;
        cols.push(Arc::new(sa.clone()));
    }

    let mut builder = StringBuilder::new();

    for i in 0..n_rows {
        if idx.is_null(i) {
            builder.append_null();
            continue;
        }

        let index = idx.value(i);

        // TODO: if spark.sql.ansi.enabled is true,
        //  throw ArrayIndexOutOfBoundsException for invalid indices;
        //  if false, return NULL instead (current behavior).
        if index < 1 || (index as usize) > num_values {
            builder.append_null();
            continue;
        }

        let value_idx = (index as usize) - 1;
        let col = &cols[value_idx];

        if col.is_null(i) {
            builder.append_null();
        } else {
            builder.append_value(col.value(i));
        }
    }

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

#[cfg(test)]
mod tests {