Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: Fix decompress_impl for csv with n_rows set #17118

Merged
merged 6 commits into from
Jun 22, 2024
Merged

fix: Fix decompress_impl for csv with n_rows set #17118

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ce93aaa
Adds "polars-ops/timezones" dependency for "timezones" feature
Mottl Jun 21, 2024
c4c74f2
formatting in Cargo.toml
Mottl Jun 21, 2024
aca8984
removes WIP tests
Mottl Jun 21, 2024
5ad093f
Fixes unused_imports warning
Mottl Jun 21, 2024
f61fa31
fixes decompress_impl for csv with n_rows set
Mottl Jun 21, 2024
6fb4b51
Removes test for reading a compressed csv file
Mottl Jun 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions crates/polars-io/src/csv/read/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,9 @@ fn decompress_impl<R: Read>(
},
};
}
if line_count == n_rows {
out.truncate(buf_pos); // retain only first n_rows in out
}
out
},
})
Expand Down
8 changes: 7 additions & 1 deletion crates/polars/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,13 @@ string_pad = ["polars-lazy?/string_pad", "polars-ops/string_pad"]
string_reverse = ["polars-lazy?/string_reverse", "polars-ops/string_reverse"]
string_to_integer = ["polars-lazy?/string_to_integer", "polars-ops/string_to_integer"]
take_opt_iter = ["polars-core/take_opt_iter"]
timezones = ["polars-core/timezones", "polars-lazy?/timezones", "polars-io/timezones", "polars-sql?/timezones"]
timezones = [
"polars-core/timezones",
"polars-lazy?/timezones",
"polars-io/timezones",
"polars-ops/timezones",
"polars-sql?/timezones",
]
to_dummies = ["polars-ops/to_dummies"]
top_k = ["polars-lazy?/top_k"]
trigonometry = ["polars-lazy?/trigonometry"]
Expand Down
17 changes: 16 additions & 1 deletion crates/polars/tests/it/io/csv.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ use std::io::Cursor;
use std::num::NonZeroUsize;

use polars::io::RowIndex;
use polars_core::export::chrono;
use polars_core::utils::concat_df;

use super::*;
Expand Down Expand Up @@ -41,7 +40,10 @@ fn write_csv() {
}

#[test]
#[cfg(feature = "timezones")]
fn write_dates() {
use polars_core::export::chrono;

let s0 = Series::new("date", [chrono::NaiveDate::from_yo_opt(2024, 33), None]);
let s1 = Series::new("time", [None, chrono::NaiveTime::from_hms_opt(19, 50, 0)]);
let s2 = Series::new(
Expand Down Expand Up @@ -1389,3 +1391,16 @@ fn test_read_io_reader() {
let expected = CsvReader::new(file).finish().unwrap();
assert!(df.equals(&expected))
}

#[test]
#[cfg(any(feature = "decompress", feature = "decompress-fast"))]
fn test_read_compressed() {
const COMPRESSED_CSV: &str = "../../examples/datasets/compressed.csv.gz";
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we create the compressed file on the fly (in-memory). That saves adding a file to the repo.

Copy link
Contributor Author

@Mottl Mottl Jun 22, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Current compressed.csv.gz is randomly generated and the n_rows which yields to ComputeError is found manually. I'd leave compressed.csv.gz as-is or remove test_read_compressed completely, since it's mostly here for illustrating the initial issue.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we can create a in memory random file that illustrates this. Otherwise removing the test and file is also fine. 👍

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A test_read_compressed and a correspondent csv.gz file are removed 👍

let df = CsvReadOptions::default()
.with_n_rows(Some(2367))
.try_into_reader_with_file_path(Some(COMPRESSED_CSV.into()))
.unwrap()
.finish()
.unwrap();
assert_eq!(df.shape(), (2367, 3));
}
Binary file added examples/datasets/compressed.csv.gz
View file
Open in desktop
Binary file not shown.
Loading