Arrow safe refactoring

[Hennzau]

Objective

Hi! I was adding other datatypes in the other_datatype branch, and after discussing with Marc Perlade https://github.com/Ar3sion https://git.eleves.ens.fr/mperlade, a student at ENS PSL, we found a safer way to retrieve data from the arrow::UnionArray than the previous method. So, I rewrote the Arrow part and achieved a really neat result!

Changes

If you want to convert an Arrow object into an Image, you can simply get the UnionArray, convert it into a HashMap, and then extract the fields as Vec<T>.

pub fn from_arrow(array: arrow::array::UnionArray) -> Result<Self> {
    let mut map = arrow_union_into_map(array)?;

    let width =
        get_primitive_array_from_map::<u32, arrow::datatypes::UInt32Type>("width", &mut map)?
            .first();
    let name = get_utf8_array_from_map("name", &mut map)?
        .first();
}

Code

Everything is safe!

// Convert an ArrowUnion to a HashMap
pub fn arrow_union_into_map(
    array: arrow::array::UnionArray,
) -> Result<HashMap<String, arrow::array::ArrayRef>> {
    let mut result = HashMap::new();

    let (union_fields, _, _, children) = array.into_parts();

    for (a, b) in union_fields.iter() {
        let child = children
            .get(a as usize)
            .ok_or_eyre(Report::msg(
                "Invalid union array field index. Must be >= 0 and correspond to the children index in the array.",
            ))?
            .clone();

        result.insert(b.name().to_string(), child);
    }

    Ok(result)
}

// Extract a Primitive Array without copying by consuming the `Arc<dyn Array>`
pub fn get_primitive_array_from_map<
    T: arrow::datatypes::ArrowNativeType,
    G: arrow::datatypes::ArrowPrimitiveType,
>(
    field: &str,
    map: &mut HashMap<String, arrow::array::ArrayRef>,
) -> Result<Vec<T>> {
    use arrow::array::Array;

    let array_data = map
        .remove(field)
        .ok_or_eyre(Report::msg("Invalid field for this map."))?
        .into_data();

    let array = arrow::array::PrimitiveArray::<G>::from(array_data);
    let (_, buffer, _) = array.into_parts();
    let buffer = buffer.into_inner();

    match buffer.into_vec::<T>() {
        Ok(vec) => Ok(vec),
        Err(e) => Err(Report::msg(format!(
            "T is not a valid type for this buffer. It must have the same layout as the buffer (this usually occurs when the type is incorrect or when another reference exists). Error: {:?}", e
        ))),
    }
}

// Extract a UTF8 String Array with only ONE copy per string, by consuming the `Arc<dyn Array>`
pub fn get_utf8_array_from_map(
    field: &str,
    map: &mut HashMap<String, arrow::array::ArrayRef>,
) -> Result<Vec<String>> {
    use arrow::array::Array;

    let array_data = map
        .remove(field)
        .ok_or_eyre(Report::msg("Invalid field for this map."))?
        .into_data();

    let array = arrow::array::StringArray::from(array_data);
    let (offsets, buffer, _) = array.into_parts();

    let slice = buffer.as_slice();
    let mut last_offset = 0;
    let mut iterator = offsets.iter();
    iterator.next();

    iterator
        .map(|&offset| {
            let offset = offset as usize;
            let slice = &slice[last_offset..offset];
            last_offset = offset;

            String::from_utf8(slice.to_vec()).wrap_err(Report::msg("Array is not UTF-8 encoded."))
        })
        .collect::<Result<Vec<String>>>()
}

Acknowledgments

I really want to thank Marc Perlade for his help. He’s not into sharing his work on GitHub/GitLab, but he is fantastic.

[haixuanTao]

So, I think this should work, but it should also add couple of heap calls.

I'm not against it, but it could be interesting to add some roundtrip benchmark if you have time.

[haixuanTao]

Suggested change

	Encoding::RGB8 => {
	let data = get_primitive_array_from_map::<u8, arrow::datatypes::UInt8Type>(
	"data", &mut map,
	)?;
	DataContainer::from_u8(data)
	}
	Encoding::BGR8 => {
	let data = get_primitive_array_from_map::<u8, arrow::datatypes::UInt8Type>(
	"data", &mut map,
	)?;
	DataContainer::from_u8(data)
	}
	Encoding::GRAY8 => {
	let data = get_primitive_array_from_map::<u8, arrow::datatypes::UInt8Type>(
	"data", &mut map,
	)?;
	DataContainer::from_u8(data)
	}
	Encoding::RGB8 | Encoding::BGR8 | Encoding::GRAY8 => {
	let data = get_primitive_array_from_map::<u8, arrow::datatypes::UInt8Type>(
	"data", &mut map,
	)?;
	DataContainer::from_u8(data)
	}