trs_dataframe/dataframe/

python.rs

use std::collections::HashMap;

use crate::{DataFrame, DataValue, JoinRelation, Key};
use data_value::Extract as _;
use ndarray::Array1;
use numpy::{IntoPyArray, PyArray2};
use pyo3::{
    exceptions::PyTypeError,
    prelude::*,
    types::{PyBytes, PyList},
    IntoPyObjectExt,
};
use tracing::trace;

impl DataFrame {
    fn select_data(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
    ) -> Result<ndarray::Array2<DataValue>, crate::error::Error> {
        let keys = keys
            .map(|x| x.into_iter().map(Key::from).collect::<Vec<Key>>())
            .unwrap_or(self.keys());
        if transposed.unwrap_or(false) {
            self.select(Some(keys.as_slice()))
        } else {
            self.select_transposed(Some(keys.as_slice()))
        }
    }
}

enum DfOrDict {
    DataFrame(DataFrame),
    Dict(HashMap<String, DataValue>),
}

impl DfOrDict {
    pub fn new(object: Bound<'_, PyAny>) -> Result<DfOrDict, PyErr> {
        if let Ok(df) = object.extract::<DataFrame>() {
            Ok(DfOrDict::DataFrame(df))
        } else {
            let dict: HashMap<String, DataValue> = object.extract()?;
            Ok(DfOrDict::Dict(dict))
        }
    }
}

#[pymethods]
impl DataFrame {
    /// Create a new empty DataFrame.
    #[new]
    pub fn init() -> Self {
        Self::default()
    }

    /// Create a DataFrame from a polars dataframe in python.
    /// ```text
    /// import polars as pl
    /// df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    /// tei_df = tdf.DataFrame.from_polars(df)
    /// ```
    #[cfg(feature = "polars-df")]
    #[staticmethod]
    pub fn from_polars(df: pyo3_polars::PyDataFrame) -> Self {
        df.0.into()
    }

    /// Create a DataFrame from a dictionary.
    /// ```text
    /// df = tdf.DataFrame.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]})
    /// ```
    #[staticmethod]
    pub fn from_dict(df: HashMap<String, Vec<DataValue>>) -> Self {
        let mut result_df: Vec<(Key, Vec<DataValue>)> = Vec::new();
        for (key, value) in df.into_iter() {
            result_df.push((key.as_str().into(), value));
        }
        result_df.into()
    }

    /// Returns the keys of the DataFrame.
    pub fn keys(&self) -> Vec<Key> {
        self.dataframe.keys().to_vec()
    }

    /// Convert the DataFrame to polars DataFrame.
    /// This requires the `polars-df` feature to be enabled.
    /// ```text
    /// import polars as pl
    /// original_df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
    /// df = tdf.DataFrame.from_polars(original_df)
    /// polars_df = df.as_polars()
    /// assert polars_df.frame_equal(original_df)
    /// ```  
    #[cfg(feature = "polars-df")]
    #[pyo3(name = "as_polars")]
    pub fn py_as_polars(&self) -> PyResult<pyo3_polars::PyDataFrame> {
        let df = self
            .as_polars()
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot prepare polars DF: {e}")))?;
        Ok(pyo3_polars::PyDataFrame(df))
    }

    /// Apply a function to the DataFrame.
    /// The function should accept a DataFrame and return a DataFrame.
    /// ```text
    /// def my_function(df):
    ///    # Perform some operations on the DataFrame
    ///   return df
    /// /// df = tdf.DataFrame.init()
    /// df.apply(my_function)
    /// ```
    pub fn apply(&mut self, function: Bound<'_, PyAny>) -> Result<(), PyErr> {
        let df: DataFrame = pyo3::Python::attach(|py| {
            let self_ = self
                .clone()
                .into_pyobject(py)
                .expect("BUG: cannot convert to PyObject");
            let result = function.call1((self_,)).expect("BUG: cannot call function");
            result
                .extract::<Bound<DataFrame>>()
                .expect("BUG: cannot extract data frame")
                .unbind()
                .extract(py)
                .expect("BUG: cannot extract data frame")
        });
        self.dataframe = df.dataframe;
        Ok(())
    }

    /// Returns slice from dataframe as numpy.array of uint32 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_u32(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.uint32))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_u32<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<u32>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| u32::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of uint64 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_u64(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.uint64))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_u64<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<u64>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| u64::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of int32 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_i32(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.int32))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_i32<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<i32>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| i32::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of int64 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_i64(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.int64))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_i64<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<i64>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| i64::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of float32 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_f32(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.float32))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_f32<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<f32>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| f32::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of float64 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// a_np = df.as_numpy_f64(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([[1, 11], [2, 21]], dtype=np.float64))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_f64<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<f64>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        Ok(PyArray2::from_array(py, &data.mapv(|x| f64::extract(&x))))
    }

    /// Returns slice from dataframe as numpy.array of objects of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": "a", "key2": "b"})
    /// df.push({"key1": "d", "key2": "c"})
    /// a_np = df.as_numpy_f64(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([["a", d""], ["b", "c"]], dtype=np.object))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy_str<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<Py<PyAny>>>> {
        self.as_numpy(keys, transposed, py)
    }

    /// Returns slice from dataframe as numpy.array of float64 of the given keys.
    /// If `transposed` is true, the keys will be transposed.
    /// If `keys` is None, all keys will be used.
    /// ```text
    /// import numpy as np
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": "a", "key2": 1})
    /// df.push({"key1": "d", "key2": 2})
    /// a_np = df.as_numpy_f64(['key1', 'key2'])
    /// assert np.array_equal(a_np, np.array([["a", d""], [1, 2]], dtype=np.object))
    /// ```
    #[pyo3(signature = (keys=None, transposed=None))]
    pub fn as_numpy<'py>(
        &self,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, numpy::PyArray2<Py<PyAny>>>> {
        let data = self
            .select_data(keys, transposed)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?;
        let data = data.mapv(|x| {
            String::extract(&x)
                .into_py_any(py)
                .expect("cannot convert string to py object")
        });
        Ok(data.into_pyarray(py))
    }

    #[pyo3(name = "shrink")]
    pub fn py_shrink(&mut self) {
        self.dataframe.shrink();
    }

    #[pyo3(name = "add_metadata")]
    pub fn py_add_metadata(&mut self, key: String, value: DataValue) {
        self.metadata.insert(key, value);
    }

    #[pyo3(name = "get_metadata")]
    pub fn py_get_metadata(&self, key: &str) -> Option<DataValue> {
        self.metadata.get(key).cloned()
    }

    #[pyo3(name = "rename_key")]
    pub fn py_rename_key(&mut self, key: &str, new_name: &str) -> Result<(), PyErr> {
        // fixme this may have a problem when the type is different and checked
        self.dataframe
            .rename_key(key, new_name.into())
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("{e}")))
    }

    #[pyo3(name = "add_alias")]
    pub fn py_add_alias(&mut self, key: &str, new_name: &str) -> Result<(), PyErr> {
        self.dataframe
            .add_alias(key, new_name)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("{e}")))
    }

    /// Selects data from the DataFrame.
    /// If `keys` is None, all keys will be used.
    /// If `keys` is provided, only the specified keys will be selected.
    /// Returns a list of lists, where each inner list represents a row of data.
    /// ```text
    /// import trs_dataframe as tdf
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// # selected = df.select(["key1", "key2"])
    /// # assert selected == [[1, 2], [11, 21]]
    /// # selected = df.select()
    #[pyo3(name = "select", signature = (keys=None, transposed=None))]
    pub fn py_select<'py>(
        &self,
        py: Python<'py>,
        keys: Option<Vec<String>>,
        transposed: Option<bool>,
    ) -> Result<Bound<'py, PyList>, PyErr> {
        let keys = keys
            .map(|x| x.into_iter().map(Key::from).collect::<Vec<Key>>())
            .unwrap_or(self.keys());

        let selected = if transposed.unwrap_or_default() {
            self.select_transposed(Some(keys.as_slice()))
                .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?
        } else {
            self.select(Some(keys.as_slice()))
                .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot select data: {e}")))?
        };

        let list = PyList::empty(py);
        for rows in selected.rows() {
            let row = PyList::empty(py);
            for value in rows.iter() {
                row.append(value.clone())
                    .expect("BUG: cannot append to list");
            }
            list.append(row).expect("BUG: cannot append to list");
        }
        Ok(list)
    }

    /// Selects a column from the DataFrame.
    /// If the column does not exist, it will raise a TypeError.
    /// Returns a list of values in the selected column.
    /// ```text
    /// import trs_dataframe as tdf
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// # selected = df.select_column("key1")
    /// # assert selected == [1, 11]
    /// # selected = df.select_column("key2")
    /// # assert selected == [2, 21]
    /// # selected = df.select_column("non_existing_key")  # Raises TypeError
    /// ```
    #[pyo3(name = "select_column")]
    pub fn py_select_column<'py>(
        &self,
        py: Python<'py>,
        key: String,
    ) -> Result<Bound<'py, PyList>, PyErr> {
        let selected = self
            .select_column(Key::from(key))
            .ok_or_else(|| PyErr::new::<PyTypeError, _>("Cannot select column"))?;

        let list = PyList::empty(py);
        for x in selected.to_vec().into_iter() {
            list.append(x)?;
        }

        Ok(list)
    }

    /// Joins the current DataFrame with another DataFrame.
    /// The join type is specified by the `join_type` parameter.
    /// see [`JoinRelation`] for available join types.
    /// ```text
    /// import trs_dataframe as tdf
    /// df1 = tdf.DataFrame.init()
    /// df1.push({"key1": 1, "key2": 2})
    /// df1.push({"key1": 11, "key2": 21})
    /// df2 = tdf.DataFrame.init()
    /// df2.push({"key1": 1, "key2": 3})
    /// df2.push({"key1": 11, "key2": 23})
    /// df1.join(df2, tei.JoinRelation.extend())
    /// assert df1.select(["key1", "key2"]) == [[1, 2], [11, 21], [1, 3], [11, 23]]
    /// ```
    #[pyo3(name = "join")]
    pub fn py_join(&mut self, other: DataFrame, join_type: JoinRelation) -> Result<(), PyErr> {
        self.dataframe
            .join(other.dataframe, &join_type)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot join data: {e}")))?;

        Ok(())
    }

    /// Pushes a new row of data into the DataFrame.
    /// The data should be provided as a dictionary where keys are column names and values are the corresponding data values.
    /// ```text
    /// import trs_dataframe as tdf
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// ```
    #[pyo3(name = "push")]
    pub fn py_push(&mut self, data: HashMap<Key, DataValue>) -> Result<(), PyErr> {
        self.dataframe
            .push(data)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot join data: {e}")))?;
        Ok(())
    }

    /// Adds a new column to the DataFrame.
    /// The column is specified by a key and a vector of data values.
    /// If the length of the data vector does not match the number of rows in the DataFrame, it will raise a TypeError.
    /// ```text
    /// import trs_dataframe as tdf
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// df.add_column("key3", [3, 4])
    /// assert df.select(["key1", "key2", "key3"]) == [[1, 2, 3], [11, 21, 4]]
    /// ```
    #[pyo3(name = "add_column")]
    pub fn py_add_column(&mut self, key: Key, data: Vec<DataValue>) -> Result<(), PyErr> {
        self.dataframe
            .add_single_column(key, Array1::from_vec(data))
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot join data: {e}")))?;
        Ok(())
    }

    pub fn add_constant(&mut self, key: Key, feature: DataValue) -> Result<(), PyErr> {
        self.constants.insert(key, feature);
        Ok(())
    }

    /// Filters the DataFrame by a given expression.
    /// The expression should be a string that can be parsed by the DataFrame's filter method
    ///
    /// ```text
    /// import trs_dataframe as tdf
    /// df = tdf.DataFrame.init()
    /// df.push({"key1": 1, "key2": 2})
    /// df.push({"key1": 11, "key2": 21})
    /// df.filter_by_expression("key1 > 5")
    /// assert df.select(["key1", "key2"]) == [[11, 21 ]]
    /// ```
    pub fn filter_by_expression(&mut self, expression: String) -> Result<Self, PyErr> {
        let filter = crate::filter::FilterRules::try_from(expression.as_str())
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot parse expression: {e}")))?;
        self.filter(&filter)
            .map_err(|e| PyErr::new::<PyTypeError, _>(format!("Cannot filter data: {e}")))
    }

    fn __repr__(&self) -> String {
        self.to_string()
    }

    fn __str__(&self) -> String {
        self.to_string()
    }

    pub fn __iadd__(&mut self, object: Bound<'_, PyAny>) -> Result<(), PyErr> {
        trace!("{object:?}");
        let df_or_dict = DfOrDict::new(object)?;
        match df_or_dict {
            DfOrDict::DataFrame(df) => {
                self.dataframe += df.dataframe;
            }
            DfOrDict::Dict(dict) => {
                self.dataframe += dict;
            }
        }
        Ok(())
    }

    pub fn __isub__(&mut self, object: Bound<'_, PyAny>) -> Result<(), PyErr> {
        trace!("{object:?}");

        let df_or_dict = DfOrDict::new(object)?;
        match df_or_dict {
            DfOrDict::DataFrame(df) => {
                self.dataframe -= df.dataframe;
            }
            DfOrDict::Dict(dict) => {
                self.dataframe -= dict;
            }
        }
        Ok(())
    }

    pub fn __imul__(&mut self, object: Bound<'_, PyAny>) -> Result<(), PyErr> {
        trace!("{object:?}");
        let df_or_dict = DfOrDict::new(object)?;
        match df_or_dict {
            DfOrDict::DataFrame(df) => {
                self.dataframe *= df.dataframe;
            }
            DfOrDict::Dict(dict) => {
                self.dataframe *= dict;
            }
        }
        Ok(())
    }

    pub fn __itruediv__(&mut self, object: Bound<'_, PyAny>) -> Result<(), PyErr> {
        trace!("{object:?}");
        let df_or_dict = DfOrDict::new(object)?;
        match df_or_dict {
            DfOrDict::DataFrame(df) => {
                self.dataframe /= df.dataframe;
            }
            DfOrDict::Dict(dict) => {
                self.dataframe /= dict;
            }
        }
        Ok(())
    }

    pub fn __len__(&mut self) -> Result<usize, PyErr> {
        Ok(self.dataframe.len())
    }

    pub fn serialize_to_json_string(&self) -> String {
        serde_json::to_string(self).expect("Cannot serialize to strinng")
    }

    #[staticmethod]
    pub fn deserialize_from_json_string(json_df: String) -> Self {
        let mut df: DataFrame =
            serde_json::from_str(json_df.as_str()).expect("Cannot deserialize from str");
        let _ = df.dataframe.try_fix_dtype();

        df
    }

    // derive Serialize and Deserialize
    pub fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
        let s: DataFrame = rmp_serde::decode::from_slice(state.as_bytes()).map_err(|e| {
            pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                "Cannot deserialize object {e}"
            ))
        })?;
        *self = s;
        self.dataframe.try_fix_dtype().map_err(|e| {
            pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                "Cannot deserialize object {e}"
            ))
        })?;
        Ok(())
    }
    pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let buf = rmp_serde::encode::to_vec(self).map_err(|e| {
            // let buf = serde_json::to_string(self).map_err(|e| {
            pyo3::PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(format!(
                "Cannot deserialize object {e}"
            ))
        })?;
        Ok(PyBytes::new(py, &buf))
    }

    pub fn __del__(&mut self) {
        self.dataframe = Default::default();
    }

    // pub fn __reduce__(&self, py: Python<'_>) -> PyResult<(Py<PyAny>, Py<PyAny>)> {
    //     let cls = py.get_type::<Self>();
    //     Ok((
    //         cls.into(),
    //         pyo3::types::PyTuple::new(py, &[self.__getstate__(py)?])?
    //             .into_any()
    //             .unbind(),
    //     ))
    // }
}

#[cfg(test)]
mod test {

    use super::*;
    use crate::DataType;
    use data_value::{stdhashmap, DataValue};
    use halfbrown::hashmap;
    use pyo3::ffi::c_str;
    use rstest::*;
    use tracing_test::traced_test;

    #[fixture]
    fn df() -> DataFrame {
        let mut df = DataFrame::init();
        assert!(df
            .push(hashmap! {
                Key::new("key1", DataType::U32) => DataValue::U32(1),
                Key::new("key2", DataType::U32) => DataValue::U32(2),
            })
            .is_ok());
        assert!(df
            .push(hashmap! {
                Key::from("key1") => DataValue::U32(11),
                Key::from("key2") => DataValue::U32(21),
            })
            .is_ok());
        df
    }

    #[fixture]
    fn hm() -> HashMap<String, DataValue> {
        stdhashmap!(
            "key1".to_string() => DataValue::U32(2),
            "key2".to_string() => DataValue::U32(3),
        )
    }

    #[rstest]
    fn serde_py(df: DataFrame) {
        let str_df = df.serialize_to_json_string();
        assert!(!str_df.is_empty());

        let loaded = DataFrame::deserialize_from_json_string(str_df);

        assert_eq!(loaded, df);
    }
    #[cfg(feature = "python")]
    #[rstest]
    fn pickle_py(df: DataFrame) {
        pyo3::Python::attach(|py| {
            let bytes = df.__getstate__(py);
            assert!(bytes.is_ok());

            let mut deser = DataFrame::default();
            assert!(deser.__setstate__(bytes.unwrap().into()).is_ok());
            assert_eq!(deser, df);
        });
    }
    #[rstest]
    fn test_select_data(df: DataFrame) {
        let data = df.select_data(Some(vec!["key1".into(), "key2".into()]), Some(false));
        assert!(data.is_ok());
        assert_eq!(
            data.unwrap(),
            ndarray::array![[1u32.into(), 11u32.into()], [2u32.into(), 21u32.into()]]
        );

        let data = df.select_data(Some(vec!["key1".into(), "key2".into()]), Some(true));
        assert!(data.is_ok());
        assert_eq!(
            data.unwrap(),
            ndarray::array![[1u32.into(), 2u32.into()], [11u32.into(), 21u32.into()]]
        );
    }

    #[cfg(feature = "python")]
    #[rstest]
    fn test_from_create() {
        pyo3::Python::attach(|_py| {
            let mut hm: HashMap<String, Vec<DataValue>> = Default::default();
            let value: Vec<DataValue> = vec![1i32.into(), 22i32.into()];
            hm.insert("a".into(), value);

            let df = DataFrame::from_dict(hm);
            assert_eq!(
                df.select(Some(&["a".into()])),
                Ok(ndarray::array![
                    [DataValue::from(1i32)],
                    [DataValue::from(22i32)]
                ]),
            );
        });
        #[cfg(feature = "polars-df")]
        {
            let pdf = polars::df!(
                "a" => [1u64, 2u64, 3u64],
                "b" => [4f64, 5f64, 6f64],
                "c" => [7i64, 8i64, 9i64]
            )
            .expect("BUG: should be ok");
            let df = DataFrame::from_polars(pyo3_polars::PyDataFrame(pdf));
            assert_eq!(
                df.select(Some(&["a".into(), "b".into(), "c".into()])),
                crate::df! {
                    "a" => [1u64, 2u64, 3u64],
                    "b" => [4f64, 5f64, 6f64],
                    "c" => [7i64, 8i64, 9i64]
                }
                .select(Some(&["a".into(), "b".into(), "c".into()])),
            );
        }
    }

    #[rstest]
    #[traced_test]
    fn basic_ops_add(mut df: DataFrame, hm: HashMap<String, DataValue>) {
        let mut df_expect = df.clone();
        let df2 = df.clone();
        let exec = Python::attach(|py| -> PyResult<()> {
            df.__iadd__(df.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe += df2.dataframe;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            df.__iadd__(hm.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe += hm;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            Ok(())
        });

        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    #[traced_test]
    fn basic_ops_sub(mut df: DataFrame, hm: HashMap<String, DataValue>) {
        let mut df_expect = df.clone();
        let df2 = df.clone();
        let exec = Python::attach(|py| -> PyResult<()> {
            df.__isub__(df.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe -= df2.dataframe;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            df.__isub__(hm.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe -= hm;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            Ok(())
        });

        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    #[traced_test]
    fn basic_ops_mul(mut df: DataFrame, hm: HashMap<String, DataValue>) {
        let mut df_expect = df.clone();
        let df2 = df.clone();
        let exec = Python::attach(|py| -> PyResult<()> {
            df.__imul__(df.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe *= df2.dataframe;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            df.__imul__(hm.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe *= hm;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);
            Ok(())
        });

        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    #[traced_test]
    fn basic_ops_div(mut df: DataFrame, hm: HashMap<String, DataValue>) {
        let mut df_expect = df.clone();
        let df2 = df.clone();
        let exec = Python::attach(|py| -> PyResult<()> {
            df.__itruediv__(df.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe /= df2.dataframe;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);

            df.__itruediv__(hm.clone().into_pyobject(py)?.into_any())?;
            df_expect.dataframe /= hm;
            tracing::trace!("{} vs {}", df, df_expect);
            assert_eq!(df.dataframe, df_expect.dataframe);
            Ok(())
        });

        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    #[traced_test]
    #[rstest]
    fn test_numpy(mut df: DataFrame) {
        let exec = Python::attach(|py| -> PyResult<()> {
            let code = c_str!(
                r#"
def example(df):
    import numpy as np
    a_np = df.as_numpy_f32(['key1', 'key2'])
    print(a_np)
    b_np = df.as_numpy_u32(['key1', 'key'])
    print(b_np)
    b_np = df.as_numpy_i32(['key1', 'key'])
    print(b_np)
    b_np = df.as_numpy_i64(['key1', 'key'])
    print(b_np)
    b_np = df.as_numpy_u64(['key1', 'key'])
    print(b_np)
    b_np = df.as_numpy_f64(['key1', 'key'])
    print(b_np)
    b_np = df.as_numpy_f64(['key1', 'key'], transposed=True)
    print(b_np)
    b_np = df.as_numpy(['key1', 'key'], transposed=True)
    print(b_np)
    b_np = df.as_numpy_str(['key1', 'key'], transposed=True)
    print(b_np)
    return df
            "#
            );
            let fun: Py<PyAny> = PyModule::from_code(py, code, c_str!(""), c_str!(""))?
                .getattr("example")?
                .into();
            let result = fun.call1(py, (df.clone(),));
            assert!(df.py_join(df.clone(), JoinRelation::default()).is_ok());
            // user may not have installed polars, we need to get an error in that
            // case
            if py.import("numpy").is_ok() {
                assert!(result.is_ok(), "{:?}", result);
            } else {
                assert!(result.is_err(), "{:?}", result);
            }
            Ok(())
        });
        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    #[traced_test]
    #[rstest]
    fn test_fill_from_python(df: DataFrame) {
        let exec = Python::attach(|_py| -> PyResult<()> {
            let hm = stdhashmap!(
                Key::from("key1") => DataValue::U32(1),
                Key::from("key2") => DataValue::U32(2),
            );
            let mut df2 = DataFrame::init();
            assert!(df2.py_push(hm).is_ok());
            assert!(df2
                .py_push(stdhashmap!(
                    Key::from("key1") => DataValue::U32(11),
                    Key::from("key2") => DataValue::U32(21),
                ))
                .is_ok());

            assert_eq!(df, df2);

            let mut df2 = DataFrame::init();
            assert!(df2
                .py_add_column(
                    Key::from("key1"),
                    vec![DataValue::U32(1), DataValue::U32(11)]
                )
                .is_ok());
            assert!(df2
                .py_add_column(
                    Key::from("key2"),
                    vec![DataValue::U32(2), DataValue::U32(21)]
                )
                .is_ok());

            assert_eq!(df, df2);
            Ok(())
        });
        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    fn basic_python_dataframe(mut df: DataFrame) {
        let exec = Python::attach(|py| -> PyResult<()> {
            let fun: Py<PyAny> = PyModule::from_code(
                py,
                c_str!(
                    "
def example(df):
    print(df)
    df.shrink()
    assert len(df) == 2
    df.add_alias('key1', 'key1-alias')
    a = df.select(['key1', 'key2'])
    print(a)
    b = df.select(['key1-alias', 'key2'])
    print(b)
    df.rename_key('key1', 'key1new')
    df.rename_key('key1new', 'key1')
    assert a == [[1, 2], [11, 21]]
    assert a == b
    df.add_metadata('test', 1)
    m = df.get_metadata('test')
    assert m == 1
    b = df.select_transposed(['key1', 'key2'])
    print(b)
    assert b == [[1, 11], [2, 21]]
    c = df.select_column('key1')
    print(c)
    assert c == [1, 11]

    a += b
    print(a)
    assert a == [[2, 13], [4, 23]]
    a -= b
    print(a)
    assert e == a
    f = e * b
    print(f)
    assert f == [[1, 22], [44, 441]]
    g = f / b
    print(g)
    assert g == e

                "
                ),
                c_str!(""),
                c_str!(""),
            )?
            .getattr("example")?
            .into();
            let _ = fun.call1(py, (df.clone(),));
            assert!(df.py_join(df.clone(), JoinRelation::default()).is_ok());
            Ok(())
        });
        assert!(exec.is_ok(), "{:?}", exec);
    }

    #[rstest]
    fn dummy_test_apply(mut df: DataFrame) {
        let exec = Python::attach(|py| -> PyResult<()> {
            let fun: Py<PyAny> = PyModule::from_code(
                py,
                c_str!(
                    r#"
def multiply_by_ten(x):
    print(x)
    x *= {"key1": 10}
    print(x)
    return x

def example(df):
    print(df)
    df.apply(multiply_by_ten)
                "#
                ),
                c_str!(""),
                c_str!(""),
            )?
            .getattr("example")?
            .into();
            let _ = fun.call1(py, (df.clone(),));
            assert!(df.py_join(df.clone(), JoinRelation::default()).is_ok());
            Ok(())
        });
        assert!(exec.is_ok(), "{:?}", exec);
    }
}