From 97ef657a50f7de8a55ef90a3722b0369f4be2e71 Mon Sep 17 00:00:00 2001
From: Hayden Heroux <haydenheroux@gmail.com>
Date: Mon, 21 Sep 2026 13:14:02 -0400
Subject: [PATCH] Query collaborators for artist.

---
 src/main.rs | 42 ++++++++++++++++++++++++++++++------------
 1 file changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index bbec0df..26c3f7f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -40,7 +40,7 @@ fn write_parquet(lf: &LazyFrame, path: PlRefPath) -> Result<(), PolarsError> {
     Ok(())
 }
 
-#[derive(Clone, ValueEnum)]
+#[derive(PartialEq, Clone, ValueEnum)]
 enum SourceKind {
     CSV,
     Parquet,
@@ -48,9 +48,10 @@ enum SourceKind {
 
 #[derive(Parser)]
 struct Cli {
-    source: SourceKind,
     from: PlRefPath,
     name: String,
+    #[arg(value_enum, default_value = "parquet")]
+    source: SourceKind,
     to: Option<PlRefPath>,
 }
 
@@ -62,22 +63,39 @@ fn main() -> Result<(), PolarsError> {
         SourceKind::Parquet => read_parquet(args.from)?,
     };
 
-    if let Some(to) = args.to {
-        write_parquet(&all_data, to)?;
+    if args.source == SourceKind::CSV {
+        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 is_writer = col("RoleType").eq(lit("W"));
+    let contains_name = col("Name").str().contains_literal(lit(args.name.clone()));
 
-    let song_titles = all_data
-        .filter(writer_is_name)
+    let artist_song_ids = all_data
+        .clone()
+        .filter(is_writer.clone().and(contains_name.clone()))
+        .select([col("SongID")])
         .unique(None, UniqueKeepStrategy::Any)
-        .sort(["Title"], Default::default())
-        .select([col("SongID"), col("Title")])
         .collect()?;
 
-    println!("{}", song_titles);
+    let collaborators = all_data
+        .filter(is_writer)
+        .join(
+            artist_song_ids.lazy(),
+            [col("SongID")],
+            [col("SongID")],
+            JoinArgs::new(JoinType::Inner),
+        )
+        .filter(contains_name.not())
+        .group_by([col("Name")])
+        .agg([len().alias("Count")])
+        .sort(
+            ["Count"],
+            SortMultipleOptions::default().with_order_descending(true),
+        );
+
+    println!("{}", collaborators.collect()?);
 
     Ok(())
 }