jinschoi
commented
11 months ago Reading the s3 file into a std::io::Cursor and creating a ParquetReader from it only takes 5s as well:
use std::io::{Read, Cursor};
use cloud::AmazonS3ConfigKey as Key;
use dotenvy::dotenv;
use eyre::{Result, Report};
use polars::prelude::*;
use rusoto_core::Region;
use rusoto_s3::{GetObjectRequest, S3Client, S3};
use tokio::task;
async fn cursor_for_s3(key: &str) -> Result<Cursor<Vec<u8>>> {
dbg!(key);
let s3_client = S3Client::new(Region::default());
let req = GetObjectRequest {
bucket: "s3bucket".to_owned(),
key: key.to_owned(),
..Default::default()
};
let mut reader = s3_client
.get_object(req)
.await?
.body
.unwrap()
.into_blocking_read();
let mut buf = vec![];
let join_handle: task::JoinHandle<Result<Cursor<Vec<u8>>, Report>> = task::spawn_blocking(move || {
reader.read_to_end(&mut buf)?;
Ok(Cursor::new(buf))
});
join_handle.await?
}
macro_rules! time {
($msg:expr, $($stmts:stmt);+) => {
let now = std::time::Instant::now();
$($stmts)*
println!("{}: {:.2?}", $msg, now.elapsed());
};
}
pub async fn dataframe_from_s3_key_rusoto(key: &str) -> Result<DataFrame> {
time!("parquet read",
let mut cursor = cursor_for_s3(key).await?);
time!("df conversion",
let df = ParquetReader::new(&mut cursor).finish()?);
Ok(df)
}
#[tokio::main]
async fn main() -> Result<()> {
dotenv()?;
let df = dataframe_from_s3_key_rusoto("xxx").await?;
dbg!(&df);
Ok(())
}
allinux
ritchie46
Checks
[X] I have checked that this issue has not already been reported.
[X] I have confirmed this bug exists on the latest version of Polars.
Reproducible example
Log output
Issue description
Running the above code on a 60MB parquet file in s3 takes 27s. The equivalent Python code:
takes ~5s. While running with POLARS_VERBOSE=1, each of those BatchedParquetReader lines outputs about one per second.
Expected behavior
The rust code should be at least as fast as the Python version.
Installed versions