rustalib/util/
dataframe_utils.rs

1use polars::frame::column::Column;
2use polars::prelude::*;
3
4/// Ensure a column in a DataFrame is of Float64 type
5///
6/// # Arguments
7///
8/// * `df` - DataFrame to modify
9/// * `column_name` - Name of the column to convert
10///
11/// # Returns
12///
13/// Returns a PolarsResult indicating success or failure
14///
15/// # Example
16///
17/// ```
18/// use polars::prelude::*;
19/// use ta_lib_in_rust::util::dataframe_utils::ensure_f64_column;
20///
21/// let mut df = DataFrame::new(vec![
22///     Series::new("price", &[1, 2, 3]),
23/// ]).unwrap();
24/// ensure_f64_column(&mut df, "price").unwrap();
25/// assert_eq!(df.column("price").unwrap().dtype(), &DataType::Float64);
26/// ```
27pub fn ensure_f64_column(df: &mut DataFrame, column_name: &str) -> PolarsResult<()> {
28    // 1) Wrap the existing Series in a Column for in-place mutation
29    let s: Series = df.column(column_name)?.as_materialized_series().clone();
30    let mut col: Column = s.into_column();
31
32    // 2) Materialize and get a &mut Series to cast in place
33    let series_mut: &mut Series = col.into_materialized_series();
34    *series_mut = series_mut.cast(&DataType::Float64)?;
35
36    // 3) Convert the Column back into a Series and replace it in the DataFrame
37    let series: Series = col.take_materialized_series();
38    df.replace(column_name, series)?;
39
40    Ok(())
41}
42
43/// Check if a DataFrame has enough rows for a given window size
44///
45/// # Arguments
46///
47/// * `df` - The DataFrame to check
48/// * `window` - The window size required
49/// * `indicator_name` - Name of the indicator (for error message)
50///
51/// # Returns
52///
53/// Returns a PolarsResult<()> or an error if there are not enough rows
54///
55/// # Example
56///
57/// ```
58/// use polars::prelude::*;
59/// use ta_lib_in_rust::util::dataframe_utils::check_window_size;
60///
61/// let df = DataFrame::new(vec![Series::new("close", &[1.0, 2.0, 3.0, 4.0])]).unwrap();
62/// assert!(check_window_size(&df, 3, "test").is_ok());
63/// assert!(check_window_size(&df, 5, "test").is_err());
64/// ```
65pub fn check_window_size(df: &DataFrame, window: usize, indicator_name: &str) -> PolarsResult<()> {
66    if df.height() < window {
67        return Err(PolarsError::ComputeError(
68            format!(
69                "Not enough data points ({}) for {} window ({})",
70                df.height(),
71                indicator_name,
72                window
73            )
74            .into(),
75        ));
76    }
77    Ok(())
78}
rustalib/util/dataframe_utils.rs

rustalib/util/
dataframe_utils.rs
