From f64e589058d04b4382d3082261c0fe75a2067dec Mon Sep 17 00:00:00 2001 From: Hayden Heroux Date: Mon, 21 Sep 2026 11:53:11 -0400 Subject: [PATCH] Read and write to Parquet. Uses `clap` to parse command line arguments. "REBERGEN WILLEM" query: 105.1s -> 4.45s. 23x improvement. --- .gitignore | 1 + Cargo.lock | 169 ++++++++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 3 +- src/main.rs | 62 ++++++++++++++++--- 4 files changed, 227 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 703b47e..b98097a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ ASCAP_CATALOG.csv +*.parquet /target diff --git a/Cargo.lock b/Cargo.lock index 831d7ea..79ee2d9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -32,6 +47,56 @@ dependencies = [ "libc", ] +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + [[package]] name = "ar_archive_writer" version = "0.5.3" @@ -67,6 +132,7 @@ checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" name = "ascap" version = "0.1.0" dependencies = [ + "clap", "polars", "polars-lazy", ] @@ -201,6 +267,27 @@ version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "36f64beae40a84da1b4b26ff2761a5b895c12adc41dc25aaee1c4f2bbfe97a6e" +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -302,6 +389,52 @@ dependencies = [ "phf", ] +[[package]] +name = "clap" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8876b300ab35ba921adea3dfd70157a46249b33f95c9084ae5709785478946" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0797fb7aeb1406c84efac526901f7ec3ead2124f946b494e72879d4b54704d" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9c751b79415d4e559e3d1fcf128e09e720eb673a06d26cf6f392d37d75b66e0" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "clap_lex" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + [[package]] name = "comfy-table" version = "7.2.2" @@ -1096,6 +1229,12 @@ version = "2.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "791930b43c0d5973160d90a8f3894509f2b273430f5c5c73b668636d0287c5c0" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + [[package]] name = "itertools" version = "0.14.0" @@ -1348,6 +1487,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1820,10 +1965,13 @@ checksum = "d33a17979d2afa0309437478662cf996361a6e1768bf467db82bd4f4a9443064" dependencies = [ "async-stream", "base64", + "brotli", "bytemuck", "ethnum", + "flate2", "futures", "hashbrown 0.17.1", + "lz4", "num-traits", "polars-arrow", "polars-buffer", @@ -1835,7 +1983,9 @@ dependencies = [ "regex", "serde", "simdutf8", + "snap", "streaming-decompression", + "zstd", ] [[package]] @@ -1877,6 +2027,7 @@ dependencies = [ "polars-error", "polars-io", "polars-ops", + "polars-parquet", "polars-time", "polars-utils", "rayon", @@ -2618,6 +2769,12 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9be42f50aa861c555654aa3a37f52f4b1074bacf4e48fe0ef7fa584e80f1f0f" +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + [[package]] name = "socket2" version = "0.6.5" @@ -2696,6 +2853,12 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "strum_macros" version = "0.28.0" @@ -3028,6 +3191,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.26.1" diff --git a/Cargo.toml b/Cargo.toml index d1bbd62..67f297f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,5 +4,6 @@ version = "0.1.0" edition = "2024" [dependencies] -polars = { version = "0.55.2", features = ["lazy", "strings", "regex"] } +clap = { version = "4.0", features = ["derive"] } +polars = { version = "0.55.2", features = ["lazy", "strings", "regex", "parquet"] } polars-lazy = { version = "0.55.2", features = ["csv", "cum_agg"] } diff --git a/src/main.rs b/src/main.rs index ed2c671..bbec0df 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,7 @@ +use clap::{Parser, ValueEnum}; use polars::prelude::*; use polars_lazy::frame::LazyCsvReader; +use std::sync::Arc; fn read_csv(path: PlRefPath) -> Result { let schema = Schema::from_iter(vec![ @@ -10,20 +12,66 @@ fn read_csv(path: PlRefPath) -> Result { 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")]); + 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::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, +} + fn main() -> Result<(), PolarsError> { - let path: PlRefPath = PlRefPath::new("ASCAP_CATALOG.csv"); - let name: &str = "REBERGEN WILLEM"; + let args = Cli::parse(); - let lf = read_csv(path)?; + let all_data = match args.source { + SourceKind::CSV => read_csv(args.from)?, + SourceKind::Parquet => read_parquet(args.from)?, + }; - let writer_is_name = col("RoleType").eq(lit("W")).and(col("Name").str().contains_literal(lit(name))); + if let Some(to) = args.to { + write_parquet(&all_data, to)?; + } - let song_titles = lf.filter(writer_is_name) + 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")])