[Python] polars scan_csv bad file resiliency

This might be a bug, but I also have a more general question about scan_csv (and scan_parquet, too) resiliency to bad files. First, here's some sample code. If I try saving files using Option A, I get the error:

polars.exceptions.ColumnNotFoundError: did not find column column_1, consider enabling `allow_missing_columns`


Which isn't so helpful; allow_missing_columns is a kwarg for only read/scan_parquet.

If I use Option B, it works (I know the naming pattern is different, that's not important). The problem is that Option B is slower and uses significantly more memory (obviously this toy problem doesn't use much memory, but a meaningfully big problem could).

import numpy as np
import polars as pl

if __name__ == "__main__":
base_data_fname = "/tmp/base_data.csv"
base_data = pl.DataFrame(
{"column_1": np.arange(1000), "column_2": np.random.random(1000)},
)
base_data.write_csv(base_data_fname, include_header=False)

newline_fname = "/tmp/newline.csv"
with open(newline_fname, "w") as fp:
fp.write("\n")

schema = {
"column_1": pl.Int64,
"column_2": pl.Float64,
}

data = pl.scan_csv(
[base_data_fname, newline_fname],
has_header=False,
ignore_errors=True,
schema_overrides=schema,
infer_schema=False,
raise_if_empty=False,
)

# To reduce the number of output groups.
data = data.with_columns(
[pl.col("column_1").mod(17).alias("column_1")],
)

# Option A: Doesn't work with newline file.
# Comment this section to use Option B
data.sink_parquet(
pl.PartitionByKey(
base_path="/tmp/deleteme/",
file_path=lambda ctx: (
f"column_1={ctx.keys[0].str_value}/0.parquet"
),
by="column_1",
),
mkdir=True,
)
# End Option A

# Option B: Does work with newline file.
data = data.collect()
for c1, subdata in data.group_by("column_1"):
subdata.write_parquet(
f"/tmp/deleteme/{c1[0]}.parquet",
)
# End Option B


First, is this a bug? I think that the raise_if_empty flag is working here, the newline file isn't empty. So maybe it's not a bug, but there's no way to catch this functionally empty file as far as I can tell.

Therefore, my higher-level question is if there is a way to handle problem files like this gracefully with scan_csv (and scan_parquet, too) sticking with a LazyFrame (Option A)? If I'm processing thousands of files and only one of them is a problem, is there a way to just not use that file? I want any file that might be broken, empty, or inaccessible (e.g. if the retries count is hit) to be thrown out so that the rest of the process works.

Thanks!

Continue reading...