legume-numeric 0.8.11

use parquet::basic::Type as ParquetType;
use parquet::basic::{Compression, ConvertedType, ZstdLevel};
use parquet::data_type::ByteArray;
use parquet::data_type::{ByteArrayType, DoubleType, FloatType, Int32Type, Int64Type};
use parquet::file::properties::WriterProperties;
use parquet::file::reader::{FileReader, SerializedFileReader};
use parquet::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
use parquet::record::RowAccessor;
use parquet::schema::types::Type as SchemaType;
use rustc_hash::FxHashSet as HashSet;
use std::any::TypeId;
use std::fs::File;
use std::sync::Arc;

/// get field names by peeking into `file_path`
pub fn peek_parquet_field_names(file_path: &str) -> anyhow::Result<Vec<Box<str>>> {
    let file = File::open(file_path)?;
    let reader = SerializedFileReader::new(file)?;
    let metadata = reader.metadata();
    let fields = metadata.file_metadata().schema().get_fields();

    Ok(fields
        .iter()
        .map(|f| f.name().to_string().into_boxed_str())
        .collect())
}

/// The first string (`BYTE_ARRAY`) column of a parquet file, by index, or
/// `None` when every column is numeric.
///
/// This is how a table's NAME column is found: by type, not by position. A
/// table written without its index has no string column, and a reader that
/// took column 0 regardless would stringify the first data column into names.
pub fn first_string_column(file_path: &str) -> anyhow::Result<Option<usize>> {
    let file = File::open(file_path)?;
    let reader = SerializedFileReader::new(file)?;
    let fields = reader
        .metadata()
        .file_metadata()
        .schema()
        .get_fields()
        .to_vec();
    Ok(fields
        .iter()
        .position(|f| f.get_physical_type() == ParquetType::BYTE_ARRAY))
}

/// How many numeric (non-string) columns a parquet file has: the width of a
/// named table, read from the footer without decoding a row.
pub fn parquet_numeric_column_count(file_path: &str) -> anyhow::Result<usize> {
    let file = File::open(file_path)?;
    let reader = SerializedFileReader::new(file)?;
    Ok(reader
        .metadata()
        .file_metadata()
        .schema()
        .get_fields()
        .iter()
        .filter(|f| f.get_physical_type() != ParquetType::BYTE_ARRAY)
        .count())
}

/// Read one string (`BYTE_ARRAY`) column out of a parquet file.
///
/// [`ParquetReader`] is a *matrix* reader: it needs at least one numeric column
/// and bails with "no available columns" otherwise. A name table (a gene list, a
/// `gene,celltype` marker file) is all strings, so it needs this instead.
pub fn read_parquet_string_column(
    file_path: &str,
    column_index: usize,
) -> anyhow::Result<Vec<Box<str>>> {
    let file = File::open(file_path)?;
    let reader = SerializedFileReader::new(file)?;
    let metadata = reader.metadata();
    let nrows = metadata.file_metadata().num_rows() as usize;
    let fields = metadata.file_metadata().schema().get_fields();

    let field = fields.get(column_index).ok_or_else(|| {
        anyhow::anyhow!(
            "{file_path}: column index {column_index} out of range ({} column(s))",
            fields.len()
        )
    })?;
    anyhow::ensure!(
        field.get_physical_type() == ParquetType::BYTE_ARRAY,
        "{file_path}: column `{}` is {:?}, not a string column",
        field.name(),
        field.get_physical_type()
    );

    let mut out: Vec<Box<str>> = Vec::with_capacity(nrows);
    for record in reader.get_row_iter(None)? {
        out.push(record?.get_string(column_index)?.clone().into_boxed_str());
    }
    Ok(out)
}

/// Read several named string columns in ONE pass, returning one
/// `Vec<Box<str>>` per requested name, in request order.
///
/// Prefer this over calling [`read_parquet_string_column`] once per column: that
/// reopens the file and walks every row again for each column, so reading `k`
/// columns costs `k` full scans instead of one.
///
/// Unlike [`read_parquet_string_column`], a non-string cell is not an error:
/// a number is formatted (`INT64 7` → `"7"`, see
/// [`crate::matrix::table::field_to_string`]) and a null yields `""`. The
/// callers are annotation and label tables, where a numeric label (a cluster
/// id) is still a label and a missing cell means "unlabelled"; a missing
/// *column* is still an error, since that is a schema mismatch rather than a
/// gap in the data.
pub fn read_parquet_string_columns_by_name(
    file_path: &str,
    wanted: &[&str],
) -> anyhow::Result<Vec<Vec<Box<str>>>> {
    let file = File::open(file_path).map_err(|e| anyhow::anyhow!("opening {file_path}: {e}"))?;
    let reader = SerializedFileReader::new(file)?;
    let fields = reader
        .metadata()
        .file_metadata()
        .schema()
        .get_fields()
        .to_vec();
    let idx: Vec<usize> = wanted
        .iter()
        .map(|w| {
            fields
                .iter()
                .position(|f| f.name() == *w)
                .ok_or_else(|| anyhow::anyhow!("column '{w}' not found in {file_path}"))
        })
        .collect::<anyhow::Result<_>>()?;

    let mut out: Vec<Vec<Box<str>>> = vec![Vec::new(); wanted.len()];
    for record in reader.get_row_iter(None)? {
        let row = record?;
        let cells: Vec<_> = row.get_column_iter().map(|(_, f)| f).collect();
        for (k, &j) in idx.iter().enumerate() {
            let v = cells.get(j).map_or_else(
                || Box::from(""),
                |f| crate::matrix::table::field_to_string(f),
            );
            out[k].push(v);
        }
    }
    Ok(out)
}

/// String columns and numeric columns of a table, each in request order.
pub type TableColumns = (Vec<Vec<Box<str>>>, Vec<Vec<f64>>);

/// Read a tidy mixed-type table (as written by [`write_named_table`] /
/// [`write_table`]): the named string columns and the named numeric columns,
/// each returned in request order. A missing column of either kind is an
/// error; a string column requested as numeric (or vice versa) is too.
pub fn read_table_columns(
    file_path: &str,
    string_cols: &[&str],
    numeric_cols: &[&str],
) -> anyhow::Result<TableColumns> {
    let strings = if string_cols.is_empty() {
        Vec::new()
    } else {
        read_parquet_string_columns_by_name(file_path, string_cols)?
    };
    if numeric_cols.is_empty() {
        return Ok((strings, Vec::new()));
    }
    let wanted: Vec<Box<str>> = numeric_cols.iter().map(|&c| c.into()).collect();
    let reader = ParquetReader::new(file_path, None, None, Some(&wanted))?;
    let ncols = reader.column_names.len();
    let nrows = if ncols == 0 {
        0
    } else {
        reader.row_major_data.len() / ncols
    };
    let numbers = numeric_cols
        .iter()
        .map(|&c| {
            let j = reader
                .column_names
                .iter()
                .position(|n| n.as_ref() == c)
                .ok_or_else(|| anyhow::anyhow!("numeric column '{c}' not found in {file_path}"))?;
            Ok((0..nrows)
                .map(|i| reader.row_major_data[i * ncols + j])
                .collect())
        })
        .collect::<anyhow::Result<Vec<Vec<f64>>>>()?;
    Ok((strings, numbers))
}

pub struct ParquetReader {
    pub row_major_data: Vec<f64>,
    pub row_names: Vec<Box<str>>,
    pub column_names: Vec<Box<str>>,
}

impl ParquetReader {
    /// Create a new parquet reader for a matrix with row and column
    /// names.
    ///
    /// * `row_name_index`: if `None`, no column is treated as row names
    ///   (row names will be generated as "0", "1", "2", ...)
    ///
    /// * `select_column_index`: if `None`, use all the other columns
    ///
    /// * `select_column_names`: if `None`, use all the other columns
    pub fn new(
        file_path: &str,
        row_name_index: Option<usize>,
        select_columns_index: Option<&[usize]>,
        select_columns_names: Option<&[Box<str>]>,
    ) -> anyhow::Result<Self> {
        let file = File::open(file_path)?;
        let reader = SerializedFileReader::new(file)?;
        let metadata = reader.metadata();
        let nrows = metadata.file_metadata().num_rows() as usize;
        let fields = metadata.file_metadata().schema().get_fields();

        let select_columns: HashSet<usize> = {
            let mut indices: HashSet<usize> = Default::default();

            // Add indices from `select_columns_index` if provided
            if let Some(select) = select_columns_index {
                indices.extend(select.iter().copied());
            }

            // Add indices from `select_columns_names` if provided
            if let Some(names) = select_columns_names {
                indices.extend(fields.iter().enumerate().filter_map(|(j, f)| {
                    if names.iter().any(|name| name.as_ref() == f.name()) {
                        Some(j)
                    } else {
                        None
                    }
                }));
            }

            // Default to all columns if neither is provided
            if indices.is_empty() {
                (0..fields.len()).collect()
            } else {
                indices
            }
        };

        // Get the type of the row name column for later use (if specified)
        let row_name_type = row_name_index.map(|idx| fields[idx].get_physical_type());

        let select_indices = fields
            .iter()
            .enumerate()
            .filter_map(|(j, f)| {
                // Exclude row_name_index column if specified
                let is_row_name_col = row_name_index == Some(j);
                if select_columns.contains(&j) && !is_row_name_col {
                    let tt = f.get_physical_type();
                    match tt {
                        parquet::basic::Type::FLOAT
                        | parquet::basic::Type::DOUBLE
                        | parquet::basic::Type::INT32
                        | parquet::basic::Type::INT64 => Some((tt, j)),
                        _ => None,
                    }
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        if select_indices.is_empty() {
            return Err(anyhow::anyhow!("no available columns"));
        }

        let ncols = select_indices.len();

        let column_names: Vec<Box<str>> = select_indices
            .iter()
            .map(|&(_, j)| fields[j].name().to_string().into_boxed_str())
            .collect();

        let row_iter = reader.get_row_iter(None)?;
        let mut row_names: Vec<Box<str>> = Vec::with_capacity(nrows);
        let mut row_major_data: Vec<f64> = Vec::with_capacity(nrows * ncols);

        for (row_counter, record) in row_iter.enumerate() {
            let row = record?;
            // Handle different column types for row names
            let row_name: Box<str> = match (row_name_index, row_name_type) {
                (Some(idx), Some(parquet::basic::Type::BYTE_ARRAY)) => {
                    row.get_string(idx)?.clone().into_boxed_str()
                }
                (Some(idx), Some(parquet::basic::Type::DOUBLE)) => {
                    row.get_double(idx)?.to_string().into_boxed_str()
                }
                (Some(idx), Some(parquet::basic::Type::FLOAT)) => {
                    row.get_float(idx)?.to_string().into_boxed_str()
                }
                (Some(idx), Some(parquet::basic::Type::INT32)) => {
                    row.get_int(idx)?.to_string().into_boxed_str()
                }
                (Some(idx), Some(parquet::basic::Type::INT64)) => {
                    row.get_long(idx)?.to_string().into_boxed_str()
                }
                (Some(idx), Some(_)) => {
                    // Fallback: try string, or use row index
                    row.get_string(idx)
                        .map(|s| s.clone().into_boxed_str())
                        .unwrap_or_else(|_| row_counter.to_string().into_boxed_str())
                }
                // No row name column specified, generate numeric names
                (None, _) | (_, None) => row_counter.to_string().into_boxed_str(),
            };
            row_names.push(row_name);

            let numbers: anyhow::Result<Vec<f64>> =
                select_indices
                    .iter()
                    .try_fold(Vec::with_capacity(ncols), |mut acc, &(tt, j)| {
                        let x = match tt {
                            parquet::basic::Type::DOUBLE => row.get_double(j)?,
                            parquet::basic::Type::FLOAT => row.get_float(j)? as f64,
                            parquet::basic::Type::INT32 => row.get_int(j)? as f64,
                            parquet::basic::Type::INT64 => row.get_long(j)? as f64,
                            _ => {
                                unimplemented!("we just support integer and float/double for now")
                            }
                        };
                        acc.push(x);
                        Ok(acc)
                    });

            row_major_data.extend(numbers?);
        }

        Ok(Self {
            row_major_data,
            row_names,
            column_names,
        })
    }
}

pub struct ParquetWriter {
    file: std::fs::File,
    schema: Arc<SchemaType>,
    writer_properties: Arc<WriterProperties>,
    row_names: Vec<ByteArray>,
}

impl ParquetWriter {
    /// Create a new parquet writer for a matrix with row and column
    /// names.
    ///
    /// * `file_path`: output file path
    ///
    /// * `shape`: number of rows and columns
    ///
    /// * `names`: for row and column names, respectively; if `None`, just add `[0, n)` numbers.
    ///
    #[allow(clippy::type_complexity)]
    pub fn new(
        file_path: &str,
        shape: (usize, usize),
        names: (Option<&[Box<str>]>, Option<&[Box<str>]>),
        column_types: Option<&[ParquetType]>,
        row_column_name: Option<&str>,
    ) -> anyhow::Result<Self> {
        let (nrows, ncols) = shape;
        let (row_names, column_names) = names;

        let schema = build_columns_schema(ncols, column_names, column_types, row_column_name)?;

        let file = std::fs::File::create(file_path)?;

        let zstd_level = ZstdLevel::try_new(5)?;
        let writer_properties = std::sync::Arc::new(
            WriterProperties::builder()
                .set_compression(Compression::ZSTD(zstd_level))
                .build(),
        );

        let row_names: Vec<ByteArray> = match row_names {
            Some(row_names) => row_names
                .iter()
                .map(|r| ByteArray::from(r.as_ref()))
                .collect(),
            None => (0..nrows)
                .map(|i| ByteArray::from(i.to_string().as_bytes()))
                .collect(),
        };

        Ok(Self {
            file,
            schema,
            writer_properties,
            row_names,
        })
    }

    pub fn row_names_vec(&self) -> &Vec<ByteArray> {
        &self.row_names
    }

    pub fn get_writer(&self) -> anyhow::Result<SerializedFileWriter<File>> {
        Ok(SerializedFileWriter::new(
            self.file.try_clone()?,
            self.schema.clone(),
            self.writer_properties.clone(),
        )?)
    }
}

/// write down a string vector of `Box<str>` to `row_group_writer` by creating a
/// new `column_writer`
pub fn parquet_add_bytearray<'a>(
    row_group_writer: &mut SerializedRowGroupWriter<'a, File>,
    data: &[ByteArray],
) -> anyhow::Result<()> {
    if let Some(mut column_writer) = row_group_writer.next_column()? {
        let typed_writer = column_writer.typed::<ByteArrayType>();
        typed_writer.write_batch(data, None, None)?;
        column_writer.close()?;
    }

    Ok(())
}

/// write down a string vector of `Box<str>` to `row_group_writer` by creating a
/// new `column_writer`
pub fn parquet_add_string_column<'a>(
    row_group_writer: &mut SerializedRowGroupWriter<'a, File>,
    data: &[Box<str>],
) -> anyhow::Result<()> {
    let data_bytearray = data
        .iter()
        .map(|x| ByteArray::from(x.as_ref()))
        .collect::<Vec<_>>();

    parquet_add_bytearray(row_group_writer, &data_bytearray)?;
    Ok(())
}

/// write down a numeric vector to `row_group_writer` by creating a
/// new `column_writer`
pub fn parquet_add_numeric_column<'a, T: 'static + num_traits::ToPrimitive>(
    row_group_writer: &mut SerializedRowGroupWriter<'a, File>,
    data: &[T],
) -> anyhow::Result<()> {
    if TypeId::of::<T>() == TypeId::of::<f64>() {
        if let Some(mut column_writer) = row_group_writer.next_column()? {
            let typed_writer = column_writer.typed::<DoubleType>();
            let data: Vec<f64> = data
                .iter()
                .map(|x| x.to_f64().expect("Failed to convert to f64"))
                .collect();
            typed_writer.write_batch(&data, None, None)?;
            column_writer.close()?;
        }
    } else if TypeId::of::<T>() == TypeId::of::<f32>() {
        if let Some(mut column_writer) = row_group_writer.next_column()? {
            let typed_writer = column_writer.typed::<FloatType>();
            let data: Vec<f32> = data
                .iter()
                .map(|x| x.to_f32().expect("Failed to convert to f32"))
                .collect();
            typed_writer.write_batch(&data, None, None)?;
            column_writer.close()?;
        }
    } else if TypeId::of::<T>() == TypeId::of::<i32>()
        || TypeId::of::<T>() == TypeId::of::<u32>()
        || TypeId::of::<T>() == TypeId::of::<usize>()
    {
        if let Some(mut column_writer) = row_group_writer.next_column()? {
            let typed_writer = column_writer.typed::<Int32Type>();
            let data: Vec<i32> = data
                .iter()
                .map(|x| x.to_i32().expect("Failed to convert to i32"))
                .collect();
            typed_writer.write_batch(&data, None, None)?;
            column_writer.close()?;
        }
    } else if TypeId::of::<T>() == TypeId::of::<i64>() || TypeId::of::<T>() == TypeId::of::<u64>() {
        if let Some(mut column_writer) = row_group_writer.next_column()? {
            let typed_writer = column_writer.typed::<Int64Type>();
            let data: Vec<i64> = data
                .iter()
                .map(|x| x.to_i64().expect("Failed to convert to i64"))
                .collect();
            typed_writer.write_batch(&data, None, None)?;
            column_writer.close()?;
        }
    } else {
        return Err(anyhow::anyhow!("Unsupported data type"));
    }

    Ok(())
}

/// One column of a tidy table for [`write_named_table`]. Borrows its data.
pub enum Column<'a> {
    Str(&'a [Box<str>]),
    F32(&'a [f32]),
    I32(&'a [i32]),
    I64(&'a [i64]),
}

/// Write a tidy mixed-type table to parquet: a leading string key column
/// (`row_col_name` / `row_names`) followed by `columns` of heterogeneous
/// type. All columns must have the same length as `row_names`.
pub fn write_named_table(
    file_path: &str,
    row_col_name: &str,
    row_names: &[Box<str>],
    columns: &[(Box<str>, Column)],
) -> anyhow::Result<()> {
    let ncols = columns.len();
    let col_names: Vec<Box<str>> = columns.iter().map(|(n, _)| n.clone()).collect();
    let col_types: Vec<ParquetType> = columns
        .iter()
        .map(|(_, c)| match c {
            Column::Str(_) => ParquetType::BYTE_ARRAY,
            Column::F32(_) => ParquetType::FLOAT,
            Column::I32(_) => ParquetType::INT32,
            Column::I64(_) => ParquetType::INT64,
        })
        .collect();

    let writer = ParquetWriter::new(
        file_path,
        (row_names.len(), ncols),
        (Some(row_names), Some(&col_names)),
        Some(&col_types),
        Some(row_col_name),
    )?;
    let row_ba = writer.row_names_vec().clone();
    let mut fw = writer.get_writer()?;
    let mut rg = fw.next_row_group()?;
    parquet_add_bytearray(&mut rg, &row_ba)?;
    for (_, col) in columns {
        match col {
            Column::Str(d) => parquet_add_string_column(&mut rg, d)?,
            Column::F32(d) => parquet_add_numeric_column(&mut rg, d)?,
            Column::I32(d) => parquet_add_numeric_column(&mut rg, d)?,
            Column::I64(d) => parquet_add_numeric_column(&mut rg, d)?,
        }
    }
    rg.close()?;
    fw.close()?;
    Ok(())
}

/// Write a mixed-type table to parquet with NO leading key column: exactly the
/// columns given, in that order. The keyless sibling of [`write_named_table`],
/// for tables whose name column sits elsewhere, or that have none.
pub fn write_table(file_path: &str, columns: &[(Box<str>, Column)]) -> anyhow::Result<()> {
    let n = columns.first().map_or(0, |(_, c)| c.len());
    anyhow::ensure!(
        columns.iter().all(|(_, c)| c.len() == n),
        "write_table: every column must have the same length"
    );
    let names: Vec<Box<str>> = columns.iter().map(|(n, _)| n.clone()).collect();
    let types: Vec<ParquetType> = columns.iter().map(|(_, c)| c.parquet_type()).collect();
    let schema = Arc::new(
        SchemaType::group_type_builder("2dMatrix")
            .with_fields(column_fields(&names, &types))
            .build()?,
    );
    let file = std::fs::File::create(file_path)?;
    let props = Arc::new(
        WriterProperties::builder()
            .set_compression(Compression::ZSTD(ZstdLevel::try_new(5)?))
            .build(),
    );
    let mut fw = SerializedFileWriter::new(file, schema, props)?;
    let mut rg = fw.next_row_group()?;
    for (_, col) in columns {
        match col {
            Column::Str(d) => parquet_add_string_column(&mut rg, d)?,
            Column::F32(d) => parquet_add_numeric_column(&mut rg, d)?,
            Column::I32(d) => parquet_add_numeric_column(&mut rg, d)?,
            Column::I64(d) => parquet_add_numeric_column(&mut rg, d)?,
        }
    }
    rg.close()?;
    fw.close()?;
    Ok(())
}

impl Column<'_> {
    /// Number of rows this column carries.
    pub fn len(&self) -> usize {
        match self {
            Column::Str(d) => d.len(),
            Column::F32(d) => d.len(),
            Column::I32(d) => d.len(),
            Column::I64(d) => d.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    fn parquet_type(&self) -> ParquetType {
        match self {
            Column::Str(_) => ParquetType::BYTE_ARRAY,
            Column::F32(_) => ParquetType::FLOAT,
            Column::I32(_) => ParquetType::INT32,
            Column::I64(_) => ParquetType::INT64,
        }
    }
}

/// One REQUIRED primitive field per (name, type); strings carry the UTF8
/// converted type.
fn column_fields(names: &[Box<str>], types: &[ParquetType]) -> Vec<Arc<SchemaType>> {
    names
        .iter()
        .zip(types)
        .map(|(name, &ty)| {
            let b = SchemaType::primitive_type_builder(name, ty)
                .with_repetition(parquet::basic::Repetition::REQUIRED);
            let b = if ty == ParquetType::BYTE_ARRAY {
                b.with_converted_type(ConvertedType::UTF8)
            } else {
                b
            };
            Arc::new(b.build().unwrap())
        })
        .collect()
}

fn build_columns_schema(
    ncols: usize,
    column_names: Option<&[Box<str>]>,
    column_types: Option<&[ParquetType]>,
    row_column_name: Option<&str>,
) -> anyhow::Result<Arc<SchemaType>> {
    if let Some(column_names) = column_names {
        if column_names.len() != ncols {
            return Err(anyhow::anyhow!(
                "Column names length ({}) does not match number of columns ({})",
                column_names.len(),
                ncols
            ));
        }
    }

    let row_col_name: Box<str> = row_column_name.unwrap_or("rowname").into();
    let _column_names: Vec<Box<str>> = (0..ncols).map(|x| x.to_string().into_boxed_str()).collect();
    let _column_types = (0..ncols).map(|_x| ParquetType::FLOAT).collect::<Vec<_>>();
    let column_names: &[Box<str>] = column_names.unwrap_or(&_column_names);
    let column_types: &[ParquetType] = column_types.unwrap_or(&_column_types);

    // The key column first, then the data columns, through one field builder.
    let mut names: Vec<Box<str>> = Vec::with_capacity(ncols + 1);
    let mut types: Vec<ParquetType> = Vec::with_capacity(ncols + 1);
    names.push(row_col_name);
    types.push(ParquetType::BYTE_ARRAY);
    names.extend(column_names.iter().cloned());
    types.extend(column_types.iter().copied());

    let schema = Arc::new(
        SchemaType::group_type_builder("2dMatrix")
            .with_fields(column_fields(&names, &types))
            .build()?,
    );

    Ok(schema)
}

#[cfg(test)]
#[path = "parquet_tests.rs"]
mod parquet_tests;