Read and write to Parquet.

Uses `clap` to parse command line arguments.
"REBERGEN WILLEM" query: 105.1s -> 4.45s. 23x improvement.
This commit is contained in:
Hayden Heroux 2026-09-21 11:53:11 -04:00
parent db03e51cfe
commit f64e589058
4 changed files with 227 additions and 8 deletions

1
.gitignore vendored
View file

@ -1,2 +1,3 @@
ASCAP_CATALOG.csv
*.parquet
/target

169
Cargo.lock generated
View file

@ -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"

View file

@ -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"] }

View file

@ -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<LazyFrame, PolarsError> {
let schema = Schema::from_iter(vec![
@ -10,20 +12,66 @@ fn read_csv(path: PlRefPath) -> Result<LazyFrame, PolarsError> {
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, 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 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")])