Search code examples
rustrust-polars

How to properly apply a MAP function


Need some help with a map function.

I want to take a DataType::Date and store the corresponding weekday as a string column.

I have it working starting with string -> date-type -> string (case #1).

What I am looking for is date-type -> string (case#2).

Here is the working code for the first case ... any suggestions on how to get this to work for my second case?

My challenge with this stems from my lack of proper understanding of how map is supposed to work in this instance.

use chrono::{Date, Datelike, NaiveDate, Utc};
use polars::prelude::*;

fn main() {
    let days = df!("column_1" => &["Tuesday"],
                                "column_2" => &["1900-01-02"]);

    let options = StrpTimeOptions {
        date_dtype: DataType::Date,
        fmt: Some("%Y-%m-%d".into()),
        strict: false,
        exact: true,
    };

    // convert column_2-string into dtype(date) and put into new column "date"
    let days = days
        .unwrap()
        .lazy()
        .with_column(col("column_2").alias("date").str().strptime(options));

    let o = GetOutput::from_type(DataType::Utf8);
    fn str_to_weekday(str_val: Series) -> Result<Series> {
        let x = str_val
            .utf8()
            .unwrap()
            .into_iter()
            // your actual custom function would be in this map
            .map(|opt_date: Option<&str>| {
                opt_date.map(|date: &str| {
                    // for DEBUG purpose only:
                    println! {"Date-String: {:?}", date};
                    NaiveDate::parse_from_str(date, "%Y-%m-%d")
                        .unwrap()
                        .format("%A")
                        .to_string()
                })
            })
            .collect::<Utf8Chunked>();
        Ok(x.into_series())
    }

    // column_2 to weekday-string ... into new column "weekday"
    let days = days
        .with_column(col("column_2").alias("weekday").apply(str_to_weekday, o))
        .collect()
        .unwrap()
        .lazy();

    println!("{:?}", days.clone().collect());
}


Solution

  • Got it to work :)

    With a simplified approach ...

    use polars::prelude::*;
    
    fn main() {
        let days = df!("column_1" => &["Tuesday"],
                                    "column_2" => &["1900-01-02"]);
    
        let options = StrpTimeOptions {
            date_dtype: DataType::Date,
            fmt: Some("%Y-%m-%d".into()),
            strict: false,
            exact: true,
        };
    
        // convert column_2-string into dtype(date) and put into new column "date"
        let days = days
            .unwrap()
            .lazy()
            .with_column(col("column_2").alias("date").str().strptime(options));
    
        println!("{:?}", days.clone().collect());
    
        let o = GetOutput::from_type(DataType::Utf8);
    
        let days = days.with_column(
            col("date")
                .alias("weekday")
                .map(|x| Ok(x.strftime("%A").unwrap()), o),
        );
    
        println!("{:?}", days.collect());
    }