alopex_dataframe/ops/

nulls.rs

use std::sync::Arc;

use arrow::array::{
    Array, BooleanArray, BooleanBuilder, Float32Array, Float32Builder, Float64Array,
    Float64Builder, Int16Array, Int16Builder, Int32Array, Int32Builder, Int64Array, Int64Builder,
    Int8Array, Int8Builder, StringArray, StringBuilder, UInt16Array, UInt16Builder, UInt32Array,
    UInt32Builder, UInt64Array, UInt64Builder, UInt8Array, UInt8Builder,
};
use arrow::datatypes::{DataType, Field, Schema};
use arrow::record_batch::RecordBatch;

use crate::expr::Scalar;
use crate::ops::{FillNull, FillNullStrategy};
use crate::{DataFrameError, Result};

pub fn fill_null_batches(input: Vec<RecordBatch>, fill: &FillNull) -> Result<Vec<RecordBatch>> {
    let batch = concat_batches(&input)?;
    if batch.num_rows() == 0 {
        return Ok(vec![batch]);
    }

    let arrays = match fill {
        FillNull::Value(value) => fill_with_scalar(&batch, value)?,
        FillNull::Strategy(strategy) => fill_with_strategy(&batch, *strategy)?,
    };

    let batch = RecordBatch::try_new(batch.schema(), arrays).map_err(|e| {
        DataFrameError::schema_mismatch(format!("failed to build RecordBatch: {e}"))
    })?;
    Ok(vec![batch])
}

pub fn drop_nulls_batches(
    input: Vec<RecordBatch>,
    subset: Option<&[String]>,
) -> Result<Vec<RecordBatch>> {
    let batch = concat_batches(&input)?;
    if batch.num_rows() == 0 {
        return Ok(vec![batch]);
    }
    let indices = resolve_subset(&batch, subset)?;

    let mut mask_builder = BooleanBuilder::with_capacity(batch.num_rows());
    for row in 0..batch.num_rows() {
        let mut keep = true;
        for idx in &indices {
            if batch.column(*idx).is_null(row) {
                keep = false;
                break;
            }
        }
        mask_builder.append_value(keep);
    }
    let mask = mask_builder.finish();
    let batch = arrow::compute::filter_record_batch(&batch, &mask)
        .map_err(|source| DataFrameError::Arrow { source })?;
    Ok(vec![batch])
}

pub fn null_count_batches(input: Vec<RecordBatch>) -> Result<Vec<RecordBatch>> {
    if input.is_empty() {
        return Ok(vec![RecordBatch::new_empty(Arc::new(Schema::empty()))]);
    }

    let schema = input[0].schema();
    let mut counts = vec![0_u64; schema.fields().len()];
    for batch in &input {
        for (idx, col) in batch.columns().iter().enumerate() {
            counts[idx] += col.null_count() as u64;
        }
    }

    let mut fields = Vec::with_capacity(counts.len());
    let mut arrays = Vec::with_capacity(counts.len());
    for (idx, field) in schema.fields().iter().enumerate() {
        fields.push(Field::new(field.name(), DataType::UInt64, false));
        let array = UInt64Array::from(vec![counts[idx]]);
        arrays.push(Arc::new(array) as _);
    }

    let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).map_err(|e| {
        DataFrameError::schema_mismatch(format!("failed to build RecordBatch: {e}"))
    })?;
    Ok(vec![batch])
}

fn concat_batches(batches: &[RecordBatch]) -> Result<RecordBatch> {
    if batches.is_empty() {
        return Ok(RecordBatch::new_empty(Arc::new(Schema::empty())));
    }
    let schema = batches[0].schema();
    if batches.len() == 1 {
        return Ok(batches[0].clone());
    }
    arrow::compute::concat_batches(&schema, batches)
        .map_err(|source| DataFrameError::Arrow { source })
}

fn resolve_subset(batch: &RecordBatch, subset: Option<&[String]>) -> Result<Vec<usize>> {
    let schema = batch.schema();
    let indices = match subset {
        Some(cols) => {
            if cols.is_empty() {
                return Err(DataFrameError::invalid_operation(
                    "drop_nulls subset must be non-empty",
                ));
            }
            cols.iter()
                .map(|name| {
                    schema
                        .fields()
                        .iter()
                        .position(|f| f.name() == name)
                        .ok_or_else(|| DataFrameError::column_not_found(name.clone()))
                })
                .collect::<Result<Vec<_>>>()?
        }
        None => (0..schema.fields().len()).collect(),
    };
    Ok(indices)
}

fn fill_with_scalar(batch: &RecordBatch, value: &Scalar) -> Result<Vec<Arc<dyn Array>>> {
    let mut arrays = Vec::with_capacity(batch.num_columns());
    for col in batch.columns() {
        let filled: Arc<dyn Array> = match col.data_type() {
            DataType::Boolean => Arc::new(fill_boolean(col, value)?),
            DataType::Int8 => Arc::new(fill_int8(col, value)?),
            DataType::Int16 => Arc::new(fill_int16(col, value)?),
            DataType::Int32 => Arc::new(fill_int32(col, value)?),
            DataType::Int64 => Arc::new(fill_int64(col, value)?),
            DataType::UInt8 => Arc::new(fill_uint8(col, value)?),
            DataType::UInt16 => Arc::new(fill_uint16(col, value)?),
            DataType::UInt32 => Arc::new(fill_uint32(col, value)?),
            DataType::UInt64 => Arc::new(fill_uint64(col, value)?),
            DataType::Float32 => Arc::new(fill_float32(col, value)?),
            DataType::Float64 => Arc::new(fill_float64(col, value)?),
            DataType::Utf8 => Arc::new(fill_utf8(col, value)?),
            other => {
                return Err(DataFrameError::type_mismatch(
                    None::<String>,
                    other.to_string(),
                    format!("{value:?}"),
                ))
            }
        };
        arrays.push(filled);
    }
    Ok(arrays)
}

fn fill_with_strategy(
    batch: &RecordBatch,
    strategy: FillNullStrategy,
) -> Result<Vec<Arc<dyn Array>>> {
    let mut arrays = Vec::with_capacity(batch.num_columns());
    for col in batch.columns() {
        let filled: Arc<dyn Array> = match strategy {
            FillNullStrategy::Forward => Arc::new(fill_forward(col)?),
            FillNullStrategy::Backward => Arc::new(fill_backward(col)?),
            FillNullStrategy::Min => Arc::new(fill_min(col)?),
            FillNullStrategy::Max => Arc::new(fill_max(col)?),
            FillNullStrategy::Mean => Arc::new(fill_mean(col)?),
            FillNullStrategy::Zero => Arc::new(fill_numeric_constant(col, 0)?),
            FillNullStrategy::One => Arc::new(fill_numeric_constant(col, 1)?),
        };
        arrays.push(filled);
    }
    Ok(arrays)
}

fn fill_boolean(col: &Arc<dyn Array>, value: &Scalar) -> Result<BooleanArray> {
    let array = col
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad BooleanArray downcast"))?;
    let fill = match value {
        Scalar::Boolean(v) => *v,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Boolean".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = BooleanBuilder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_int8(col: &Arc<dyn Array>, value: &Scalar) -> Result<Int8Array> {
    let array = col
        .as_any()
        .downcast_ref::<Int8Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int8Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => i8::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "Int8", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Int8".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Int8Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_int16(col: &Arc<dyn Array>, value: &Scalar) -> Result<Int16Array> {
    let array = col
        .as_any()
        .downcast_ref::<Int16Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int16Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => i16::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "Int16", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Int16".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Int16Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_int32(col: &Arc<dyn Array>, value: &Scalar) -> Result<Int32Array> {
    let array = col
        .as_any()
        .downcast_ref::<Int32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int32Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => i32::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "Int32", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Int32".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Int32Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_int64(col: &Arc<dyn Array>, value: &Scalar) -> Result<Int64Array> {
    let array = col
        .as_any()
        .downcast_ref::<Int64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int64Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => *v,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Int64".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Int64Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_uint8(col: &Arc<dyn Array>, value: &Scalar) -> Result<UInt8Array> {
    let array = col
        .as_any()
        .downcast_ref::<UInt8Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt8Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => u8::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "UInt8", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "UInt8".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = UInt8Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_uint16(col: &Arc<dyn Array>, value: &Scalar) -> Result<UInt16Array> {
    let array = col
        .as_any()
        .downcast_ref::<UInt16Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt16Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => u16::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "UInt16", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "UInt16".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = UInt16Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_uint32(col: &Arc<dyn Array>, value: &Scalar) -> Result<UInt32Array> {
    let array = col
        .as_any()
        .downcast_ref::<UInt32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt32Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => u32::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "UInt32", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "UInt32".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = UInt32Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_uint64(col: &Arc<dyn Array>, value: &Scalar) -> Result<UInt64Array> {
    let array = col
        .as_any()
        .downcast_ref::<UInt64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt64Array downcast"))?;
    let fill = match value {
        Scalar::Int64(v) => u64::try_from(*v)
            .map_err(|_| DataFrameError::type_mismatch(None::<String>, "UInt64", "out of range"))?,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "UInt64".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = UInt64Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_float32(col: &Arc<dyn Array>, value: &Scalar) -> Result<Float32Array> {
    let array = col
        .as_any()
        .downcast_ref::<Float32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Float32Array downcast"))?;
    let fill = match value {
        Scalar::Float64(v) => *v as f32,
        Scalar::Int64(v) => *v as f32,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Float32".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Float32Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_float64(col: &Arc<dyn Array>, value: &Scalar) -> Result<Float64Array> {
    let array = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Float64Array downcast"))?;
    let fill = match value {
        Scalar::Float64(v) => *v,
        Scalar::Int64(v) => *v as f64,
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Float64".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = Float64Builder::with_capacity(array.len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_utf8(col: &Arc<dyn Array>, value: &Scalar) -> Result<StringArray> {
    let array = col
        .as_any()
        .downcast_ref::<StringArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad StringArray downcast"))?;
    let fill = match value {
        Scalar::Utf8(v) => v.as_str(),
        Scalar::Null => return Ok(array.clone()),
        other => {
            return Err(DataFrameError::type_mismatch(
                None::<String>,
                "Utf8".to_string(),
                format!("{other:?}"),
            ))
        }
    };

    let mut builder = StringBuilder::with_capacity(array.len(), array.value_data().len());
    for i in 0..array.len() {
        if array.is_null(i) {
            builder.append_value(fill);
        } else {
            builder.append_value(array.value(i));
        }
    }
    Ok(builder.finish())
}

fn fill_forward(col: &Arc<dyn Array>) -> Result<Arc<dyn Array>> {
    match col.data_type() {
        DataType::Boolean => Ok(Arc::new(fill_forward_bool(col)?)),
        DataType::Int8 => Ok(Arc::new(fill_forward_i8(col)?)),
        DataType::Int16 => Ok(Arc::new(fill_forward_i16(col)?)),
        DataType::Int32 => Ok(Arc::new(fill_forward_i32(col)?)),
        DataType::Int64 => Ok(Arc::new(fill_forward_i64(col)?)),
        DataType::UInt8 => Ok(Arc::new(fill_forward_u8(col)?)),
        DataType::UInt16 => Ok(Arc::new(fill_forward_u16(col)?)),
        DataType::UInt32 => Ok(Arc::new(fill_forward_u32(col)?)),
        DataType::UInt64 => Ok(Arc::new(fill_forward_u64(col)?)),
        DataType::Float32 => Ok(Arc::new(fill_forward_f32(col)?)),
        DataType::Float64 => Ok(Arc::new(fill_forward_f64(col)?)),
        DataType::Utf8 => Ok(Arc::new(fill_forward_utf8(col)?)),
        other => Err(DataFrameError::invalid_operation(format!(
            "unsupported fill_null forward type {other:?}",
        ))),
    }
}

fn fill_backward(col: &Arc<dyn Array>) -> Result<Arc<dyn Array>> {
    match col.data_type() {
        DataType::Boolean => Ok(Arc::new(fill_backward_bool(col)?)),
        DataType::Int8 => Ok(Arc::new(fill_backward_i8(col)?)),
        DataType::Int16 => Ok(Arc::new(fill_backward_i16(col)?)),
        DataType::Int32 => Ok(Arc::new(fill_backward_i32(col)?)),
        DataType::Int64 => Ok(Arc::new(fill_backward_i64(col)?)),
        DataType::UInt8 => Ok(Arc::new(fill_backward_u8(col)?)),
        DataType::UInt16 => Ok(Arc::new(fill_backward_u16(col)?)),
        DataType::UInt32 => Ok(Arc::new(fill_backward_u32(col)?)),
        DataType::UInt64 => Ok(Arc::new(fill_backward_u64(col)?)),
        DataType::Float32 => Ok(Arc::new(fill_backward_f32(col)?)),
        DataType::Float64 => Ok(Arc::new(fill_backward_f64(col)?)),
        DataType::Utf8 => Ok(Arc::new(fill_backward_utf8(col)?)),
        other => Err(DataFrameError::invalid_operation(format!(
            "unsupported fill_null backward type {other:?}",
        ))),
    }
}

fn fill_min(col: &Arc<dyn Array>) -> Result<Arc<dyn Array>> {
    fill_numeric_stat(col, Stat::Min)
}

fn fill_max(col: &Arc<dyn Array>) -> Result<Arc<dyn Array>> {
    fill_numeric_stat(col, Stat::Max)
}

fn fill_mean(col: &Arc<dyn Array>) -> Result<Arc<dyn Array>> {
    fill_numeric_stat(col, Stat::Mean)
}

fn fill_numeric_constant(col: &Arc<dyn Array>, value: i64) -> Result<Arc<dyn Array>> {
    match col.data_type() {
        DataType::Int8 => Ok(Arc::new(fill_int8(col, &Scalar::Int64(value))?)),
        DataType::Int16 => Ok(Arc::new(fill_int16(col, &Scalar::Int64(value))?)),
        DataType::Int32 => Ok(Arc::new(fill_int32(col, &Scalar::Int64(value))?)),
        DataType::Int64 => Ok(Arc::new(fill_int64(col, &Scalar::Int64(value))?)),
        DataType::UInt8 => Ok(Arc::new(fill_uint8(col, &Scalar::Int64(value))?)),
        DataType::UInt16 => Ok(Arc::new(fill_uint16(col, &Scalar::Int64(value))?)),
        DataType::UInt32 => Ok(Arc::new(fill_uint32(col, &Scalar::Int64(value))?)),
        DataType::UInt64 => Ok(Arc::new(fill_uint64(col, &Scalar::Int64(value))?)),
        DataType::Float32 => Ok(Arc::new(fill_float32(col, &Scalar::Int64(value))?)),
        DataType::Float64 => Ok(Arc::new(fill_float64(col, &Scalar::Int64(value))?)),
        other => Err(DataFrameError::type_mismatch(
            None::<String>,
            "numeric".to_string(),
            other.to_string(),
        )),
    }
}

enum Stat {
    Min,
    Max,
    Mean,
}

fn fill_numeric_stat(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    match col.data_type() {
        DataType::Int8 => fill_stat_i8(col, stat),
        DataType::Int16 => fill_stat_i16(col, stat),
        DataType::Int32 => fill_stat_i32(col, stat),
        DataType::Int64 => fill_stat_i64(col, stat),
        DataType::UInt8 => fill_stat_u8(col, stat),
        DataType::UInt16 => fill_stat_u16(col, stat),
        DataType::UInt32 => fill_stat_u32(col, stat),
        DataType::UInt64 => fill_stat_u64(col, stat),
        DataType::Float32 => fill_stat_f32(col, stat),
        DataType::Float64 => fill_stat_f64(col, stat),
        other => Err(DataFrameError::type_mismatch(
            None::<String>,
            "numeric".to_string(),
            other.to_string(),
        )),
    }
}

fn fill_forward_bool(col: &Arc<dyn Array>) -> Result<BooleanArray> {
    let array = col
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad BooleanArray downcast"))?;
    let mut builder = BooleanBuilder::with_capacity(array.len());
    let mut last: Option<bool> = None;
    for i in 0..array.len() {
        if array.is_null(i) {
            match last {
                Some(v) => builder.append_value(v),
                None => builder.append_null(),
            }
        } else {
            let v = array.value(i);
            last = Some(v);
            builder.append_value(v);
        }
    }
    Ok(builder.finish())
}

fn fill_backward_bool(col: &Arc<dyn Array>) -> Result<BooleanArray> {
    let array = col
        .as_any()
        .downcast_ref::<BooleanArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad BooleanArray downcast"))?;
    let mut tmp: Vec<Option<bool>> = Vec::with_capacity(array.len());
    let mut next: Option<bool> = None;
    for i in (0..array.len()).rev() {
        if array.is_null(i) {
            tmp.push(next);
        } else {
            let v = array.value(i);
            next = Some(v);
            tmp.push(Some(v));
        }
    }
    tmp.reverse();
    let mut builder = BooleanBuilder::with_capacity(array.len());
    for v in tmp {
        match v {
            Some(v) => builder.append_value(v),
            None => builder.append_null(),
        }
    }
    Ok(builder.finish())
}

macro_rules! fill_forward_numeric {
    ($name:ident, $array_ty:ty, $builder_ty:ty, $cast:expr) => {
        fn $name(col: &Arc<dyn Array>) -> Result<$array_ty> {
            let array = col
                .as_any()
                .downcast_ref::<$array_ty>()
                .ok_or_else(|| DataFrameError::invalid_operation("bad array downcast"))?;
            let mut builder = <$builder_ty>::with_capacity(array.len());
            let mut last = None;
            for i in 0..array.len() {
                if array.is_null(i) {
                    match last {
                        Some(v) => builder.append_value(v),
                        None => builder.append_null(),
                    }
                } else {
                    let v = array.value(i);
                    let v = $cast(v);
                    last = Some(v);
                    builder.append_value(v);
                }
            }
            Ok(builder.finish())
        }
    };
}

macro_rules! fill_backward_numeric {
    ($name:ident, $array_ty:ty, $builder_ty:ty, $cast:expr) => {
        fn $name(col: &Arc<dyn Array>) -> Result<$array_ty> {
            let array = col
                .as_any()
                .downcast_ref::<$array_ty>()
                .ok_or_else(|| DataFrameError::invalid_operation("bad array downcast"))?;
            let mut tmp = Vec::with_capacity(array.len());
            let mut next = None;
            for i in (0..array.len()).rev() {
                if array.is_null(i) {
                    tmp.push(next);
                } else {
                    let v = $cast(array.value(i));
                    next = Some(v);
                    tmp.push(Some(v));
                }
            }
            tmp.reverse();
            let mut builder = <$builder_ty>::with_capacity(array.len());
            for v in tmp {
                match v {
                    Some(v) => builder.append_value(v),
                    None => builder.append_null(),
                }
            }
            Ok(builder.finish())
        }
    };
}

fill_forward_numeric!(fill_forward_i8, Int8Array, Int8Builder, |v: i8| v);
fill_forward_numeric!(fill_forward_i16, Int16Array, Int16Builder, |v: i16| v);
fill_forward_numeric!(fill_forward_i32, Int32Array, Int32Builder, |v: i32| v);
fill_forward_numeric!(fill_forward_i64, Int64Array, Int64Builder, |v: i64| v);
fill_forward_numeric!(fill_forward_u8, UInt8Array, UInt8Builder, |v: u8| v);
fill_forward_numeric!(fill_forward_u16, UInt16Array, UInt16Builder, |v: u16| v);
fill_forward_numeric!(fill_forward_u32, UInt32Array, UInt32Builder, |v: u32| v);
fill_forward_numeric!(fill_forward_u64, UInt64Array, UInt64Builder, |v: u64| v);
fill_forward_numeric!(fill_forward_f32, Float32Array, Float32Builder, |v: f32| v);
fill_forward_numeric!(fill_forward_f64, Float64Array, Float64Builder, |v: f64| v);

fill_backward_numeric!(fill_backward_i8, Int8Array, Int8Builder, |v: i8| v);
fill_backward_numeric!(fill_backward_i16, Int16Array, Int16Builder, |v: i16| v);
fill_backward_numeric!(fill_backward_i32, Int32Array, Int32Builder, |v: i32| v);
fill_backward_numeric!(fill_backward_i64, Int64Array, Int64Builder, |v: i64| v);
fill_backward_numeric!(fill_backward_u8, UInt8Array, UInt8Builder, |v: u8| v);
fill_backward_numeric!(fill_backward_u16, UInt16Array, UInt16Builder, |v: u16| v);
fill_backward_numeric!(fill_backward_u32, UInt32Array, UInt32Builder, |v: u32| v);
fill_backward_numeric!(fill_backward_u64, UInt64Array, UInt64Builder, |v: u64| v);
fill_backward_numeric!(fill_backward_f32, Float32Array, Float32Builder, |v: f32| v);
fill_backward_numeric!(fill_backward_f64, Float64Array, Float64Builder, |v: f64| v);

fn fill_forward_utf8(col: &Arc<dyn Array>) -> Result<StringArray> {
    let array = col
        .as_any()
        .downcast_ref::<StringArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad StringArray downcast"))?;
    let mut builder = StringBuilder::with_capacity(array.len(), array.value_data().len());
    let mut last: Option<String> = None;
    for i in 0..array.len() {
        if array.is_null(i) {
            match last.as_deref() {
                Some(v) => builder.append_value(v),
                None => builder.append_null(),
            }
        } else {
            let v = array.value(i).to_string();
            last = Some(v.clone());
            builder.append_value(v);
        }
    }
    Ok(builder.finish())
}

fn fill_backward_utf8(col: &Arc<dyn Array>) -> Result<StringArray> {
    let array = col
        .as_any()
        .downcast_ref::<StringArray>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad StringArray downcast"))?;
    let mut tmp = Vec::with_capacity(array.len());
    let mut next: Option<String> = None;
    for i in (0..array.len()).rev() {
        if array.is_null(i) {
            tmp.push(next.clone());
        } else {
            let v = array.value(i).to_string();
            next = Some(v.clone());
            tmp.push(Some(v));
        }
    }
    tmp.reverse();
    let mut builder = StringBuilder::with_capacity(array.len(), array.value_data().len());
    for v in tmp {
        match v {
            Some(v) => builder.append_value(v),
            None => builder.append_null(),
        }
    }
    Ok(builder.finish())
}

fn fill_stat_i8(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Int8Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int8Array downcast"))?;
    let (min, max, mean) = stats_signed(array.iter().flatten().map(|v| v as i128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_int8(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_i16(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Int16Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int16Array downcast"))?;
    let (min, max, mean) = stats_signed(array.iter().flatten().map(|v| v as i128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_int16(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_i32(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Int32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int32Array downcast"))?;
    let (min, max, mean) = stats_signed(array.iter().flatten().map(|v| v as i128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_int32(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_i64(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Int64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Int64Array downcast"))?;
    let (min, max, mean) = stats_signed(array.iter().flatten().map(|v| v as i128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_int64(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_u8(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<UInt8Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt8Array downcast"))?;
    let (min, max, mean) = stats_unsigned(array.iter().flatten().map(|v| v as u128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_uint8(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_u16(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<UInt16Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt16Array downcast"))?;
    let (min, max, mean) = stats_unsigned(array.iter().flatten().map(|v| v as u128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_uint16(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_u32(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<UInt32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt32Array downcast"))?;
    let (min, max, mean) = stats_unsigned(array.iter().flatten().map(|v| v as u128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_uint32(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_u64(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<UInt64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad UInt64Array downcast"))?;
    let (min, max, mean) = stats_unsigned(array.iter().flatten().map(|v| v as u128));
    let value = match stat {
        Stat::Min => min.map(|v| v as i64),
        Stat::Max => max.map(|v| v as i64),
        Stat::Mean => mean.map(|v| v as i64),
    };
    match value {
        Some(v) => Ok(Arc::new(fill_uint64(col, &Scalar::Int64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_f32(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Float32Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Float32Array downcast"))?;
    let (min, max, mean) = stats_float(array.iter().flatten().map(|v| v as f64));
    let value = match stat {
        Stat::Min => min,
        Stat::Max => max,
        Stat::Mean => mean,
    };
    match value {
        Some(v) => Ok(Arc::new(fill_float32(col, &Scalar::Float64(v))?)),
        None => Ok(col.clone()),
    }
}

fn fill_stat_f64(col: &Arc<dyn Array>, stat: Stat) -> Result<Arc<dyn Array>> {
    let array = col
        .as_any()
        .downcast_ref::<Float64Array>()
        .ok_or_else(|| DataFrameError::invalid_operation("bad Float64Array downcast"))?;
    let (min, max, mean) = stats_float(array.iter().flatten());
    let value = match stat {
        Stat::Min => min,
        Stat::Max => max,
        Stat::Mean => mean,
    };
    match value {
        Some(v) => Ok(Arc::new(fill_float64(col, &Scalar::Float64(v))?)),
        None => Ok(col.clone()),
    }
}

fn stats_signed<I>(values: I) -> (Option<i128>, Option<i128>, Option<i128>)
where
    I: Iterator<Item = i128>,
{
    let mut min: Option<i128> = None;
    let mut max: Option<i128> = None;
    let mut sum: i128 = 0;
    let mut count = 0;
    for v in values {
        min = Some(min.map_or(v, |m| m.min(v)));
        max = Some(max.map_or(v, |m| m.max(v)));
        sum += v;
        count += 1;
    }
    let mean = if count == 0 { None } else { Some(sum / count) };
    (min, max, mean)
}

fn stats_unsigned<I>(values: I) -> (Option<u128>, Option<u128>, Option<u128>)
where
    I: Iterator<Item = u128>,
{
    let mut min: Option<u128> = None;
    let mut max: Option<u128> = None;
    let mut sum: u128 = 0;
    let mut count = 0;
    for v in values {
        min = Some(min.map_or(v, |m| m.min(v)));
        max = Some(max.map_or(v, |m| m.max(v)));
        sum += v;
        count += 1;
    }
    let mean = if count == 0 { None } else { Some(sum / count) };
    (min, max, mean)
}

fn stats_float<I>(values: I) -> (Option<f64>, Option<f64>, Option<f64>)
where
    I: Iterator<Item = f64>,
{
    let mut min: Option<f64> = None;
    let mut max: Option<f64> = None;
    let mut sum = 0.0_f64;
    let mut count = 0;
    for v in values {
        min = Some(min.map_or(v, |m| m.min(v)));
        max = Some(max.map_or(v, |m| m.max(v)));
        sum += v;
        count += 1;
    }
    let mean = if count == 0 {
        None
    } else {
        Some(sum / count as f64)
    };
    (min, max, mean)
}