indexlake 0.3.0

use std::collections::HashMap;
use std::sync::Arc;

use arrow::array::*;
use arrow::datatypes::{DataType, TimeUnit, i256};
use futures::StreamExt;
use uuid::Uuid;

use crate::catalog::{
    Catalog, CatalogSchema, DataFileRecord, IndexFileRecord, InlineIndexRecord, RowValidity,
    TransactionHelper, rows_to_record_batch,
};
use crate::index::IndexBuilder;
use crate::storage::{DataFileFormat, build_parquet_writer};
use crate::table::{Table, TableSchemaRef};
use crate::utils::{
    extract_row_id_array_from_record_batch, fixed_size_binary_array_to_uuids, rewrite_batch_schema,
    serialize_array,
};
use crate::{ILError, ILResult};

pub struct TableInsertion {
    pub data: Vec<RecordBatch>,
    pub force_inline: bool,
    pub ignore_row_id: bool,
}

impl TableInsertion {
    pub fn new(data: Vec<RecordBatch>) -> Self {
        Self {
            data,
            force_inline: false,
            ignore_row_id: true,
        }
    }

    pub fn with_force_inline(mut self, force_inline: bool) -> Self {
        self.force_inline = force_inline;
        self
    }

    pub fn rewrite_columns(mut self, field_name_id_map: &HashMap<String, Uuid>) -> ILResult<Self> {
        let mut new_data = Vec::with_capacity(self.data.len());
        for batch in self.data {
            let new_batch = rewrite_batch_schema(&batch, field_name_id_map)?;
            new_data.push(new_batch);
        }
        self.data = new_data;
        Ok(self)
    }
}

pub(crate) async fn process_insert_into_inline_rows(
    tx_helper: &mut TransactionHelper,
    table: &Table,
    batches: &[RecordBatch],
) -> ILResult<()> {
    if batches.is_empty() {
        return Ok(());
    }

    let mut mergeable_index_builders = table.index_manager.new_mergeable_index_builders()?;
    let mut non_mergeable_index_builders =
        table.index_manager.new_non_mergeable_index_builders()?;

    // insert inline rows
    for batch in batches {
        let sql_values = record_batch_to_sql_values(batch, tx_helper.catalog.as_ref())?;

        let inline_field_names = batch
            .schema()
            .fields()
            .iter()
            .map(|field| field.name().clone())
            .collect::<Vec<_>>();

        tx_helper
            .insert_inline_rows(&table.table_id, &inline_field_names, sql_values)
            .await?;
    }

    // append index builders
    for batch in batches {
        for builder in mergeable_index_builders.iter_mut() {
            builder.append(batch)?;
        }
    }

    append_non_mergeable_index_builders(
        tx_helper,
        &table.table_id,
        &table.table_schema,
        &mut non_mergeable_index_builders,
    )
    .await?;

    // build inline index records
    let mut inline_index_records = Vec::new();
    for builder in mergeable_index_builders.iter_mut() {
        let mut index_data = Vec::new();
        builder.write_bytes(&mut index_data)?;
        inline_index_records.push(InlineIndexRecord {
            index_id: builder.index_def().index_id,
            index_data,
        });
    }
    for builder in non_mergeable_index_builders.iter_mut() {
        let mut index_data = Vec::new();
        builder.write_bytes(&mut index_data)?;
        inline_index_records.push(InlineIndexRecord {
            index_id: builder.index_def().index_id,
            index_data,
        });
    }

    // clean non-mergeable inline index records
    tx_helper
        .delete_inline_indexes(
            &non_mergeable_index_builders
                .iter()
                .map(|builder| builder.index_def().index_id)
                .collect::<Vec<_>>(),
        )
        .await?;

    // insert inline index records
    tx_helper
        .insert_inline_indexes(&inline_index_records)
        .await?;

    Ok(())
}

pub(crate) async fn process_bypass_insert(
    tx_helper: &mut TransactionHelper,
    table: &Table,
    batches: &[RecordBatch],
) -> ILResult<()> {
    let data_file_id = uuid::Uuid::now_v7();
    let relative_path = DataFileRecord::build_relative_path(
        &table.namespace_id,
        &table.table_id,
        &data_file_id,
        table.config.preferred_data_file_format,
    );

    let mut index_builders = table.index_manager.new_index_builders()?;

    let row_ids = match table.config.preferred_data_file_format {
        DataFileFormat::ParquetV1 | DataFileFormat::ParquetV2 => {
            write_parquet_file(table, &relative_path, batches, &mut index_builders).await?
        }
    };
    let record_count = row_ids.len();

    tx_helper
        .insert_data_files(&[DataFileRecord {
            data_file_id,
            table_id: table.table_id,
            format: table.config.preferred_data_file_format,
            relative_path: relative_path.clone(),
            record_count: record_count as i64,
            validity: RowValidity::new(record_count),
        }])
        .await?;

    // update index files
    let mut index_file_records = Vec::new();
    for index_builder in index_builders.iter_mut() {
        let index_file_id = Uuid::now_v7();
        let relative_path = IndexFileRecord::build_relative_path(
            &table.namespace_id,
            &table.table_id,
            &index_file_id,
        );
        let output_file = table.storage.create(&relative_path).await?;
        index_builder.write_file(output_file).await?;
        index_file_records.push(IndexFileRecord {
            index_file_id,
            table_id: table.table_id,
            index_id: index_builder.index_def().index_id,
            data_file_id,
            relative_path,
        });
    }

    tx_helper.insert_index_files(&index_file_records).await?;

    Ok(())
}

async fn write_parquet_file(
    table: &Table,
    relative_path: &str,
    batches: &[RecordBatch],
    index_builders: &mut Vec<Box<dyn IndexBuilder>>,
) -> ILResult<Vec<Uuid>> {
    let output_file = table.storage.create(relative_path).await?;

    let mut arrow_writer = build_parquet_writer(
        output_file,
        table.table_schema.arrow_schema.clone(),
        table.config.parquet_row_group_size,
        table.config.preferred_data_file_format,
    )?;

    let mut row_ids = Vec::new();
    for batch in batches {
        let row_id_array = extract_row_id_array_from_record_batch(batch)?;
        row_ids.extend(fixed_size_binary_array_to_uuids(&row_id_array)?);

        arrow_writer.write(batch).await?;
        for builder in index_builders.iter_mut() {
            builder.append(batch)?;
        }
    }

    arrow_writer.close().await?;

    Ok(row_ids)
}

macro_rules! extract_sql_values {
    ($array:expr, $array_ty:ty, $convert:expr) => {{
        let mut sql_values = Vec::with_capacity($array.len());
        let array = $array.as_any().downcast_ref::<$array_ty>().ok_or_else(|| {
            ILError::internal(format!(
                "Failed to downcast array to {}",
                stringify!($array_ty),
            ))
        })?;
        for v in array.iter() {
            sql_values.push(match v {
                Some(v) => $convert(v)?,
                None => "NULL".to_string(),
            });
        }
        sql_values
    }};
}

pub(crate) fn array_to_sql_literals(
    array: &dyn Array,
    catalog: &dyn Catalog,
) -> ILResult<Vec<String>> {
    let data_type = array.data_type();
    let literals = match data_type {
        DataType::Boolean => {
            extract_sql_values!(array, BooleanArray, |v: bool| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Int8 => {
            extract_sql_values!(array, Int8Array, |v: i8| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Int16 => {
            extract_sql_values!(array, Int16Array, |v: i16| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Int32 => {
            extract_sql_values!(array, Int32Array, |v: i32| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Int64 => {
            extract_sql_values!(array, Int64Array, |v: i64| Ok::<_, ILError>(v.to_string()))
        }
        DataType::UInt8 => {
            extract_sql_values!(array, UInt8Array, |v: u8| Ok::<_, ILError>(v.to_string()))
        }
        DataType::UInt16 => {
            extract_sql_values!(array, UInt16Array, |v: u16| Ok::<_, ILError>(v.to_string()))
        }
        DataType::UInt32 => {
            extract_sql_values!(array, UInt32Array, |v: u32| Ok::<_, ILError>(v.to_string()))
        }
        DataType::UInt64 => {
            extract_sql_values!(array, UInt64Array, |v: u64| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Float32 => {
            extract_sql_values!(array, Float32Array, |v: f32| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Float64 => {
            extract_sql_values!(array, Float64Array, |v: f64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Timestamp(TimeUnit::Second, _) => {
            extract_sql_values!(array, TimestampSecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Timestamp(TimeUnit::Millisecond, _) => {
            extract_sql_values!(array, TimestampMillisecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Timestamp(TimeUnit::Microsecond, _) => {
            extract_sql_values!(array, TimestampMicrosecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
            extract_sql_values!(array, TimestampNanosecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Date32 => {
            extract_sql_values!(array, Date32Array, |v: i32| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Date64 => {
            extract_sql_values!(array, Date64Array, |v: i64| Ok::<_, ILError>(v.to_string()))
        }
        DataType::Time32(TimeUnit::Second) => {
            extract_sql_values!(array, Time32SecondArray, |v: i32| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Time32(TimeUnit::Millisecond) => {
            extract_sql_values!(array, Time32MillisecondArray, |v: i32| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Time64(TimeUnit::Microsecond) => {
            extract_sql_values!(array, Time64MicrosecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Time64(TimeUnit::Nanosecond) => {
            extract_sql_values!(array, Time64NanosecondArray, |v: i64| {
                Ok::<_, ILError>(v.to_string())
            })
        }
        DataType::Binary => {
            extract_sql_values!(array, BinaryArray, |v: &[u8]| {
                Ok::<_, ILError>(catalog.sql_binary_literal(v))
            })
        }
        DataType::FixedSizeBinary(size) => {
            if *size == 16 {
                extract_sql_values!(array, FixedSizeBinaryArray, |v: &[u8]| {
                    Ok::<_, ILError>(catalog.sql_uuid_literal(&Uuid::from_slice(v)?))
                })
            } else {
                extract_sql_values!(array, FixedSizeBinaryArray, |v: &[u8]| {
                    Ok::<_, ILError>(catalog.sql_binary_literal(v))
                })
            }
        }
        DataType::LargeBinary => {
            extract_sql_values!(array, LargeBinaryArray, |v: &[u8]| {
                Ok::<_, ILError>(catalog.sql_binary_literal(v))
            })
        }
        DataType::BinaryView => {
            extract_sql_values!(array, BinaryViewArray, |v: &[u8]| {
                Ok::<_, ILError>(catalog.sql_binary_literal(v))
            })
        }
        DataType::Utf8 => {
            extract_sql_values!(array, StringArray, |v: &str| {
                Ok::<_, ILError>(catalog.sql_string_literal(v))
            })
        }
        DataType::LargeUtf8 => {
            extract_sql_values!(array, LargeStringArray, |v: &str| {
                Ok::<_, ILError>(catalog.sql_string_literal(v))
            })
        }
        DataType::Utf8View => {
            extract_sql_values!(array, StringViewArray, |v: &str| {
                Ok::<_, ILError>(catalog.sql_string_literal(v))
            })
        }
        DataType::List(inner_field) => {
            let mut sql_values = Vec::with_capacity(array.len());
            let array = array
                .as_any()
                .downcast_ref::<ListArray>()
                .ok_or_else(|| ILError::internal("Failed to downcast array to ListArray"))?;
            for v in array.iter() {
                sql_values.push(match v {
                    Some(v) => {
                        catalog.sql_binary_literal(&serialize_array(v, inner_field.clone())?)
                    }
                    None => "NULL".to_string(),
                });
            }
            sql_values
        }
        DataType::FixedSizeList(inner_field, _len) => {
            let mut sql_values = Vec::with_capacity(array.len());
            let array = array
                .as_any()
                .downcast_ref::<FixedSizeListArray>()
                .ok_or_else(|| {
                    ILError::internal("Failed to downcast array to FixedSizeListArray")
                })?;
            for v in array.iter() {
                sql_values.push(match v {
                    Some(v) => {
                        catalog.sql_binary_literal(&serialize_array(v, inner_field.clone())?)
                    }
                    None => "NULL".to_string(),
                });
            }
            sql_values
        }
        DataType::LargeList(inner_field) => {
            let mut sql_values = Vec::with_capacity(array.len());
            let array = array
                .as_any()
                .downcast_ref::<LargeListArray>()
                .ok_or_else(|| ILError::internal("Failed to downcast array to LargeListArray"))?;
            for v in array.iter() {
                sql_values.push(match v {
                    Some(v) => {
                        catalog.sql_binary_literal(&serialize_array(v, inner_field.clone())?)
                    }
                    None => "NULL".to_string(),
                });
            }
            sql_values
        }
        DataType::Decimal128(..) => {
            extract_sql_values!(array, Decimal128Array, |v: i128| {
                Ok::<_, ILError>(format!("'{v}'"))
            })
        }
        DataType::Decimal256(..) => {
            extract_sql_values!(array, Decimal256Array, |v: i256| {
                Ok::<_, ILError>(format!("'{v}'"))
            })
        }
        _ => {
            return Err(ILError::not_supported(format!(
                "Unsupported array type: {data_type}",
            )));
        }
    };
    Ok(literals)
}

pub(crate) fn record_batch_to_sql_values(
    record: &RecordBatch,
    catalog: &dyn Catalog,
) -> ILResult<Vec<Vec<String>>> {
    let mut column_values_list = Vec::with_capacity(record.num_columns());
    for array in record.columns() {
        let column_values = array_to_sql_literals(array, catalog)?;
        column_values_list.push(column_values);
    }
    Ok(column_values_list)
}

pub(crate) async fn append_non_mergeable_index_builders(
    tx_helper: &mut TransactionHelper,
    table_id: &Uuid,
    table_schema: &TableSchemaRef,
    index_builders: &mut Vec<Box<dyn IndexBuilder>>,
) -> ILResult<()> {
    let catalog_schema = Arc::new(CatalogSchema::from_arrow(&table_schema.arrow_schema)?);

    let row_stream = tx_helper
        .scan_inline_rows(table_id, &catalog_schema, &[], None, None)
        .await?;
    let mut inline_stream = row_stream.chunks(100).map(move |rows| {
        let rows = rows.into_iter().collect::<ILResult<Vec<_>>>()?;
        let batch = rows_to_record_batch(&table_schema.arrow_schema, &rows)?;
        Ok::<_, ILError>(batch)
    });

    while let Some(batch) = inline_stream.next().await {
        let batch = batch?;
        for builder in index_builders.iter_mut() {
            builder.append(&batch)?;
        }
    }
    Ok(())
}