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:
parent
db03e51cfe
commit
f64e589058
4 changed files with 227 additions and 8 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -1,2 +1,3 @@
|
|||
ASCAP_CATALOG.csv
|
||||
*.parquet
|
||||
/target
|
||||
|
|
|
|||
169
Cargo.lock
generated
169
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
62
src/main.rs
62
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<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")])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue