arrow-row 56.2.1

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::{fixed, null_sentinel, LengthTracker, RowConverter, Rows, SortField};
use arrow_array::{new_null_array, Array, FixedSizeListArray, GenericListArray, OffsetSizeTrait};
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer};
use arrow_data::ArrayDataBuilder;
use arrow_schema::{ArrowError, DataType, SortOptions};
use std::{ops::Range, sync::Arc};

pub fn compute_lengths<O: OffsetSizeTrait>(
    lengths: &mut [usize],
    rows: &Rows,
    array: &GenericListArray<O>,
) {
    let shift = array.value_offsets()[0].as_usize();

    let offsets = array.value_offsets().windows(2);
    lengths
        .iter_mut()
        .zip(offsets)
        .enumerate()
        .for_each(|(idx, (length, offsets))| {
            let start = offsets[0].as_usize() - shift;
            let end = offsets[1].as_usize() - shift;
            let range = array.is_valid(idx).then_some(start..end);
            *length += encoded_len(rows, range);
        });
}

fn encoded_len(rows: &Rows, range: Option<Range<usize>>) -> usize {
    match range {
        None => 1,
        Some(range) => {
            1 + range
                .map(|i| super::variable::padded_length(Some(rows.row(i).as_ref().len())))
                .sum::<usize>()
        }
    }
}

/// Encodes the provided `GenericListArray` to `out` with the provided `SortOptions`
///
/// `rows` should contain the encoded child elements
pub fn encode<O: OffsetSizeTrait>(
    data: &mut [u8],
    offsets: &mut [usize],
    rows: &Rows,
    opts: SortOptions,
    array: &GenericListArray<O>,
) {
    let shift = array.value_offsets()[0].as_usize();

    offsets
        .iter_mut()
        .skip(1)
        .zip(array.value_offsets().windows(2))
        .enumerate()
        .for_each(|(idx, (offset, offsets))| {
            let start = offsets[0].as_usize() - shift;
            let end = offsets[1].as_usize() - shift;
            let range = array.is_valid(idx).then_some(start..end);
            let out = &mut data[*offset..];
            *offset += encode_one(out, rows, range, opts)
        });
}

#[inline]
fn encode_one(
    out: &mut [u8],
    rows: &Rows,
    range: Option<Range<usize>>,
    opts: SortOptions,
) -> usize {
    match range {
        None => super::variable::encode_null(out, opts),
        Some(range) if range.start == range.end => super::variable::encode_empty(out, opts),
        Some(range) => {
            let mut offset = 0;
            for i in range {
                let row = rows.row(i);
                offset += super::variable::encode_one(&mut out[offset..], Some(row.data), opts);
            }
            offset += super::variable::encode_empty(&mut out[offset..], opts);
            offset
        }
    }
}

/// Decodes an array from `rows` with the provided `options`
///
/// # Safety
///
/// `rows` must contain valid data for the provided `converter`
pub unsafe fn decode<O: OffsetSizeTrait>(
    converter: &RowConverter,
    rows: &mut [&[u8]],
    field: &SortField,
    validate_utf8: bool,
) -> Result<GenericListArray<O>, ArrowError> {
    let opts = field.options;

    let mut values_bytes = 0;

    let mut offset = 0;
    let mut offsets = Vec::with_capacity(rows.len() + 1);
    offsets.push(O::usize_as(0));

    for row in rows.iter_mut() {
        let mut row_offset = 0;
        loop {
            let decoded = super::variable::decode_blocks(&row[row_offset..], opts, |x| {
                values_bytes += x.len();
            });
            if decoded <= 1 {
                offsets.push(O::usize_as(offset));
                break;
            }
            row_offset += decoded;
            offset += 1;
        }
    }
    O::from_usize(offset).expect("overflow");

    let mut null_count = 0;
    let nulls = MutableBuffer::collect_bool(rows.len(), |x| {
        let valid = rows[x][0] != null_sentinel(opts);
        null_count += !valid as usize;
        valid
    });

    let mut values_offsets = Vec::with_capacity(offset);
    let mut values_bytes = Vec::with_capacity(values_bytes);
    for row in rows.iter_mut() {
        let mut row_offset = 0;
        loop {
            let decoded = super::variable::decode_blocks(&row[row_offset..], opts, |x| {
                values_bytes.extend_from_slice(x)
            });
            row_offset += decoded;
            if decoded <= 1 {
                break;
            }
            values_offsets.push(values_bytes.len());
        }
        *row = &row[row_offset..];
    }

    if opts.descending {
        values_bytes.iter_mut().for_each(|o| *o = !*o);
    }

    let mut last_value_offset = 0;
    let mut child_rows: Vec<_> = values_offsets
        .into_iter()
        .map(|offset| {
            let v = &values_bytes[last_value_offset..offset];
            last_value_offset = offset;
            v
        })
        .collect();

    let child = converter.convert_raw(&mut child_rows, validate_utf8)?;
    assert_eq!(child.len(), 1);

    let child_data = child[0].to_data();

    // Since RowConverter flattens certain data types (i.e. Dictionary),
    // we need to use updated data type instead of original field
    let corrected_type = match &field.data_type {
        DataType::List(inner_field) => DataType::List(Arc::new(
            inner_field
                .as_ref()
                .clone()
                .with_data_type(child_data.data_type().clone()),
        )),
        DataType::LargeList(inner_field) => DataType::LargeList(Arc::new(
            inner_field
                .as_ref()
                .clone()
                .with_data_type(child_data.data_type().clone()),
        )),
        _ => unreachable!(),
    };

    let builder = ArrayDataBuilder::new(corrected_type)
        .len(rows.len())
        .null_count(null_count)
        .null_bit_buffer(Some(nulls.into()))
        .add_buffer(Buffer::from_vec(offsets))
        .add_child_data(child_data);

    Ok(GenericListArray::from(unsafe { builder.build_unchecked() }))
}

pub fn compute_lengths_fixed_size_list(
    tracker: &mut LengthTracker,
    rows: &Rows,
    array: &FixedSizeListArray,
) {
    let value_length = array.value_length().as_usize();
    tracker.push_variable((0..array.len()).map(|idx| {
        match array.is_valid(idx) {
            true => {
                1 + ((idx * value_length)..(idx + 1) * value_length)
                    .map(|child_idx| rows.row(child_idx).as_ref().len())
                    .sum::<usize>()
            }
            false => 1,
        }
    }))
}

/// Encodes the provided `FixedSizeListArray` to `out` with the provided `SortOptions`
///
/// `rows` should contain the encoded child elements
pub fn encode_fixed_size_list(
    data: &mut [u8],
    offsets: &mut [usize],
    rows: &Rows,
    opts: SortOptions,
    array: &FixedSizeListArray,
) {
    let null_sentinel = null_sentinel(opts);
    offsets
        .iter_mut()
        .skip(1)
        .enumerate()
        .for_each(|(idx, offset)| {
            let value_length = array.value_length().as_usize();
            match array.is_valid(idx) {
                true => {
                    data[*offset] = 0x01;
                    *offset += 1;
                    for child_idx in (idx * value_length)..(idx + 1) * value_length {
                        let row = rows.row(child_idx);
                        let end_offset = *offset + row.as_ref().len();
                        data[*offset..end_offset].copy_from_slice(row.as_ref());
                        *offset = end_offset;
                    }
                }
                false => {
                    data[*offset] = null_sentinel;
                    *offset += 1;
                }
            };
        })
}

/// Decodes a fixed size list array from `rows` with the provided `options`
///
/// # Safety
///
/// `rows` must contain valid data for the provided `converter`
pub unsafe fn decode_fixed_size_list(
    converter: &RowConverter,
    rows: &mut [&[u8]],
    field: &SortField,
    validate_utf8: bool,
    value_length: usize,
) -> Result<FixedSizeListArray, ArrowError> {
    let list_type = &field.data_type;
    let element_type = match list_type {
        DataType::FixedSizeList(element_field, _) => element_field.data_type(),
        _ => {
            return Err(ArrowError::InvalidArgumentError(format!(
                "Expected FixedSizeListArray, found: {list_type:?}",
            )))
        }
    };

    let len = rows.len();
    let (null_count, nulls) = fixed::decode_nulls(rows);

    let null_element_encoded = converter.convert_columns(&[new_null_array(element_type, 1)])?;
    let null_element_encoded = null_element_encoded.row(0);
    let null_element_slice = null_element_encoded.as_ref();

    let mut child_rows = Vec::new();
    for row in rows {
        let valid = row[0] == 1;
        let mut row_offset = 1;
        if !valid {
            for _ in 0..value_length {
                child_rows.push(null_element_slice);
            }
        } else {
            for _ in 0..value_length {
                let mut temp_child_rows = vec![&row[row_offset..]];
                converter.convert_raw(&mut temp_child_rows, validate_utf8)?;
                let decoded_bytes = row.len() - row_offset - temp_child_rows[0].len();
                let next_offset = row_offset + decoded_bytes;
                child_rows.push(&row[row_offset..next_offset]);
                row_offset = next_offset;
            }
        }
        *row = &row[row_offset..]; // Update row for the next decoder
    }

    let children = converter.convert_raw(&mut child_rows, validate_utf8)?;
    let child_data = children.iter().map(|c| c.to_data()).collect();
    let builder = ArrayDataBuilder::new(list_type.clone())
        .len(len)
        .null_count(null_count)
        .null_bit_buffer(Some(nulls))
        .child_data(child_data);

    Ok(FixedSizeListArray::from(builder.build_unchecked()))
}