trs_dataframe/dataframe/

column_store.rs

use ndarray::{concatenate, s, Array, Array1, Array2, ArrayView1, ArrayViewMut2, Axis};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

use crate::error::Error;
use crate::{dataframe::index::Index, CandidateData, JoinBy, JoinRelation, Key};
use data_value::{DataValue, Extract};
use tracing::*;
mod from;
mod key_index;
mod ops;
pub mod sorted_df;
pub use key_index::KeyIndex;
pub mod filter_df;

/// [`ColumnFrame`] is used to store the data for the candidates
/// The data is stored in the [`Array2`] with the [`DataValue`] values
/// The data is stored in the columns and the columns are indexed by the [`KeyIndex`]
/// The [`KeyIndex`] is used to access the data by the column [`Key`]
/// Memory layout is same like in [`ndarray`] - the data is stored in the row-major order
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct ColumnFrame {
    pub index: KeyIndex,
    pub data_frame: Array2<DataValue>,
}

enum Continue {
    Continue,
    End,
}

impl Continue {
    pub fn should_end(&self) -> bool {
        matches!(self, Self::End)
    }
}

use std::fmt;

impl fmt::Display for ColumnFrame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Display keys and and indices
        write!(f, "\n|")?;

        for key in &self.index.keys {
            write!(f, " {key} |")?;
        }

        if self.index.is_empty() {
            writeln!(f, "|")?;
        }

        // Display type for each key
        if let Some(row) = self.data_frame.axis_iter(Axis(0)).next() {
            write!(f, "\n|")?;
            for value in row.iter() {
                // Display types during first iteration
                write!(f, " {:10?} |", crate::detect_dtype(value))?;
            }
            writeln!(f)?;
        }

        writeln!(f, "---")?;

        // Display items, limit output to 256 rows
        for (n, row) in self.data_frame.axis_iter(Axis(0)).enumerate() {
            write!(f, "|")?;

            for value in row.iter() {
                // Display types during first iteration
                write!(f, " {value} |")?;
            }
            writeln!(f)?;

            if n >= 256 {
                writeln!(f, "... (dataframe is too long)")?;
                break;
            }
        }

        writeln!(f, "---")
    }
}
/// Converts a [`DataValue`] to the canonical representation for the given [`crate::DataType`].
///
/// Uses the [`Extract`] trait for numeric coercion. If `dtype` is [`crate::DataType::Unknown`]
/// and the value is not `Null`, the type is auto-detected and conversion retried.
pub fn convert_data_value(item: DataValue, dtype: crate::DataType) -> DataValue {
    let x = &item;
    match dtype {
        crate::DataType::Bool => DataValue::Bool(bool::extract(x)),
        crate::DataType::U32 => DataValue::U32(u32::extract(x)),
        crate::DataType::I32 => DataValue::I32(i32::extract(x)),
        crate::DataType::U64 => DataValue::U64(u64::extract(x)),
        crate::DataType::I64 => DataValue::I64(i64::extract(x)),
        crate::DataType::F32 => DataValue::F32(f32::extract(x)),
        crate::DataType::U128 => DataValue::U128(u128::extract(x)),
        crate::DataType::I128 => DataValue::I128(i128::extract(x)),
        crate::DataType::F64 => DataValue::F64(f64::extract(x)),
        crate::DataType::U8 => DataValue::U8(u8::extract(x)),
        crate::DataType::String => DataValue::String(String::extract(x).into()),
        crate::DataType::Bytes => item,
        crate::DataType::Map => item,
        crate::DataType::Vec => item,
        crate::DataType::Unknown => {
            if matches!(item, DataValue::Null) {
                return item;
            }
            let dtype = crate::detect_dtype(&item);
            // this situation should not ever happen
            if matches!(dtype, crate::DataType::Unknown) {
                tracing::error!("Unknown datatype {dtype:?} - {item:?}");
                return item;
            }
            convert_data_value(item, dtype)
        }
    }
}
/// Converts a [`DataValue`] to match the dtype declared by the column [`Key`].
///
/// Shorthand for `convert_data_value(item, key.ctype)`.
pub fn convert_dv_to_dtype(key: &Key, item: DataValue) -> DataValue {
    convert_data_value(item, key.ctype)
}
impl ColumnFrame {
    /// Creates a new [`ColumnFrame`] from a column index and a 2D data array.
    ///
    /// The number of columns in `data_frame` must match the number of keys in `index`.
    pub fn new<K: Into<KeyIndex>>(index: K, data_frame: Array2<DataValue>) -> Self {
        Self {
            index: index.into(),
            data_frame,
        }
    }

    /// Returns the ordered slice of column [`Key`]s in this frame.
    pub fn keys(&self) -> &[Key] {
        self.index.get_keys()
    }

    /// Returns the number of rows in this frame.
    pub fn len(&self) -> usize {
        self.data_frame.nrows()
    }

    /// Returns `true` if the frame contains no rows.
    pub fn is_empty(&self) -> bool {
        self.data_frame.nrows() == 0
    }

    /// Compacts internal storage to reduce memory usage after mutations
    /// that may have left over-allocated capacity.
    pub fn shrink(&mut self) {
        let shape = self.data_frame.shape();
        if shape[0] > 0 && shape[1] > 0 {
            let mut new_data = Vec::with_capacity(shape[0] * shape[1]);
            for elem in self.data_frame.iter() {
                new_data.push(elem.clone());
            }
            if let Ok(arr) = Array2::from_shape_vec((shape[0], shape[1]), new_data) {
                self.data_frame = arr;
            }
        }
    }

    /// This method will try to fix dtype based on data stored in each column. If dtype is [`crate::DataType::Unknown`]
    /// this method will replace dtype for "correct" one based on [`DataValue`].
    /// NOTE: flag `force` will enforce this dtype even if dtype is known
    pub fn try_fix_dtype_for_keys(&mut self, force: bool) -> Result<(), Error> {
        for i in 0..self.index.keys.len() {
            let should_fix = force || matches!(self.index.keys[i].ctype, crate::DataType::Unknown);

            if should_fix {
                let column = self
                    .get_single_column(&self.index.keys[i])
                    .ok_or_else(|| Error::EmptyData)?;
                let dtype = crate::detect_dtype(column.get(0).ok_or_else(|| Error::EmptyData)?);
                self.index.keys[i].ctype = dtype;
            }
        }

        Ok(())
    }
    /// Attempts to fix the dtype of every column by inspecting stored values.
    ///
    /// Columns whose data does not match their declared dtype are cast in-place.
    /// Returns `Err` with a list of `(key, message)` pairs for columns that failed casting.
    pub fn try_fix_dtype(&mut self) -> Result<(), Error> {
        let mut errors = vec![];
        let keys = self.index.keys.clone();
        for key in keys {
            tracing::trace!("key: {key:?}- {:?}", key.ctype);
            if let Err(e) = self.try_fix_column_by_key(&key) {
                errors.push((key, e.to_string()));
            }
        }
        if errors.is_empty() {
            Ok(())
        } else {
            Err(Error::CastFailed(errors))
        }
    }

    /// Casts all values in the column identified by `key` to match the key's declared dtype.
    ///
    /// Returns an error if the key is not found in the frame.
    pub fn try_fix_column_by_key(&mut self, key: &Key) -> Result<(), Error> {
        let idx = self
            .index
            .get_column_index(key)
            .ok_or(Error::MissingField(format!("{key}").into()))?;
        let mut col = self.data_frame.column_mut(idx);

        col.mapv_inplace(|item| convert_dv_to_dtype(key, item));
        Ok(())
    }

    /// Forces all values in the named column to the specified `dtype`, updating
    /// both the stored data and the column's [`Key`] metadata.
    ///
    /// Returns an error if the column name does not exist.
    pub fn enforce_dtype_for_column(
        &mut self,
        key: &str,
        dtype: crate::DataType,
    ) -> Result<(), Error> {
        if let Some(idx) = self.index.get_column_index_by_name(key) {
            let new_key = Key::new(key, dtype);
            let mut col = self.data_frame.column_mut(idx);

            col.mapv_inplace(|item| convert_dv_to_dtype(&new_key, item));
            self.index.rename_key(key, new_key)?;
            Ok(())
        } else {
            Err(Error::NotFound(Key::new(key, crate::DataType::Unknown)))
        }
    }

    /// Returns a mutable 2D view over the entire data array.
    ///
    /// This allows direct element-level mutation without going through
    /// [`get_mut_by_row_index`](Self::get_mut_by_row_index).
    pub fn get_mut_view(&mut self) -> ArrayViewMut2<'_, DataValue> {
        self.data_frame.view_mut()
    }

    /// Renames a column from `old` (matched by name) to `new`.
    ///
    /// The new [`Key`] can carry a different [`DataType`](crate::DataType), which
    /// effectively re-declares the column's type without casting existing values.
    pub fn rename_key(&mut self, old: &str, new: Key) -> Result<(), Error> {
        self.index.rename_key(old, new)
    }

    /// Registers `alias` as an alternative name for the column named `key`.
    ///
    /// After aliasing, `select` and friends accept both the original name and the alias.
    pub fn add_alias(&mut self, key: &str, alias: &str) -> Result<(), Error> {
        self.index.add_alias(key, alias)
    }

    /// Selects the data from the [`ColumnFrame`] by the given keys
    /// If the keys are not provided, the data is selected by the [`KeyIndex`] keys
    /// The data is returned as the [`Vec<Vec<DataValue>>`]
    /// If the keys are not found, the empty [`Vec<Vec<DataValue>>`] is returned
    /// Returns the data in the column-major order
    pub fn select_transposed_typed<D: Extract>(&self, keys: &[Key]) -> Vec<Vec<D>> {
        let selected = self.select(Some(keys));
        let mut result = Vec::with_capacity(selected.nrows());
        for row in selected.rows() {
            let mut r = Vec::with_capacity(selected.ncols());
            for value in row.iter() {
                r.push(D::extract(value));
            }
            result.push(r);
        }
        result
    }

    /// Selects the data from the [`ColumnFrame`] by the given keys
    /// If the keys are not provided, the data is selected by the [`KeyIndex`] keys
    /// The data is returned as the [`Array2`] with the [`DataValue`] values
    /// If the keys are not found, the empty [`Array2`] is returned
    /// Returns the [`Array2`] with the data in row-major order (same as ndarray default)
    /// Note: Despite the method name, the data is NOT transposed - it returns data in row-major order
    /// Use this when you want row-oriented access to the selected columns
    pub fn select_transposed(&self, keys: Option<&[Key]>) -> Result<Array2<DataValue>, Error> {
        let keys = keys.unwrap_or_else(|| self.index.get_keys());
        let key_indexes = self.index.select(keys);
        if key_indexes.is_empty() {
            return Ok(Array2::default((0, 0)));
        }
        let data_vec: Vec<Array1<DataValue>> = key_indexes
            .indexes()
            .iter()
            .map(|x| self.data_frame.column(*x).to_owned())
            .collect();
        to_array2(data_vec)
    }

    /// Selects whole column from the [`ColumnFrame`] by the given key
    /// If the key is not found, the None is returned
    /// If the key is found, the [`ArrayView1`] with the [`DataValue`] values is returned
    pub fn select_column(&self, key: &Key) -> Option<ArrayView1<'_, DataValue>> {
        self.index
            .get_column_index(key)
            .map(|x| self.data_frame.column(x))
    }

    /// Applies a user-defined closure to this frame.
    ///
    /// The closure receives the provided `keys` and a mutable reference to the
    /// frame, allowing arbitrary in-place transformations.
    pub fn apply_function<F>(&mut self, keys: &[Key], mut func: F) -> Result<(), Error>
    where
        F: FnMut(&[Key], &mut ColumnFrame) -> Result<(), Error>,
    {
        func(keys, self)
    }

    /// Validates the access to the entry by the given column [`Key`] and row index
    /// If the column is not found, an error is returned [`Error::NotFound`]
    /// If the row index is out of bounds, an error is returned [`Error::IndexOutOfRange`]
    /// Otherwise, the column index is returned
    pub fn validate_entry_access(&self, column: &Key, row_index: usize) -> Result<usize, Error> {
        if row_index >= self.data_frame.nrows() {
            return Err(Error::IndexOutOfRange(row_index, self.data_frame.nrows()));
        }
        let Some(column_index) = self.index.get_column_index(column) else {
            return Err(Error::NotFound(column.clone()));
        };
        Ok(column_index)
    }

    /// Returns the value [`DataValue`] for the given column defined by [`Key`] and row index
    /// If the column is not found, None is returned
    /// If the row index is out of bounds, None is returned
    pub fn get_by_row_index(&self, column: &Key, row_index: usize) -> Option<&DataValue> {
        trace!(
            "Column: {column} row_index: {row_index} data_frame: cols:{}-rows:{}",
            self.data_frame.len(),
            self.data_frame.nrows()
        );
        trace!("{:?}", self.data_frame);
        match self.validate_entry_access(column, row_index) {
            Ok(column_index) => self.data_frame.get((row_index, column_index)),
            Err(e) => {
                trace!("Error: {e}");
                None
            }
        }
    }

    /// Returns mutable reference for the value [`DataValue`] for the given column defined by [`Key`] and row index
    /// If the column is not found, None is returned
    /// If the row index is out of bounds, None is returned
    pub fn get_mut_by_row_index(
        &mut self,
        column: &Key,
        row_index: usize,
    ) -> Option<&mut DataValue> {
        trace!(
            "Column: {column} row_index: {row_index} data_frame: cols:{}-rows:{}",
            self.data_frame.len(),
            self.data_frame.nrows()
        );
        trace!("{:?}", self.data_frame);
        match self.validate_entry_access(column, row_index) {
            Ok(column_index) => self.data_frame.get_mut((row_index, column_index)),
            Err(e) => {
                trace!("Error: {e}");
                None
            }
        }
    }

    /// Returns the value [`HashMap<Key, Vec<DataValue>>`] for the given columns defined by [`Key`].
    /// If the keys are not provided, the data is selected by the [`KeyIndex`] keys
    /// If the keys are not found, the empty [`HashMap`] is returned
    /// Returns the [`Array2`] with the data in the row-major order
    pub fn select_as_map(&self, keys: Option<&[Key]>) -> HashMap<Key, Vec<DataValue>> {
        let keys = keys.unwrap_or_else(|| self.index.get_keys());
        let indexes = self.index.select(keys);
        if indexes.is_empty() {
            return Default::default();
        }

        let mut new_data_frame = HashMap::with_capacity(keys.len());

        for key in keys.iter() {
            if let Some(column_index_in_source) = indexes.get_column_index(key) {
                let column = self.data_frame.column(column_index_in_source);
                new_data_frame.insert(key.clone(), column.to_vec());
            }
        }

        new_data_frame
    }

    /// Returns the value [`Array2<DataValue>`] for the given columns defined by [`Key`].
    /// If the keys are not provided, the data is selected by the [`KeyIndex`] keys
    /// If the keys are not found, the empty [`Array2`] is returned
    /// Returns the [`Array2`] with the data in the row-major order
    ///
    /// This operation clones all selected data. For read-only access, consider using
    /// [`Self::select_column`] to get a view of a single column without copying.
    pub fn select(&self, keys: Option<&[Key]>) -> Array2<DataValue> {
        let keys = keys.unwrap_or_else(|| self.index.get_keys());
        let indexes = self.index.select(keys);
        if indexes.is_empty() || keys.is_empty() {
            return Array2::default((0, 0));
        }

        let nrows = self.data_frame.nrows();
        let ncols = keys.len();

        // Pre-allocate vector with exact capacity
        let mut data = Vec::with_capacity(nrows * ncols);

        // Collect valid column mappings
        let col_mappings: Vec<(usize, usize)> = keys
            .iter()
            .enumerate()
            .filter_map(|(dst_idx, key)| {
                indexes
                    .get_column_index(key)
                    .map(|src_idx| (dst_idx, src_idx))
            })
            .collect();

        // Fill with nulls first, then overwrite valid data
        data.resize(nrows * ncols, DataValue::Null);

        // Copy column by column - better for row-major layout
        for row_idx in 0..nrows {
            for (dst_col, src_col) in &col_mappings {
                let dst_idx = row_idx * ncols + dst_col;
                data[dst_idx] = self.data_frame[(row_idx, *src_col)].clone();
            }
        }

        Array2::from_shape_vec((nrows, ncols), data).unwrap_or_else(|_| Array2::default((0, 0)))
    }

    /// Returns selected columns as a typed 2D array, converting each [`DataValue`]
    /// via the [`Extract`] trait.
    ///
    /// This is the typed counterpart of [`select`](Self::select). If `keys` is `None`,
    /// all columns are returned. The data is in row-major order (rows × columns).
    ///
    /// # Type coercion
    ///
    /// The [`Extract`] trait performs best-effort numeric coercion (e.g. `I32 -> f64`).
    /// Values that cannot be meaningfully converted yield the type's default
    /// (0 for numbers, `false` for bool, empty string for `String`).
    pub fn select_typed<T: Extract>(&self, keys: Option<&[Key]>) -> Array2<T> {
        self.select(keys).mapv(|x| T::extract(&x))
    }

    fn extend_dataframe_for_column(&mut self, key: Key) -> Result<(), Error> {
        self.index.store_key(key);
        let len = self.data_frame.nrows();
        self.data_frame.push_column(Array1::default(len).view())?;
        Ok(())
    }

    /// Pushes the row candidate into the [`ColumnFrame`]
    /// If the column is not found this method will add the column to the [`ColumnFrame`]
    ///
    /// This operation clones all values from the candidate. For batch operations,
    /// consider using [`Self::extend`] instead which can be more efficient.
    pub fn push<C: CandidateData>(&mut self, row_candidate: C) -> Result<(), Error> {
        // Pre-allocate with capacity to avoid reallocations
        let num_keys = self.index.len();
        let candidate_keys = row_candidate.keys();
        let mut arr = Vec::with_capacity(num_keys.max(candidate_keys.len()));

        // First pass: extend dataframe with any new columns
        for key in &candidate_keys {
            if self.index.get_column_index(key).is_none() {
                self.extend_dataframe_for_column(key.clone())?;
            }
        }

        // Second pass: collect values - capacity may have changed after extension
        arr.reserve(self.index.len());
        for index in self.index.get_keys() {
            arr.push(
                row_candidate
                    .get_value_ref(index)
                    .cloned()
                    .unwrap_or(DataValue::Null),
            );
        }

        self.data_frame.push_row(Array::from_vec(arr).view())?;
        Ok(())
    }

    /// Removes the specified columns from this frame and returns them as a new [`ColumnFrame`].
    ///
    /// The remaining columns stay in `self` with their data intact.
    pub fn remove_column(&mut self, keys: &[Key]) -> Result<Self, Error> {
        // fixme this is naive approach
        let mut indexes = KeyIndex::default();
        // take remove data
        let data = self.select(Some(keys));
        // remove labels from the index
        for key in keys {
            if let Some((current, _idx)) = self.index.remove_key(key) {
                indexes.store_key(current);
            }
        }
        // copy the rest of the data to the new data frame with new index
        let rest = self.select(Some(self.keys()));
        let keys = self.index.get_keys().to_vec();
        self.data_frame = rest;
        self.index = KeyIndex::new(keys);

        //remove_self.data_frame = to_array2(columns)?;
        Ok(Self::new(indexes, data))
    }

    fn check_or_init_frame(&mut self, other: &Self) -> Result<Continue, Error> {
        if self.index.is_empty() {
            self.index = other.index.clone();
            self.data_frame = other.data_frame.clone();
            return Ok(Continue::End);
        }
        if other.index.is_empty() {
            return Ok(Continue::End);
        }
        if self.is_empty() {
            self.data_frame = Array2::default((other.data_frame.nrows(), self.index.len()));
        }

        Ok(Continue::Continue)
    }

    fn extend_columns_from_other(&mut self, other: &Self) -> Result<(), Error> {
        let missing_keys: Vec<Key> = other
            .index
            .get_keys()
            .iter()
            .filter(|key| self.index.get_column_index(key).is_none())
            .cloned()
            .collect();

        if missing_keys.is_empty() {
            return Ok(());
        }

        for key in missing_keys {
            self.index.store_key(key);
        }

        let nrows = self.data_frame.nrows();
        let new_cols = self.index.len() - self.data_frame.ncols();

        if new_cols > 0 {
            // Use ndarray's optimized concatenate - faster than manual copying
            let new_data = Array2::default((nrows, new_cols));
            self.data_frame = concatenate(Axis(1), &[self.data_frame.view(), new_data.view()])?;
        }

        Ok(())
    }

    fn try_extend(&mut self, other: Self) -> Result<(), Error> {
        let mut joined_keys = self.index.clone();
        for key in other.keys() {
            if self.index.get_column_index(key).is_none() {
                joined_keys.store_key(key.clone());
            }
        }

        let sum_len = self.data_frame.nrows() + other.data_frame.nrows();
        let mut arr = Array2::default((sum_len, joined_keys.len()));
        let increment = self.data_frame.nrows();

        for key in joined_keys.get_keys() {
            let index_result = joined_keys.get_column_index(key).ok_or_else(|| {
                Error::UnknownError(format!(
                    "Index lookup failed for key '{}' in try_extend",
                    key.name()
                ))
            })?;

            let mut col = arr.column_mut(index_result);

            if let Some(index) = self.index.get_column_index(key) {
                let src_col = self.data_frame.column(index);
                col.slice_mut(s![..increment]).assign(&src_col);
            }

            if let Some(index) = other.index.get_column_index(key) {
                let src_col = other.data_frame.column(index);
                col.slice_mut(s![increment..]).assign(&src_col);
            }
        }

        *self = ColumnFrame::new(joined_keys, arr);
        Ok(())
    }

    /// Extends the [`ColumnFrame`] with the data from the other [`ColumnFrame`]
    /// If the [`KeyIndex`] is empty, the [`ColumnFrame`] is replaced with the other [`ColumnFrame`]
    /// If the other [`KeyIndex`] is empty, nothing happens
    /// If the length of the [`KeyIndex`] of the other data frame is greater then current,
    /// an error is returned [`Error::DataSetSizeDoesntMatch`]
    /// If [`Key`] from other data frame - extends the [`KeyIndex`] and add column to the current [`ColumnFrame`]
    ///
    pub fn extend(&mut self, mut other: Self) -> Result<(), Error> {
        if self.check_or_init_frame(&other)?.should_end() {
            return Ok(());
        }

        if self.index.check_order_of_indexes(&other.index).is_err() {
            return self.try_extend(other);
        }

        trace!(
            "Extend columns from other {:?} vs {:?}",
            other.index.get_keys(),
            self.index.get_keys()
        );

        if other.data_frame.ncols() < self.data_frame.ncols() {
            other.extend_columns_from_other(self)?;
        } else {
            self.extend_columns_from_other(&other)?;
        }
        self.data_frame = concatenate(Axis(0), &[self.data_frame.view(), other.data_frame.view()])?;

        Ok(())
    }

    /// Replace the [`ColumnFrame`] with the other [`ColumnFrame`]
    /// If the current [`KeyIndex`] is empty, the [`ColumnFrame`] is replaced with the other [`ColumnFrame`]
    /// If the other [`KeyIndex`] is empty, nothing happens
    /// If the [`KeyIndex`] of the other data frame and current doesn't match an error is returned [`Error::DataSetSizeDoesntMatch`]
    /// If the [`Key`] from other data frame is not present in the current [`ColumnFrame`] - extends the [`KeyIndex`] and add column to the current [`ColumnFrame`]
    pub fn replace(&mut self, other: Self) -> Result<(), Error> {
        if self.check_or_init_frame(&other)?.should_end() {
            return Ok(());
        }

        if self.data_frame.len() > other.data_frame.len() {
            return Err(Error::DataSetSizeDoesntMatch(
                self.data_frame.len(),
                other.data_frame.len(),
            ));
        }
        self.index = other.index;
        self.data_frame = other.data_frame;

        Ok(())
    }

    /// Joins the candidates by the keys in the `JoinRelation::JoinById` struct.
    /// This function creates [`Index`] for the keys and then joins the candidates by the keys.
    ///
    pub fn join_by_id_inner(&mut self, right: Self, keys: &[Key]) -> Result<(), Error> {
        if self.check_or_init_frame(&right)?.should_end() {
            return Ok(());
        }

        let timer = std::time::Instant::now();
        let new_columns = right.index.get_complement_keys(self.index.get_keys());

        // add new columns and keys into the column frame index
        self.extend_columns_from_other(&right)?;
        tracing::debug!("Extend took {}ns", timer.elapsed().as_nanos());

        // Pre-compute column index mappings AFTER extension to get correct indices
        let column_mappings: Vec<(usize, usize)> = new_columns
            .iter()
            .filter_map(|key| {
                let left_idx = self.index.get_column_index(key)?;
                let right_idx = right.index.get_column_index(key)?;
                Some((left_idx, right_idx))
            })
            .collect();

        // get the indexes for the keys
        let timer = std::time::Instant::now();
        let index = Index::new(keys.to_vec(), self);
        tracing::debug!("Left index build took: {}ns", timer.elapsed().as_nanos());
        tracing::trace!("Index {index:?}");

        let timer = std::time::Instant::now();
        let right_index = Index::new(keys.to_vec(), &right);
        let joined_idx = index.join(right_index);
        tracing::debug!(
            "Right index build and join took: {}ns",
            timer.elapsed().as_nanos()
        );

        // Instead of copying all data, we only fill in the new columns
        // The existing columns are already correct from extend_columns_from_other
        let timer = std::time::Instant::now();
        let joined_idx_len = joined_idx.len();

        for (left_col_idx, right_col_idx) in &column_mappings {
            let mut left_col = self.data_frame.column_mut(*left_col_idx);
            let right_col = right.data_frame.column(*right_col_idx);

            for (left_indices, right_indices) in &joined_idx {
                for right_row_idx in right_indices {
                    for left_idx in left_indices {
                        left_col[*left_idx] = right_col[*right_row_idx].clone();
                    }
                }
            }
        }

        let elapsed = timer.elapsed();
        tracing::debug!(
            "Filled {} rows in {}ms|{}s",
            joined_idx_len,
            elapsed.as_millis(),
            elapsed.as_secs()
        );

        Ok(())
    }

    /// Adds the single column to the current [`ColumnFrame`]
    /// If the column is already present, an error is returned [`Error::ColumnAlreadyExists`]
    /// If the length of the column is different from the current data frame, an error is returned [`Error::DataSetSizeDoesntMatch`]
    pub fn add_single_column<K: Into<Key>>(
        &mut self,
        key: K,
        column: Array1<DataValue>,
    ) -> Result<(), Error> {
        let key = key.into();
        if self.index.get_column_index(&key).is_some() {
            return Err(Error::ColumnAlreadyExists(key));
        }
        if self.len() != column.len() && !self.is_empty() {
            return Err(Error::DataSetSizeDoesntMatch(self.len(), column.len()));
        }

        self.index.store_key(key.clone());
        let rows = column.len();
        let column_index = self
            .index
            .get_column_index(&key)
            .ok_or(Error::UnknownError(format!("Column {key} should exists")))?;
        if self.is_empty() && self.index.len() == 1 {
            self.data_frame = column.into_shape_clone((rows, 1))?;
            assert_eq!(self.data_frame.column(column_index).len(), rows);
        } else if self.is_empty() {
            self.data_frame = Array2::default((column.len(), self.index.len() - 1));
            self.data_frame.push_column(column.view())?;
            assert_eq!(self.data_frame.column(column_index).len(), rows);
        } else {
            self.data_frame.push_column(column.view())?;
        }
        assert_eq!(self.data_frame.column(column_index).len(), rows);

        Ok(())
    }
    /// Adds the columns from the other [`ColumnFrame`] to the current [`ColumnFrame`]
    /// If the current [`KeyIndex`] is empty, the [`ColumnFrame`] is replaced with the other [`ColumnFrame`]
    /// If the other [`KeyIndex`] is empty, nothing happens
    pub fn add_columns(&mut self, other: Self) -> Result<(), Error> {
        if self.check_or_init_frame(&other)?.should_end() {
            return Ok(());
        }

        self.extend_columns_from_other(&other)?;
        for (idx, key) in other.index.get_keys().iter().enumerate() {
            if let Some(index) = self.index.get_column_index(key) {
                trace!("Other array = {:?}", other.data_frame.dim());
                if other.data_frame.dim() == (0, 0) {
                    self.data_frame.column_mut(index).fill(DataValue::Null);
                    continue;
                }
                let arr = other.data_frame.column(idx);
                trace!(
                    "Adding column {key:?} at index {idx} vs {index} datasize: self:{} vs other:{}",
                    self.data_frame.nrows(),
                    arr.len()
                );
                if arr.len() != self.data_frame.nrows() {
                    self.data_frame.column_mut(index).fill(DataValue::Null);
                } else {
                    self.data_frame.column_mut(index).assign(&arr);
                }
            }
        }
        Ok(())
    }

    /// Broadcasts the data from the other [`ColumnFrame`] to the current [`ColumnFrame`]
    /// If the current [`KeyIndex`] is empty, the [`ColumnFrame`] is replaced with the other [`ColumnFrame`]
    /// If the other [`KeyIndex`] is empty, nothing happens
    /// If the length (number of rows) of the other data frame is greater then 1 an error is returned [`Error::CannotBroadcast`]
    pub fn broadcast(&mut self, other: Self) -> Result<(), Error> {
        if self.check_or_init_frame(&other)?.should_end() {
            return Ok(());
        }
        if other.data_frame.nrows() != 1 {
            return Err(Error::CannotBroadcast);
        }

        // Get all keys (current + new from other)
        let all_keys = self.index.get_keys().to_vec();
        let other_keys: Vec<_> = other
            .index
            .get_keys()
            .iter()
            .filter(|k| self.index.get_column_index(k).is_none())
            .cloned()
            .collect();

        // Add new keys to index
        for key in &other_keys {
            self.index.store_key(key.clone());
        }

        let nrows = self.len();
        let ncols = self.index.len();
        let ncols_old = all_keys.len();

        // Pre-allocate result buffer
        let mut data = Vec::with_capacity(nrows * ncols);

        // Build result row by row - single pass allocation
        for row_idx in 0..nrows {
            // Copy existing columns
            for col_idx in 0..ncols_old {
                data.push(self.data_frame[(row_idx, col_idx)].clone());
            }
            // Broadcast values from other dataframe (only 1 row in other)
            for key in &other_keys {
                if let Some(other_idx) = other.index.get_column_index(key) {
                    data.push(other.data_frame[(0, other_idx)].clone());
                }
            }
        }

        self.data_frame = Array2::from_shape_vec((nrows, ncols), data)
            .map_err(|e| Error::UnknownError(format!("Broadcast reshape failed: {e}")))?;

        Ok(())
    }

    /// Computes the Cartesian product of the input structures,
    /// resulting in all possible combinations of elements.
    /// The data is stored in the row-major order
    /// The keys are stored in the order they are added - the order is preserved - new keys from the `other` [`ColumnFrame`] are added to the end
    pub fn cartesian_product(&mut self, other: Self) -> Result<(), Error> {
        if self.check_or_init_frame(&other)?.should_end() {
            return Ok(());
        }
        // extend the columns
        // let mut new_index = self.index.clone();
        for other_key in other.keys() {
            if self.index.get_column_index(other_key).is_none() {
                self.index.store_key(other_key.clone());
            } else {
                self.index.store_key(Key::new(
                    format!("{}-{}", other_key, other_key.id()).as_str(),
                    other_key.ctype,
                ));
            }
        }
        let max_rows = self.len() * other.len();
        let ncols = self.index.len();
        // create new data frame
        let mut new_df = Array2::default((max_rows, ncols));

        let mut cur_idx = 0;
        for cur_row in self.data_frame.rows() {
            for other_row in other.data_frame.rows() {
                new_df
                    .slice_mut(s![cur_idx, ..])
                    .assign(&concatenate(Axis(0), &[cur_row, other_row])?);
                cur_idx += 1;
            }
        }
        self.data_frame = new_df;
        Ok(())
    }

    /// Joins the candidates with the other candidates by the [`JoinRelation`] policy.
    /// For [`JoinBy::AddColumns`] the columns are added to the existing structure via [`Self::add_columns`]
    /// For [`JoinBy::Replace`] the columns are replaced with the new columns
    /// For [`JoinBy::Extend`] the candidates are extended via [`Self::extend`]
    /// For [`JoinBy::Broadcast`] each candidate is extended with the values of the other candidates `Self::broadcast`
    /// For [`JoinBy::CartesianProduct`] the candidates are multiplied by the other candidates
    /// For [`JoinBy::JoinById`] the candidates are joined by the keys in the `JoinRelation::JoinById` struct see [`Self::join_by_id_inner`]
    pub fn join(&mut self, right: Self, join_type: &JoinRelation) -> Result<(), Error> {
        use JoinBy::*;
        match &join_type.join_type {
            AddColumns => self.add_columns(right),
            Replace => self.replace(right),
            Extend => self.extend(right),
            Broadcast => self.broadcast(right),
            CartesianProduct => self.cartesian_product(right),
            JoinById(join) => self.join_by_id_inner(right, &join.keys),
        }
    }

    /// Returns a read-only view of a single column, or `None` if the key is absent.
    ///
    /// This is a zero-copy operation — the returned [`ArrayView1`] borrows directly
    /// from the underlying storage.
    pub fn get_single_column(&self, key: &Key) -> Option<ArrayView1<'_, DataValue>> {
        self.index
            .get_column_index(key)
            .map(|x| self.data_frame.column(x))
    }

    /// Returns a column extracted into a typed [`Array1<T>`], where each [`DataValue`]
    /// is converted via the [`Extract`] trait.
    ///
    /// This is a convenience wrapper around [`get_single_column`](Self::get_single_column)
    /// that maps every element through `T::extract`, producing an owned array of the
    /// target type. Returns `None` if the key does not exist in the frame.
    ///
    /// # Type coercion
    ///
    /// The [`Extract`] trait performs best-effort numeric coercion (e.g. `I32 -> f64`).
    /// Values that cannot be meaningfully converted yield the type's default
    /// (0 for numbers, `false` for bool, empty string for `String`).
    pub fn get_single_column_typed<T: Extract>(&self, key: &Key) -> Option<Array1<T>> {
        self.get_single_column(key)
            .map(|x| x.mapv(|x| T::extract(&x)))
    }

    /// Returns a [`SortedDataFrame`](sorted_df::SortedDataFrame) view sorted ascending by the given column.
    ///
    /// `Null` values are pushed to the end. Ties are broken by original row order.
    /// Use [`SortedDataFrame::topn`](sorted_df::SortedDataFrame::topn) to efficiently
    /// extract the first/last N rows.
    pub fn sorted(&self, key: &Key) -> Result<sorted_df::SortedDataFrame<'_>, Error> {
        let index = self
            .index
            .get_column_index(key)
            .ok_or(Error::NotFound(key.clone()))?;
        let column = self.data_frame.column(index);
        let mut data_with_index = column.iter().enumerate().collect::<Vec<_>>();
        tracing::trace!("Sorting by key: {key:?} vals {data_with_index:?}");
        data_with_index.sort_by(
            |(a_idx, a_val), (b_idx, b_val)| match a_val.partial_cmp(b_val) {
                Some(ordering) => ordering.then_with(|| a_idx.cmp(b_idx)),
                None => {
                    let a_null = matches!(a_val, DataValue::Null);
                    let b_null = matches!(b_val, DataValue::Null);
                    match (a_null, b_null) {
                        (true, true) => std::cmp::Ordering::Equal.then_with(|| a_idx.cmp(b_idx)),
                        (true, false) => std::cmp::Ordering::Greater.then_with(|| a_idx.cmp(b_idx)),
                        (false, true) => std::cmp::Ordering::Less.then_with(|| a_idx.cmp(b_idx)),
                        (false, false) => std::cmp::Ordering::Equal.then_with(|| a_idx.cmp(b_idx)),
                    }
                }
            },
        );

        tracing::trace!("Sorted by key: {key:?} vals {data_with_index:?}");
        let indicies = data_with_index
            .into_iter()
            .map(|(idx, _)| idx)
            .collect::<Vec<_>>();

        Ok(sorted_df::SortedDataFrame::new(self, indicies))
    }

    /// Returns a new [`ColumnFrame`] containing only rows that match the filter expression.
    ///
    /// The filter is applied against each row and matching row indices are collected,
    /// de-duplicated, and used to build the result.
    pub fn filter(&self, filter: &crate::filter::FilterRules) -> Result<Self, Error> {
        let mut final_indices = Vec::new();
        let filter_df = filter_df::ColumnFrameFiltering { column_frame: self };
        for rule in &filter.rules {
            final_indices.extend(crate::filter::filter_combination(&filter_df, rule)?);
        }

        final_indices.sort_unstable();
        final_indices.dedup();
        let mut new_df = ColumnFrame::new(
            self.index.clone(),
            Array2::default((final_indices.len(), self.index.len())),
        );
        final_indices
            .iter()
            .enumerate()
            .for_each(|(cur_idx, row_idx)| {
                new_df
                    .data_frame
                    .slice_mut(s![cur_idx, ..])
                    .assign(&self.data_frame.slice(s![*row_idx, ..]));
            });

        Ok(new_df)
    }
}

/// Converts a `Vec` of 1D arrays into a 2D array.
///
/// Each [`Array1`] in `source` becomes one row of the resulting [`Array2`].
pub fn to_array2<T: Clone>(source: Vec<Array1<T>>) -> Result<Array2<T>, Error> {
    let width = source.len();
    let flattened: Array1<T> = source.into_iter().flat_map(|row| row.to_vec()).collect();
    let height = flattened.len() / width;
    Ok(flattened.into_shape_with_order((width, height))?)
}
#[macro_export]
macro_rules! df {
    ($($everything:tt)*) => {
        $crate::DataFrame::new($crate::column_frame!($($everything)*))
    };
}

#[macro_export]
macro_rules! column_frame {
    // case { "a" => 1, }
    ($($key:expr => $value:expr,)+) => { $crate::column_frame!($($key => $value),+) };
    // case { "a" => vec![1, 2, 3] }
    ($($key:expr => vec![$($value:expr),*]),*) => {
        $crate::column_frame!($($key => [$($value),*]),*)
    };
    // case { "a" => [1, 2, 3] }
    ($($key:expr => [$($value:expr),*]),*) => {
        {
           let data = ::ndarray::arr2(&[$(
                [$($value.into(),)*],
            )*]);

           let _keys = vec![$($key.into(),)*];

            $crate::ColumnFrame::new(
                $crate::KeyIndex::new(_keys),
                data.reversed_axes()
            )
        }
    };
    // case { "a" => 1 }
    ($($key:expr => $value:expr),*) => {
        {
           let _data = ::ndarray::arr2(&[[$($value.into(),)*]]);
           let _keys = vec![$($key.into(),)*];

            $crate::ColumnFrame::new(
                $crate::KeyIndex::new(_keys),
                _data,
            )
        }
    };
}

#[cfg(test)]
mod test {
    use crate::{filter::FilterRules, JoinById};

    use super::*;
    use data_value::stdhashmap;
    use ndarray::ArrayView;
    use rstest::*;
    use tracing_test::traced_test;

    #[rstest]
    #[case(
        column_frame! {
            "t" => [1751001987000000u64, 1752001987000000u64, 1753001987000000u64],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [1752001987000000u64],
            "b" => [5],
            "c" => [8]
        },
        FilterRules::try_from("t.to_datetime_us() == '2025-07-08 19:13:07'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => [1751001987000000f64, 1752001987000000f64, 1753001987000000f64],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [1752001987000000f64],
            "b" => [5],
            "c" => [8]
        },
        FilterRules::try_from("t.to_datetime_us() == '2025-07-08 19:13:07'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => [1751001987000000i64, 1752001987000000i64, 1753001987000000i64],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [1752001987000000i64],
            "b" => [5],
            "c" => [8]
        },
        FilterRules::try_from("t.to_datetime_us() == '2025-07-08 19:13:07'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => [1751001987000000u64, 1752001987000000u64, 1753001987000000u64],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [1751001987000000u64],
            "b" => [4],
            "c" => [7]
        },
        FilterRules::try_from("t.to_datetime_us() < '2025-07-08 19:13:07'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => ["2025-07-08 18:13:07", "2025-07-08 19:13:07", "2025-07-08 20:13:07"],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => ["2025-07-08 18:13:07"],
            "b" => [4],
            "c" => [7]
        },
        FilterRules::try_from("t.to_datetime_us() < '2025-07-08 19:13:07'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => ["2025-07-08 18:13:07", "2025-07-08 19:13:07", "2025-07-08 20:13:07"],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [],
            "b" => [],
            "c" => []
        },
        FilterRules::try_from("t.len() < 10u64").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => ["2025-07-08 18:13:07", "2025-07-08 19:13:07", "2025-07-08 20:13:07"],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => ["2025-07-08 18:13:07", "2025-07-08 19:13:07", "2025-07-08 20:13:07"],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        FilterRules::try_from("t.len() > 10u64").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => [DataValue::Vec(vec![1.into(), 2.into(), 3.into()]), DataValue::Vec(vec![]), DataValue::Vec(vec![1.into()])],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [DataValue::Vec(vec![])],
            "b" => [5],
            "c" => [ 8]
        },
        FilterRules::try_from("t.len() == 0u64").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "t" => [DataValue::Vec(vec![1.into(), 2.into(), 3.into()]), DataValue::Vec(vec![]), DataValue::Vec(vec![1.into()])],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "t" => [DataValue::Vec(vec![1.into()])],
            "b" => [6],
            "c" => [9]
        },
        FilterRules::try_from("t.len() == 1u64").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [1, 2],
            "b" => [4, 5],
            "c" => [7, 8]
        },
        FilterRules::try_from("a <= 2i32").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [2],
            "b" => [5],
            "c" => [8]
        },
        FilterRules::try_from("a <= 2i32 && c > 7i32").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [],
            "b" => [],
            "c" => []
        },
        FilterRules::try_from("a <= 2i32 && c > 9i32").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [1, 2],
            "b" => [4, 5],
            "c" => [7, 8]
        },
        FilterRules::try_from("a <= 2i32 || c > 9i32").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [2],
            "b" => [5],
            "c" => [8]
        },
        FilterRules::try_from("a <= 2i32 && (c > 9i32 || b == 5i32)").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => ["abcd", "ab", "abcdefg"],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => ["abcd","abcdefg"],
            "b" => [4, 6],
            "c" => [7, 9]
        },
        FilterRules::try_from("a ~= 'abcd.*'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [1],
            "b" => [4],
            "c" => [7]
        },
        FilterRules::try_from("a in [1u32, 1i32]'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [2, 3],
            "b" => [5, 6],
            "c" => [8, 9]
        },
        FilterRules::try_from("a notIn [1u32, 1i32]'").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1f64, 2f64, 3f64],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a" => [1f64, 2f64],
            "b" => [4, 5],
            "c" => [7, 8]
        },
        FilterRules::try_from("a < 3f64 || (a < 3f64 && b <= 5i32)").expect("BUG: cannot create filter rules"),
    )]
    #[case(
        column_frame! {
            "a" => [1f64, 2f64, 3f64],
            "b" => [4i64, 5i64, 6i64],
            "c" => [7i64, 8i64, 9i64]
        },
        column_frame! {
            "a" => [1f64, 2f64],
            "b" => [4i64, 5i64],
            "c" => [7i64, 8i64]
        },
        FilterRules::try_from("a >= 1f64 && (b <= 5 || c <= 8) && b >= 4").expect("BUG: cannot create filter rules"),
    )]
    #[traced_test]
    fn filter_test(
        #[case] df: ColumnFrame,
        #[case] expected: ColumnFrame,
        #[case] filter: FilterRules,
    ) {
        let filtered = df.filter(&filter).expect("BUG: cannot filter");
        assert_eq!(filtered, expected);
    }

    #[rstest]
    #[traced_test]
    fn test_macro() {
        let df = column_frame! {
            "a" => 1,
            "b" => 2,
            "c" => 3,
            "d" => 4,
        };

        assert_eq!(df.len(), 1);
        assert_eq!(df.keys(), &["a".into(), "b".into(), "c".into(), "d".into()]);
        let f = Array2::from_shape_vec((1, 4), vec![1.into(), 2.into(), 3.into(), 4.into()])
            .expect("BUG: cannot create array");
        assert_eq!(df.select(None), f);

        let df = column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        };

        assert_eq!(df.len(), 3);
        assert_eq!(df.keys(), &["a".into(), "b".into(), "c".into()]);
        let f = Array2::from_shape_vec(
            (3, 3),
            vec![
                1.into(),
                4.into(),
                7.into(),
                2.into(),
                5.into(),
                8.into(),
                3.into(),
                6.into(),
                9.into(),
            ],
        )
        .expect("BUG: cannot create array");
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(selected, f);

        let df1 = df! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        };

        // just to make sure we've tested Display
        let formatted = format!("{}", df);
        debug!("{}", formatted);

        assert_eq!(df1, crate::DataFrame::from(df));
    }

    #[rstest]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        column_frame! {
            "a_new" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        vec!["a_new", "b", "c"].into_iter().map(|x| x.into()).collect(),
        vec![("a", "a_new".into())]
    )]
    #[traced_test]
    fn rename_test(
        #[case] df: ColumnFrame,
        #[case] expected: ColumnFrame,
        #[case] keys: Vec<Key>,
        #[case] renames: Vec<(&str, Key)>,
    ) {
        let mut df = df;
        for (old, new) in renames {
            df.rename_key(old, new).expect("BUG: cannot rename key");
        }
        assert_eq!(df, expected);
        assert_eq!(df.keys(), keys.as_slice());
    }

    #[rstest]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::I32),
        column_frame!("a" => [1i32, 2i32, 3i32])
    )]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::U32),
        column_frame!("a" => [1u32, 2u32, 3u32])
    )]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::I64),
        column_frame!("a" => [1i64, 2i64, 3i64])
    )]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::U64),
        column_frame!("a" => [1u64, 2u64, 3u64])
    )]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::F64),
        column_frame!("a" => [1f64, 2f64, 3f64])
    )]
    #[case(
        column_frame!("a" => [1, 2, 3]),
        Key::new("a", crate::DataType::F32),
        column_frame!("a" => [1f32, 2f32, 3f32])
    )]
    // #[case(
    //     column_frame!("a" => [1, 2, 3]),
    //     Key::new("a", crate::DataType::String),
    //     column_frame!("a" => ["1", "2", "3"])
    // )]
    fn test_try_fix_dtype(
        #[case] mut df: ColumnFrame,
        #[case] key: Key,
        #[case] expected: ColumnFrame,
    ) {
        assert!(df.try_fix_column_by_key(&key).is_ok());
        assert_eq!(
            df.select(Some(&[key.clone()])),
            expected.select(Some(&[key.clone()]))
        );
    }

    #[fixture]
    fn unknown_df() -> ColumnFrame {
        let mut hm: HashMap<String, Vec<DataValue>> = HashMap::new();

        hm.insert("a".into(), vec![1u32.into()]);
        hm.insert("b".into(), vec![3i64.into()]);
        hm.insert("c".into(), vec![1f64.into()]);
        hm.insert("d".into(), vec![1u64.into()]);

        hm.into()
    }
    #[rstest]
    #[case(stdhashmap!(
        "a" => crate::DataType::U32,
        "b" => crate::DataType::I64,
        "c" => crate::DataType::F64,
        "d" => crate::DataType::U64)
    )]
    fn test_try_fix_dtype_unknown(
        mut unknown_df: ColumnFrame,
        #[case] dtypes: HashMap<String, crate::DataType>,
    ) {
        for dtype in dtypes.iter() {
            let t: &Key = unknown_df
                .keys()
                .iter()
                .find(|x| x.name() == dtype.0)
                .unwrap();
            assert_ne!(t.ctype, crate::DataType::Unknown);
        }
        assert!(unknown_df.try_fix_dtype_for_keys(false).is_ok());
        for dtype in dtypes.iter() {
            let t: &Key = unknown_df
                .keys()
                .iter()
                .find(|x| x.name() == dtype.0)
                .unwrap();
            assert_eq!(t.ctype, *dtype.1);
            assert!(unknown_df.try_fix_dtype_for_keys(false).is_ok());
        }
        assert!(unknown_df.try_fix_dtype_for_keys(true).is_ok());
    }

    #[rstest]
    #[case(
        column_frame!(Key::new("a", crate::DataType::F32) => [1, 2, 3]),
        Key::new("a", crate::DataType::F32),
        column_frame!("a" => [1f32, 2f32, 3f32])
    )]
    #[traced_test]
    fn test_try_fix(#[case] mut df: ColumnFrame, #[case] key: Key, #[case] expected: ColumnFrame) {
        assert!(df.try_fix_dtype().is_ok());
        assert_eq!(
            df.select(Some(&[key.clone()])),
            expected.select(Some(&[key]))
        )
    }

    #[rstest]
    #[traced_test]
    fn test_not_key_fix() {
        let mut cf = column_frame!("a" => [1]);
        let non_existing = Key::new("b", crate::DataType::I32);
        assert!(cf.try_fix_column_by_key(&non_existing).is_err());
    }

    #[rstest]
    #[case(
        column_frame! {
            "a" => [1, 2, 3],
            "b" => [4, 5, 6],
            "c" => [7, 8, 9]
        },
        vec!["a_alias", "b", "c"].into_iter().map(|x| x.into()).collect(),
        vec![("a", "a_alias")]
    )]
    #[traced_test]
    fn alias_test(
        #[case] df: ColumnFrame,
        #[case] keys: Vec<Key>,
        #[case] aliases: Vec<(&str, &str)>,
    ) {
        let mut df = df;
        for (old, new) in aliases {
            df.add_alias(old, new).expect("BUG: cannot rename key");
        }
        let origin_keys = df.keys().to_vec();
        let selected_aliases = df.select(Some(keys.as_slice()));
        let selected = df.select(Some(origin_keys.as_slice()));
        assert_eq!(selected, selected_aliases);
    }

    #[rstest]
    #[traced_test]
    fn test_mut_view() {
        let data = vec![
            DataValue::from(1f64),
            DataValue::from(4f32),
            DataValue::from(2f64),
            DataValue::from(f32::NAN),
            DataValue::from(f64::NAN),
            DataValue::from(f32::INFINITY),
        ];
        let keys: Vec<Key> = vec!["a".into(), "b".into()];

        let index = KeyIndex::new(keys.clone());
        let df = Array2::from_shape_vec((3, keys.len()), data).expect("BUG: cannot create array");
        let mut df = ColumnFrame::new(index.clone(), df);
        df.get_mut_view().mapv_inplace(|x| match x {
            DataValue::F32(f) if f.is_infinite() || f.is_nan() => DataValue::F32(0f32),
            DataValue::F64(f) if f.is_infinite() || f.is_nan() => DataValue::F64(0f64),
            e => e,
        });
        let data = vec![
            DataValue::from(1f64),
            DataValue::from(4f32),
            DataValue::from(2f64),
            DataValue::from(0f32),
            DataValue::from(0f64),
            DataValue::from(0f32),
        ];
        let expected = ColumnFrame::new(
            index,
            Array2::from_shape_vec((3, keys.len()), data).expect("BUG: cannot create ndarray"),
        );
        assert_eq!(df, expected);
    }

    #[rstest]
    #[traced_test]
    fn dummy_test() {
        let data = vec![
            DataValue::U32(1),
            DataValue::I32(2),
            DataValue::I64(3),
            DataValue::U64(4),
        ];

        let keys: Vec<Key> = vec!["a".into(), "b".into(), "c".into(), "d".into()];

        let index = KeyIndex::new(keys.clone());
        let mut data_frame = Array2::default((1, keys.len()));
        for (idx, entry) in data.iter().enumerate() {
            data_frame
                .column_mut(idx)
                .assign(&ArrayView::from(&[entry.clone()]));
        }

        let frame = ColumnFrame::new(index, data_frame);
        assert_eq!(
            frame.get_by_row_index(&"a".into(), 0),
            Some(&DataValue::U32(1))
        );
        assert_eq!(frame.get_by_row_index(&"aa".into(), 0), None);
        assert_eq!(frame.get_by_row_index(&"a".into(), 1), None);
        assert_eq!(
            frame.select(Some(&["a".into(), "b".into()])),
            Array2::from_shape_vec((1, 2), vec![DataValue::U32(1), DataValue::I32(2)])
                .expect("BUG: cannot create array")
        );
    }

    #[rstest]
    #[traced_test]
    fn dummy_test_multiple_rows() {
        let data = vec![
            DataValue::U32(1),
            DataValue::I32(2),
            DataValue::I64(3),
            DataValue::U64(4),
            DataValue::U32(12),
            DataValue::I32(22),
            DataValue::I64(32),
            DataValue::U64(42),
        ];

        let keys: Vec<Key> = vec!["a".into(), "b".into(), "c".into(), "d".into()];

        let index = KeyIndex::new(keys.clone());
        let data_frame =
            Array2::from_shape_vec((2, keys.len()), data).expect("BUG: cannot create array");

        let frame = ColumnFrame::new(index, data_frame);
        assert_eq!(
            frame.get_by_row_index(&"a".into(), 0),
            Some(&DataValue::U32(1))
        );
        assert_eq!(frame.get_by_row_index(&"aa".into(), 0), None);
        assert_eq!(frame.get_by_row_index(&"a".into(), 3), None);
        let arr = Array2::from_shape_vec(
            (2, 2),
            vec![
                DataValue::U32(1),
                DataValue::I32(2),
                DataValue::U32(12),
                DataValue::I32(22),
            ],
        )
        .expect("BUG: cannot create array");
        trace!("{arr:?}");
        assert_eq!(frame.select(Some(&["a".into(), "b".into()])), arr);
    }

    #[rstest]
    #[traced_test]
    fn dummy_test_multiple_rows_push() {
        let data = vec![
            DataValue::U32(1),
            DataValue::I32(2),
            DataValue::I64(3),
            DataValue::U64(4),
            DataValue::U32(12),
            DataValue::I32(22),
            DataValue::I64(32),
            DataValue::U64(42),
        ];
        let keys: Vec<Key> = vec!["a".into(), "b".into(), "c".into(), "d".into()];

        let index = KeyIndex::new(keys.clone());
        let data_frame =
            Array2::from_shape_vec((2, keys.len()), data).expect("BUG: cannot create array");

        let mut frame = ColumnFrame::new(index, data_frame);
        assert!(frame
            .push(data_value::stdhashmap!(
                "a" => DataValue::U32(2),
                "b" => DataValue::I32(3),
                "c" => DataValue::I64(4),
                "d" => DataValue::U64(5)
            ))
            .is_ok());
        let arr = Array2::from_shape_vec(
            (3, 2),
            vec![
                DataValue::U32(1),
                DataValue::I32(2),
                DataValue::U32(12),
                DataValue::I32(22),
                DataValue::U32(2),
                DataValue::I32(3),
            ],
        )
        .expect("BUG: cannot create array");
        trace!("{arr:?}");
        assert_eq!(frame.select(Some(&["a".into(), "b".into()])), arr);
        let result = frame.push(data_value::stdhashmap!(
            "a" => DataValue::U32(34),
            "b" => DataValue::I32(44),
            "c" => DataValue::I64(54),
            "e" => DataValue::F32(6f32)
        ));
        assert!(result.is_ok(), "{result:?}");
        let arr = Array2::from_shape_vec(
            (4, 2),
            vec![
                DataValue::U64(4),
                DataValue::Null,
                DataValue::U64(42),
                DataValue::Null,
                DataValue::U64(5),
                DataValue::Null,
                DataValue::Null,
                DataValue::F32(6f32),
            ],
        )
        .expect("BUG: cannot create array");
        trace!("{arr:?}");
        assert_eq!(frame.select(Some(&["d".into(), "e".into()])), arr);
    }

    #[rstest]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![3, 4]
        },
        Some(vec![Key::from("group_id")]),
        ndarray::array!([1.into()], [2.into()])
    )]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![3, 4]
        },
        Some(vec!["group_id".into(), "feed_tag".into()]),
        ndarray::array!([1.into(), 3.into()], [2.into(), 4.into()])
    )]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![3, DataValue::Null]
        },
        Some(vec!["feed_tag".into()]),
        ndarray::array![[3.into()], [DataValue::Null]]
    )]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![1, DataValue::Null]
        },
        Some(vec!["feed_tag2".into()]),
       Array2::<DataValue>::default((0, 0))
    )]
    #[traced_test]
    fn test_select(
        #[case] input: ColumnFrame,
        #[case] keys: Option<Vec<Key>>,
        #[case] expected: Array2<DataValue>,
    ) {
        trace!("input={input:?}");
        let keys_slice = keys.as_deref();
        let selected = input.select(keys_slice);
        trace!("selected={selected:?}");
        assert_eq!(selected, expected);
        let selected = input.select_transposed(keys_slice);
        trace!("selected_transposed={selected:?}");
        assert!(selected.is_ok());
        assert_eq!(selected.unwrap(), expected.t());
    }

    #[rstest]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![3, 4]
        },
        Key::from("group_id"),
        Some(ndarray::array!(1.into(), 2.into()))
    )]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2, 5, 6],
            "feed_tag" => vec![3, 4, 7, 8]
        },
        Key::from("group_id"),
        Some(ndarray::array!(1.into(), 2.into(), 5.into(), 6.into()))
    )]
    #[case(
        column_frame! {
            "group_id" => vec![1, 2],
            "feed_tag" => vec![1, 1]
        },
        Key::from("feed_tag1"),
        None
    )]
    #[traced_test]
    fn test_select_column(
        #[case] input: ColumnFrame,
        #[case] key: Key,
        #[case] expected: Option<Array1<DataValue>>,
    ) {
        let selected = input.select_column(&key);
        trace!("selected={selected:?}");
        match expected {
            Some(expected) => {
                assert!(selected.is_some());
                assert_eq!(selected.expect("BUG: checked above"), expected);
            }
            None => assert!(selected.is_none()),
        }
    }

    #[test]
    #[traced_test]
    fn empty_join_test() {
        let join = JoinRelation::add_columns();
        let mut column_frame = ColumnFrame::default();
        column_frame
            .add_single_column("group_id", Array1::from_vec(vec![]))
            .expect("BUG: cannot add column");
        let column_frame2 = column_frame! {
            "group_id" => vec![2, 1, 3],
            "feed_tag" => vec![1, 1, 1],
            "clicks" => vec![100, 10, 10],
            "imps" => vec![1000, 200, 200]
        };
        assert!(column_frame.join(ColumnFrame::default(), &join).is_ok());

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&[
                "group_id".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into()
            ])),
            ndarray::array!(
                [2.into(), 1.into(), 100.into(), 1000.into()],
                [1.into(), 1.into(), 10.into(), 200.into()],
                [3.into(), 1.into(), 10.into(), 200.into()],
            )
        );

        let mut column_frame2 = column_frame! {
            "feed_tag" => vec![1, 1, 1],
            "clicks" => vec![100, 10, 10],
            "imps" => vec![1000, 200, 200]
        };
        let mut column_frame = ColumnFrame::default();
        column_frame
            .add_single_column("group_id", Array1::from_vec(vec![]))
            .expect("BUG: cannot add column");
        let joined = column_frame2.join(column_frame, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame2:?}");
        assert_eq!(
            column_frame2.select(Some(&[
                "group_id".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into()
            ])),
            ndarray::array!(
                [DataValue::Null, 1.into(), 100.into(), 1000.into()],
                [DataValue::Null, 1.into(), 10.into(), 200.into()],
                [DataValue::Null, 1.into(), 10.into(), 200.into()],
            )
        );

        let mut column_frame = ColumnFrame::default();
        column_frame.index = KeyIndex::new(vec!["group_id2".into()]);
        let joined = column_frame2.join(column_frame, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame2:?}");
        assert_eq!(
            column_frame2.select(Some(&[
                "group_id2".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into()
            ])),
            ndarray::array!(
                [DataValue::Null, 1.into(), 100.into(), 1000.into()],
                [DataValue::Null, 1.into(), 10.into(), 200.into()],
                [DataValue::Null, 1.into(), 10.into(), 200.into()],
            )
        );
    }

    #[test]
    #[traced_test]
    fn join_test_multiple() {
        let join = JoinRelation::new(JoinBy::JoinById(JoinById::new(vec!["group_id".into()])));
        let mut column_frame = column_frame! {
            "group_id" => vec![1, 1, 3]
        };
        let column_frame2 = column_frame! {
            "group_id" => vec![2, 1, 1],
            "clicks" => vec![100, 10, 10],
            "imps" => vec![1000, 200, 200]
        };

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&["group_id".into(), "clicks".into(), "imps".into(),])),
            ndarray::array!(
                [1.into(), 10.into(), 200.into()],
                [1.into(), 10.into(), 200.into()],
                [3.into(), DataValue::Null, DataValue::Null],
            )
        )
    }

    #[test]
    #[traced_test]
    fn join_test_no_matches() {
        let join = JoinRelation::new(JoinBy::JoinById(JoinById::new(vec!["group_id".into()])));
        let mut column_frame = column_frame! {
            "group_id" => vec![DataValue::I32(1), DataValue::I32(2), DataValue::I32(3)]
        };
        let column_frame2 = column_frame! {
            "group_id" => vec![DataValue::I32(4), DataValue::I32(5), DataValue::I32(6)],
            "clicks" => vec![DataValue::I32(100), DataValue::I32(200), DataValue::I32(300)],
        };

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&["group_id".into(), "clicks".into()])),
            ndarray::array!(
                [DataValue::I32(1), DataValue::Null],
                [DataValue::I32(2), DataValue::Null],
                [DataValue::I32(3), DataValue::Null],
            )
        )
    }
    #[test]
    #[traced_test]
    fn join_test() {
        let join = JoinRelation::new(JoinBy::JoinById(JoinById::new(vec![
            "group_id".into(),
            "feed_tag".into(),
        ])));
        let mut column_frame = column_frame! {
            "group_id" => vec![1, 2, 8],
            "feed_tag" => vec![1, 1, 10]
        };
        let column_frame2 = column_frame! {
            "group_id" => vec![2, 1, 3],
            "feed_tag" => vec![1, 1, 1],
            "clicks" => vec![100, 10, 10],
            "imps" => vec![1000, 200, 200]
        };
        assert!(column_frame.join(ColumnFrame::default(), &join).is_ok());

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&[
                "group_id".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into()
            ])),
            ndarray::array!(
                [1.into(), 1.into(), 10.into(), 200.into()],
                [2.into(), 1.into(), 100.into(), 1000.into()],
                [8.into(), 10.into(), DataValue::Null, DataValue::Null]
            )
        )
    }

    #[test]
    #[traced_test]
    fn join_test_with_additional() {
        let join = JoinRelation::new(JoinBy::JoinById(JoinById::new(vec![
            "group_id".into(),
            "feed_tag".into(),
        ])));
        let mut column_frame = column_frame! {
            "group_id" => vec![1, 2, 8],
            "feed_tag" => vec![1, 1, 10],
            "clicked" => vec![0, 0, 1]
        };
        let column_frame2 = column_frame! {
            "group_id" => vec![2, 1, 3],
            "feed_tag" => vec![1, 1, 1],
            "clicks" => vec![100, 10, 10],
            "imps" => vec![1000, 200, 200]
        };
        assert!(column_frame.join(ColumnFrame::default(), &join).is_ok());

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&[
                "group_id".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into(),
                "clicked".into()
            ])),
            ndarray::array!(
                [1.into(), 1.into(), 10.into(), 200.into(), 0.into()],
                [2.into(), 1.into(), 100.into(), 1000.into(), 0.into()],
                [
                    8.into(),
                    10.into(),
                    DataValue::Null,
                    DataValue::Null,
                    1.into()
                ]
            )
        )
    }

    #[test]
    #[traced_test]
    fn join_test_with_additional_single() {
        let join = JoinRelation::new(JoinBy::JoinById(JoinById::new(vec![
            "group_id".into(),
            "feed_tag".into(),
        ])));
        let mut column_frame = column_frame! {
            "group_id" => vec![1, 2, 8],
            "feed_tag" => vec![1, 1, 10],
            "clicked" => vec![0, 0, 1]
        };
        let column_frame2 = column_frame! {
            "a" => vec![1],
            "group_id" => vec![2],
            "feed_tag" => vec![1],
            "clicks" => vec![10],
            "imps" => vec![200]
        };
        assert!(column_frame.join(ColumnFrame::default(), &join).is_ok());

        let joined = column_frame.join(column_frame2, &join);
        assert!(joined.is_ok(), "{joined:?}");

        trace!("{column_frame:?}");
        assert_eq!(
            column_frame.select(Some(&[
                "group_id".into(),
                "feed_tag".into(),
                "clicks".into(),
                "imps".into(),
                "clicked".into()
            ])),
            ndarray::array!(
                [
                    1.into(),
                    1.into(),
                    DataValue::Null,
                    DataValue::Null,
                    0.into(),
                ],
                [2.into(), 1.into(), 10.into(), 200.into(), 0.into()],
                [
                    8.into(),
                    10.into(),
                    DataValue::Null,
                    DataValue::Null,
                    1.into()
                ]
            )
        )
    }

    #[rstest]
    #[traced_test]
    fn cartesian_product_join() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "zone_id" => vec![111111, 111133],
            "zone_avg_ctr" => vec![0.1, 0.001]
        };
        assert!(df
            .join(
                ColumnFrame::default(),
                &JoinRelation::new(JoinBy::CartesianProduct)
            )
            .is_ok());
        let join = JoinRelation::new(JoinBy::CartesianProduct);
        let result = df.join(df2, &join);
        assert!(result.is_ok(), "{result:?}");
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [1.into(), 1.into(), 111111.into(), 0.1.into()],
                [1.into(), 1.into(), 111133.into(), 0.001.into()],
                [2.into(), 2.into(), 111111.into(), 0.1.into()],
                [2.into(), 2.into(), 111133.into(), 0.001.into()],
                [3.into(), 3.into(), 111111.into(), 0.1.into()],
                [3.into(), 3.into(), 111133.into(), 0.001.into()],
            )
        );

        let df2 = column_frame! {
            "zone_id" => vec![111]
        };
        let result = df.join(df2, &join);
        assert!(result.is_ok(), "{result:?}");
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [1.into(), 1.into(), 111111.into(), 0.1.into(), 111.into()],
                [1.into(), 1.into(), 111133.into(), 0.001.into(), 111.into()],
                [2.into(), 2.into(), 111111.into(), 0.1.into(), 111.into()],
                [2.into(), 2.into(), 111133.into(), 0.001.into(), 111.into()],
                [3.into(), 3.into(), 111111.into(), 0.1.into(), 111.into()],
                [3.into(), 3.into(), 111133.into(), 0.001.into(), 111.into()],
            )
        );
    }

    #[rstest]
    #[traced_test]
    fn broadcast_join() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "zone_id" => vec![111111]
        };
        assert!(df
            .join(
                ColumnFrame::default(),
                &JoinRelation::new(JoinBy::Broadcast)
            )
            .is_ok());
        let join = JoinRelation::new(JoinBy::Broadcast);
        assert!(df.join(df2, &join).is_ok());
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [1.into(), 1.into(), 111111.into()],
                [2.into(), 2.into(), 111111.into()],
                [3.into(), 3.into(), 111111.into()]
            )
        );
    }
    #[rstest]
    #[traced_test]
    fn merge_test() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "group_id" => vec![11, 21, 31],
            "feed_tag" => vec![12, 22, 32]
        };

        let join = JoinRelation::new(JoinBy::Replace);
        assert!(df.join(df2, &join).is_ok());
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [11.into(), 12.into()],
                [21.into(), 22.into()],
                [31.into(), 32.into()]
            )
        );
    }

    #[rstest]
    #[traced_test]
    fn extend_test() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "group_id" => vec![11, 21, 31],
            "feed_tag" => vec![5, 6, 7]
        };
        assert!(df
            .join(ColumnFrame::default(), &JoinRelation::new(JoinBy::Extend))
            .is_ok());

        let join = JoinRelation::new(JoinBy::Extend);
        assert!(df.join(df2, &join).is_ok());
        let selected = df.select(Some(&["feed_tag".into(), "group_id".into()]));
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [1.into(), 1.into()],
                [2.into(), 2.into()],
                [3.into(), 3.into()],
                [5.into(), 11.into()],
                [6.into(), 21.into()],
                [7.into(), 31.into()]
            )
        );
        let as_map = df.select_as_map(Some(&["feed_tag".into(), "group_id".into()]));
        trace!("{as_map:?}");
        assert_eq!(
            as_map,
            stdhashmap!(
                "feed_tag" => vec![1, 2, 3, 5, 6, 7],
                "group_id" => vec![1, 2, 3, 11, 21, 31]
            )
        );

        let as_map = df.select_as_map(Some(&["feed_tag1".into()]));
        trace!("{as_map:?}");
        assert_eq!(as_map, HashMap::default());
    }

    #[rstest]
    #[traced_test]
    fn extend_test_with_non_existing_cols() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let mut df2 = column_frame! {
            "group_id" => vec![11, 21, 31],
            "feed_tag" => vec![5, 6, 7],
            "clicks" => vec![100, 200, 300],
            "impressions" => vec![1000, 2000, 3000]
        };
        let df_bckp = df.clone();
        let join = JoinRelation::new(JoinBy::Extend);
        assert!(df.join(df2.clone(), &join).is_ok());
        let selected = df.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [1.into(), 1.into(), DataValue::Null, DataValue::Null],
                [2.into(), 2.into(), DataValue::Null, DataValue::Null],
                [3.into(), 3.into(), DataValue::Null, DataValue::Null],
                [11.into(), 5.into(), 100.into(), 1000.into()],
                [21.into(), 6.into(), 200.into(), 2000.into()],
                [31.into(), 7.into(), 300.into(), 3000.into()]
            )
        );
        let join = JoinRelation::new(JoinBy::Extend);
        let r = df2.join(df_bckp, &join);
        assert!(r.is_ok(), "{r:?}");
        let selected = df2.select(None);
        trace!("{selected:?}");
        assert_eq!(
            selected,
            ndarray::array!(
                [11.into(), 5.into(), 100.into(), 1000.into()],
                [21.into(), 6.into(), 200.into(), 2000.into()],
                [31.into(), 7.into(), 300.into(), 3000.into()],
                [1.into(), 1.into(), DataValue::Null, DataValue::Null],
                [2.into(), 2.into(), DataValue::Null, DataValue::Null],
                [3.into(), 3.into(), DataValue::Null, DataValue::Null]
            )
        );
    }

    #[rstest]
    #[traced_test]
    fn extend_test_with_non_existing_cols_wrong_order() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "feed_tag" => vec![5, 6, 7],
            "group_id" => vec![11, 21, 31]
        };
        let join = JoinRelation::new(JoinBy::Extend);
        let err = df.join(df2, &join);
        assert!(err.is_ok(), "{err:?}");
    }

    #[rstest]
    #[traced_test]
    fn test_replace_not_compatible() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "feed_tag" => vec![5, 6],
            "group_id" => vec![11, 21]
        };
        let join = JoinRelation::new(JoinBy::Replace);
        let err = df.join(df2, &join);
        assert!(err.is_err(), "{err:?}");
        let empty = ColumnFrame::default();
        let err = df.join(empty, &join);
        assert!(err.is_ok(), "{err:?}");
    }

    #[rstest]
    #[traced_test]
    fn test_different_data() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let df2 = column_frame! {
            "group_id" => vec![11, 21],
            "a" => vec![5, 6]
        };
        let join = JoinRelation::new(JoinBy::Extend);
        let err = df.join(df2, &join);
        assert!(err.is_ok(), "{err:?}");
        println!("{df:?}");
        let expected_df = ColumnFrame::new(
            KeyIndex::from(vec!["group_id".into(), "feed_tag".into(), "a".into()]),
            ndarray::array!(
                [1.into(), 1.into(), DataValue::Null],
                [2.into(), 2.into(), DataValue::Null],
                [3.into(), 3.into(), DataValue::Null],
                [11.into(), DataValue::Null, 5.into()],
                [21.into(), DataValue::Null, 6.into()]
            ),
        );
        assert_eq!(df, expected_df)
    }

    #[rstest]
    #[traced_test]
    fn serde_column_frame() {
        let df = column_frame! {
            "group_id" => vec![1u64, 2u64, 3u64],
            "feed_tag" => vec![1u64, 2u64, 3u64]
        };
        let key_idx = df.index.clone();
        let serialized = serde_json::to_string(&key_idx).expect("BUG: cannot serialize");
        let deserialized: KeyIndex =
            serde_json::from_str(&serialized).expect("BUG: cannot deserialize");
        assert_eq!(key_idx, deserialized);
        assert!(key_idx.get_key(0).is_some_and(|x| x == "group_id".into()));
        let serialized = serde_json::to_string(&df).expect("BUG: cannot serialize");
        let deserialized: ColumnFrame =
            serde_json::from_str(&serialized).expect("BUG: cannot deserialize");
        assert_eq!(df, deserialized);
    }

    #[rstest]
    #[traced_test]
    fn update_value() {
        let mut df = column_frame! {
            "group_id" => vec![1, 2, 3],
            "feed_tag" => vec![1, 2, 3]
        };
        let group_id: Key = "group_id".into();
        let v = df.get_mut_by_row_index(&group_id, 1);
        assert!(v.is_some());
        let v = v.unwrap();
        assert_eq!(v, &DataValue::I32(2));
        *v = DataValue::U64(22);
        let v = df.get_by_row_index(&group_id, 1);
        assert!(v.is_some());
        let v = v.unwrap();
        assert_eq!(v, &DataValue::U64(22));

        assert!(df.get_mut_by_row_index(&"group_id2".into(), 1).is_none());
    }

    #[rstest]
    fn get_single_column_typed_f64() {
        let df = column_frame! {
            "a" => [1i32, 2i32, 3i32],
            "b" => [10u64, 20u64, 30u64]
        };
        let key: Key = "a".into();
        let col = df.get_single_column_typed::<f64>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[1.0, 2.0, 3.0]));
    }

    #[rstest]
    fn get_single_column_typed_i64() {
        let df = column_frame! {
            "x" => [10u32, 20u32, 30u32]
        };
        let key: Key = "x".into();
        let col = df.get_single_column_typed::<i64>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[10i64, 20i64, 30i64]));
    }

    #[rstest]
    fn get_single_column_typed_string() {
        let df = column_frame! {
            "name" => ["alice", "bob", "carol"]
        };
        let key: Key = "name".into();
        let col = df.get_single_column_typed::<String>(&key).unwrap();
        assert_eq!(
            col,
            ndarray::arr1(&["alice".to_string(), "bob".to_string(), "carol".to_string()])
        );
    }

    #[rstest]
    fn get_single_column_typed_bool() {
        let df = column_frame! {
            "flag" => [1i32, 0i32, 1i32]
        };
        let key: Key = "flag".into();
        let col = df.get_single_column_typed::<bool>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[true, false, true]));
    }

    #[rstest]
    fn get_single_column_typed_missing_key_returns_none() {
        let df = column_frame! {
            "a" => [1, 2, 3]
        };
        let missing: Key = "nonexistent".into();
        assert!(df.get_single_column_typed::<f64>(&missing).is_none());
    }

    #[rstest]
    fn get_single_column_typed_numeric_coercion_from_mixed() {
        let df = column_frame! {
            "vals" => [1.5f64, 2.7f64, 3.9f64]
        };
        let key: Key = "vals".into();
        // f64 -> i32 truncates via `as`
        let col = df.get_single_column_typed::<i32>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[1i32, 2i32, 3i32]));
    }

    #[rstest]
    fn get_single_column_typed_selects_correct_column() {
        let df = column_frame! {
            "x" => [1, 2, 3],
            "y" => [10, 20, 30],
            "z" => [100, 200, 300]
        };
        let key: Key = "y".into();
        let col = df.get_single_column_typed::<i64>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[10i64, 20i64, 30i64]));
    }

    #[rstest]
    fn get_single_column_typed_u64_identity() {
        let df = column_frame! {
            "id" => [100u64, 200u64, 300u64]
        };
        let key: Key = "id".into();
        let col = df.get_single_column_typed::<u64>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[100u64, 200u64, 300u64]));
    }

    #[rstest]
    fn get_single_column_typed_single_row() {
        let df = column_frame! {
            "x" => [42i32]
        };
        let key: Key = "x".into();
        let col = df.get_single_column_typed::<f64>(&key).unwrap();
        assert_eq!(col, ndarray::arr1(&[42.0f64]));
    }

    #[rstest]
    fn get_single_column_typed_empty_frame() {
        let df = ColumnFrame::default();
        let key: Key = "x".into();
        let col = df.get_single_column_typed::<f64>(&key);
        assert!(col.is_none());
    }

    #[rstest]
    fn select_typed_all_columns() {
        let df = column_frame! {
            "a" => [1i32, 2i32],
            "b" => [3i32, 4i32]
        };
        let result = df.select_typed::<f64>(None);
        assert_eq!(result.nrows(), 2);
        assert_eq!(result.ncols(), 2);
        assert_eq!(result[[0, 0]], 1.0);
        assert_eq!(result[[0, 1]], 3.0);
        assert_eq!(result[[1, 0]], 2.0);
        assert_eq!(result[[1, 1]], 4.0);
    }

    #[rstest]
    fn select_typed_subset_of_columns() {
        let df = column_frame! {
            "a" => [10u64, 20u64],
            "b" => [30u64, 40u64],
            "c" => [50u64, 60u64]
        };
        let keys: Vec<Key> = vec!["a".into(), "c".into()];
        let result = df.select_typed::<i64>(Some(&keys));
        assert_eq!(result.nrows(), 2);
        assert_eq!(result.ncols(), 2);
        assert_eq!(result[[0, 0]], 10i64);
        assert_eq!(result[[0, 1]], 50i64);
        assert_eq!(result[[1, 0]], 20i64);
        assert_eq!(result[[1, 1]], 60i64);
    }

    #[rstest]
    fn select_typed_nonexistent_keys_returns_empty() {
        let df = column_frame! {
            "a" => [1i32, 2i32]
        };
        let keys: Vec<Key> = vec!["z".into()];
        let result = df.select_typed::<f64>(Some(&keys));
        assert_eq!(result.shape(), &[0, 0]);
    }

    #[rstest]
    fn select_typed_string_extraction() {
        let df = column_frame! {
            "name" => ["hello", "world"]
        };
        let result = df.select_typed::<String>(None);
        assert_eq!(result[[0, 0]], "hello");
        assert_eq!(result[[1, 0]], "world");
    }

    #[rstest]
    fn select_typed_matches_manual_mapv() {
        let df = column_frame! {
            "x" => [1i32, 2i32, 3i32],
            "y" => [4i32, 5i32, 6i32]
        };
        let typed = df.select_typed::<f64>(None);
        let manual = df.select(None).mapv(|v| f64::extract(&v));
        assert_eq!(typed, manual);
    }
}