ascap/src/main.rs

use clap::{Parser, ValueEnum};
use polars::prelude::*;
use polars_lazy::frame::LazyCsvReader;
use std::sync::Arc;

fn read_csv(path: PlRefPath) -> Result<LazyFrame, PolarsError> {
    let schema = Schema::from_iter(vec![
        Field::new("Title".into(), DataType::String),
        Field::new("RoleType".into(), DataType::String),
        Field::new("Name".into(), DataType::String),
        Field::new("Shares".into(), DataType::String),
        Field::new("Note".into(), DataType::String),
    ]);

    let lf = LazyCsvReader::new(path)
        .with_has_header(true)
        .with_schema(Some(Arc::new(schema)))
        .finish()?;
    let with_song_id = lf.with_columns([col("RoleType")
        .eq(lit("ASCAP"))
        .cast(DataType::Int64)
        .cum_sum(false)
        .alias("SongID")]);
    Ok(with_song_id)
}

fn read_parquet(path: PlRefPath) -> Result<LazyFrame, PolarsError> {
    LazyFrame::scan_parquet(path, ScanArgsParquet::default())
}

fn write_parquet(lf: &LazyFrame, path: PlRefPath) -> Result<(), PolarsError> {
    let sink = lf.clone().sink(
        SinkDestination::File {
            target: SinkTarget::Path(path),
        },
        FileWriteFormat::Parquet(Arc::new(ParquetWriteOptions::default())),
        UnifiedSinkArgs::default(),
    )?;
    sink.collect()?;
    Ok(())
}

#[derive(Clone, ValueEnum)]
enum SourceKind {
    CSV,
    Parquet,
}

#[derive(Parser)]
struct Cli {
    source: SourceKind,
    from: PlRefPath,
    name: String,
    to: Option<PlRefPath>,
}

fn main() -> Result<(), PolarsError> {
    let args = Cli::parse();

    let all_data = match args.source {
        SourceKind::CSV => read_csv(args.from)?,
        SourceKind::Parquet => read_parquet(args.from)?,
    };

    if let Some(to) = args.to {
        write_parquet(&all_data, to)?;
    }

    let writer_is_name = col("RoleType")
        .eq(lit("W"))
        .and(col("Name").str().contains_literal(lit(args.name)));

    let song_titles = all_data
        .filter(writer_is_name)
        .unique(None, UniqueKeepStrategy::Any)
        .sort(["Title"], Default::default())
        .select([col("SongID"), col("Title")])
        .collect()?;

    println!("{}", song_titles);

    Ok(())
}