search

pola-rs / polars

Dataframes powered by a multithreaded, vectorized query engine, written in Rust
https://docs.pola.rs
Other
30.52k stars 1.98k forks source link

Imperfect behavior with `scan_csv()` on `zstd` compressed file with `new_columns` param #19916

Open kscott-1 opened 2 days ago

kscott-1 commented 2 days ago

Checks

Reproducible example

import numpy as np
import polars as pl
import zstandard as zstd

#n_rows = 6 # works
n_rows = 7 # (or higher) doesn't work

df = pl.DataFrame(
    {
        "foo": np.random.randint(1, 10, n_rows),
        "bar": np.random.randint(1, 10, n_rows),
        "baz": np.random.randint(1, 10, n_rows),
        "qux": np.random.randint(1, 10, n_rows),
        "quux": np.random.randint(1, 10, n_rows),
        "corge": np.random.randint(1, 10, n_rows),
    }
)

txtdf = df.write_csv(None, separator="\t", include_header=False)
with open("foobar.txt.zst", "wb") as zfile:
    zfile.write(zstd.ZstdCompressor().compress(txtdf.encode()))

pl.scan_csv(
    "foobar.txt.zst",
    separator="\t",
    has_header=False,
    new_columns=["foo", "bar", "baz", "qux", "quux", "corge"],
).collect()

Log output

Traceback (most recent call last):
  File "[...]", line 22, in <module>
    pl.scan_csv(
    ~~~~~~~~~~~^
        "foobar.txt.zst",
        ^^^^^^^^^^^^^^^^^
    ...<2 lines>...
        new_columns=["foo", "bar", "baz", "qux", "quux", "corge"],
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ).collect()
    ^
  File "[...]/lib/python3.13/site-packages/polars/_utils/deprecation.py", line 92, in wrapper
    return function(*args, **kwargs)
  File "[...]/lib/python3.13/site-packages/polars/_utils/deprecation.py", line 92, in wrapper
    return function(*args, **kwargs)
  File "[...]/lib/python3.13/site-packages/polars/_utils/deprecation.py", line 92, in wrapper
    return function(*args, **kwargs)
  File "[...]/lib/python3.13/site-packages/polars/io/csv/functions.py", line 1296, in scan_csv
    return _scan_csv_impl(
        source,
    ...<30 lines>...
        include_file_paths=include_file_paths,
    )
  File "[...]/lib/python3.13/site-packages/polars/io/csv/functions.py", line 1394, in _scan_csv_impl
    pylf = PyLazyFrame.new_from_csv(
        source,
    ...<30 lines>...
        include_file_paths=include_file_paths,
    )
polars.exceptions.ShapeError: The length of the new names list should be equal to or less than the original column length

Issue description

When scanning zstd compressed text files, there is strange behavior on whether or not it is possible to define a new set of column names for a headerless file without error. This is related to #17841 which seems to have brought forth the support that didn't exist for compressed files in scan_csv() previously.

Expected behavior

A list of column names should populate a lazy frame during scanning when the number of columns is equal to the number of elements in the list.

Installed versions

``` --------Version info--------- Polars: 1.14.0 Index type: UInt32 Platform: macOS-15.1-arm64-arm-64bit-Mach-O Python: 3.13.0 (main, Oct 7 2024, 05:02:14) [Clang 16.0.0 (clang-1600.0.26.3)] LTS CPU: False ----Optional dependencies---- adbc_driver_manager altair boto3 cloudpickle connectorx deltalake fastexcel fsspec gevent google.auth great_tables matplotlib nest_asyncio numpy 2.1.3 openpyxl pandas pyarrow pydantic pyiceberg sqlalchemy torch xlsx2csv xlsxwriter ```