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)]
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, "---")
    }
}

impl ColumnFrame {
    pub fn new<K: Into<KeyIndex>>(index: K, data_frame: Array2<DataValue>) -> Self {
        Self {
            index: index.into(),
            data_frame,
        }
    }

    pub fn keys(&self) -> &[Key] {
        self.index.get_keys()
    }

    pub fn len(&self) -> usize {
        self.data_frame.nrows()
    }

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

    pub fn shrink(&mut self) {
        //fixme
    }

    /// 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> {
        let mut keys = self.index.keys.clone();
        for key in keys.iter_mut() {
            if !force && matches!(key.ctype, crate::DataType::Unknown) {
                let column = self
                    .get_single_column(key)
                    .ok_or_else(|| Error::EmptyData)?;
                let dtype = crate::detect_dtype(column.get(0).ok_or_else(|| Error::EmptyData)?);
                key.ctype = dtype;
            } else if force {
                let column = self
                    .get_single_column(key)
                    .ok_or_else(|| Error::EmptyData)?;
                let dtype = crate::detect_dtype(column.get(0).ok_or_else(|| Error::EmptyData)?);
                key.ctype = dtype;
            }
        }
        self.index.keys = keys;

        Ok(())
    }
    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))
        }
    }

    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| {
            let x = &item;
            match key.ctype {
                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::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 => item,
            }
        });
        Ok(())
    }

    pub fn get_mut_view(&mut self) -> ArrayViewMut2<'_, DataValue> {
        self.data_frame.view_mut()
    }

    pub fn rename_key(&mut self, old: &str, new: Key) -> Result<(), Error> {
        self.index.rename_key(old, new)
    }

    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 the column-major order
    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))
    }

    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
    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() {
            return Array2::default((0, 0));
        }
        let mut new_data_frame = Array2::default((self.data_frame.nrows(), keys.len()));

        for (idx, key) in keys.iter().enumerate() {
            if let Some(column_index_in_source) = indexes.get_column_index(key) {
                new_data_frame
                    .column_mut(idx)
                    .assign(&self.data_frame.column(column_index_in_source));
            }
        }

        new_data_frame
    }

    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`]
    ///
    pub fn push<C: CandidateData>(&mut self, row_candidate: C) -> Result<(), Error> {
        let mut arr = vec![];
        for key in row_candidate.keys() {
            if self.index.get_column_index(&key).is_none() {
                self.extend_dataframe_for_column(key)?;
            }
        }
        for index in self.index.get_keys() {
            if let Some(value) = row_candidate.get_value_ref(index) {
                arr.push(value.clone());
            } else {
                arr.push(DataValue::Null);
            }
        }
        self.data_frame.push_row(Array::from_vec(arr).view())?;
        Ok(())
    }

    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> {
        for key in other.index.get_keys() {
            if self.index.get_column_index(key).is_none() {
                self.extend_dataframe_for_column(key.clone())?;
            }
        }
        Ok(())
    }

    fn try_extend(&mut self, other: Self) -> Result<(), Error> {
        let mut joined_keys = self.index.clone();
        // extend keys
        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)
                .expect("BUG: index for this has to be defined");
            if let Some(index) = self.index.get_column_index(key) {
                for (idx, value) in self.data_frame.column(index).iter().enumerate() {
                    if let Some(x) = arr.get_mut((idx, index_result)) {
                        *x = value.to_owned();
                    }
                }
            }

            if let Some(index) = other.index.get_column_index(key) {
                for (idx, value) in other.data_frame.column(index).iter().enumerate() {
                    if let Some(x) = arr.get_mut((increment + idx, index_result)) {
                        *x = value.to_owned();
                    }
                }
            }
        }
        *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 new_columnns = 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)?;
        // get the indexes for the keys
        let index = Index::new(keys.to_vec(), self);
        tracing::trace!("Index {index:?}");
        let right_indexes = right.index.select(keys);

        let mut new_df = Array2::default((self.len(), self.index.len()));
        new_df.assign(&self.data_frame);
        trace!("Indexes: {right_indexes:?}");
        trace!("New DF: {new_df:?}");
        trace!("Right DF: {:?}", right.data_frame);
        trace!("current {:?}", self.data_frame);
        // iterate over all rows in the right data frame and find the corresponding row in the current data frame
        // then fill the new data frame with the values from the right data frame
        // if the row is not found, the row is skipped - the data is not filled - [`DataValue::Null`] is used
        for c_row in right.data_frame.rows() {
            let values = right_indexes
                .get_keys()
                .iter()
                .map(|k| {
                    c_row
                        .get(
                            right_indexes
                                .get_column_index(k)
                                .expect("BUG: Something is very wrong"),
                        )
                        .cloned()
                        .unwrap_or_default()
                })
                .collect::<Vec<_>>();
            if let Some(row_index) = index.get(&values) {
                trace!("Index: {row_index} {c_row:?} -> {values:?} vs {new_df:?}");
                for complement_key in new_columnns.iter() {
                    let column_index = self
                        .index
                        .get_column_index(complement_key)
                        .expect("BUG: Something is very wrong");
                    if let Some(right_column_index) = right.index.get_column_index(complement_key) {
                        trace!("Filling Index: [{complement_key:?}] ri:{row_index} rci:{right_column_index:?} -> {:?} vs {:?}", c_row[right_column_index], new_df.get_mut((row_index, column_index)));
                        if let Some(v) = new_df.get_mut((row_index, column_index)) {
                            *v = c_row[right_column_index].to_owned();
                        }
                    }
                }
            }
        }

        self.data_frame = new_df;

        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);
        }
        self.extend_columns_from_other(&other)?;
        let mut new_df = Array2::default((self.len(), self.index.len()));
        for (idx, key) in self.index.get_keys().iter().enumerate() {
            if let Some(other_idx) = other.index.get_column_index(key) {
                new_df
                    .column_mut(idx)
                    .assign(&other.data_frame.column(other_idx));
            } else {
                new_df.column_mut(idx).assign(&self.data_frame.column(idx));
            }
        }
        self.data_frame = new_df;
        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),
        }
    }

    pub fn get_single_column(&self, key: &Key) -> Option<ArrayView1<'_, DataValue>> {
        self.index
            .get_column_index(key)
            .map(|x| self.data_frame.column(x))
    }

    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)| {
            a_val
                .partial_cmp(b_val)
                .unwrap_or(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))
    }

    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)
    }
}

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_eq!(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() {
        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());
    }
}