From 832d876af068da753291800b14e1139d41e833a7 Mon Sep 17 00:00:00 2001 From: AC Date: Fri, 16 Oct 2026 14:32:00 +0800 Subject: [PATCH] Fix parallel Polars reads with all-missing chunks --- change_log.md | 1 + pyreadstat/_readstat_parser.pyx | 4 +-- pyreadstat/pyreadstat.py | 20 ++++++++++++ tests/test_narwhalified.py | 57 +++++++++++++++++++++++++++++++++ 4 files changed, 80 insertions(+), 2 deletions(-) diff --git a/change_log.md b/change_log.md index f5610fd..88be038 100644 --- a/change_log.md +++ b/change_log.md @@ -1,4 +1,5 @@ # 1.3.7 (github, pypi and conda 2026.xx.xx) +* Fix parallel Polars reads when a chunk contains only missing values, including datetime columns (#350) * Updated Readstat sources to 98813868140c835fd56ce61de5894c3359b15c5e * Updated source code and tests for polars 2 * variable_storage_width for SPSS string variables now returns the declared diff --git a/pyreadstat/_readstat_parser.pyx b/pyreadstat/_readstat_parser.pyx index b8c9663..555f5a5 100644 --- a/pyreadstat/_readstat_parser.pyx +++ b/pyreadstat/_readstat_parser.pyx @@ -1117,8 +1117,8 @@ cdef object dict_to_dataframe(object dict_data, data_container dc): data_frame = data_frame.with_columns( [ pl.from_epoch( - (pl.col(c) % 1 * 1e6).round().cast(pl.Int64) + ( - pl.col(c).floor() * 1e6).cast(pl.Int64) - ( + (pl.col(c).cast(pl.Float64) % 1 * 1e6).round().cast(pl.Int64) + ( + pl.col(c).cast(pl.Float64).floor() * 1e6).cast(pl.Int64) - ( pl.lit(dc.unix_to_origin_secs) * 1e6).cast(pl.Int64), time_unit='us') for c in datetime_cols if data_frame[c].len() > 0 diff --git a/pyreadstat/pyreadstat.py b/pyreadstat/pyreadstat.py index 993497d..0be7cd9 100644 --- a/pyreadstat/pyreadstat.py +++ b/pyreadstat/pyreadstat.py @@ -1197,6 +1197,26 @@ def read_file_multiprocessing( final = {key: list(chain.from_iterable(chunk[key] for chunk in chunks)) for key in keys} else: # final = pd.concat(chunks, axis=0, ignore_index=True) + if output_format == "polars": + import polars as pl + + # All-null chunks have no inferred type. Use the type from a + # populated chunk without coercing conflicting non-null types. + schemas = [chunk.schema for chunk in chunks] + populated_types = { + name: dtype + for schema in schemas + for name, dtype in schema.items() + if dtype != pl.Null + } + chunks = [ + chunk.cast({ + name: populated_types[name] + for name, dtype in schema.items() + if dtype == pl.Null and name in populated_types + }) + for chunk, schema in zip(chunks, schemas) + ] chunks = [nw.from_native(x) for x in chunks] final = nw.concat(chunks, how="vertical") ispandas = False diff --git a/tests/test_narwhalified.py b/tests/test_narwhalified.py index 9e3299f..591c3e8 100644 --- a/tests/test_narwhalified.py +++ b/tests/test_narwhalified.py @@ -18,6 +18,7 @@ from datetime import datetime, timedelta, date from concurrent.futures import ThreadPoolExecutor import unittest +from unittest.mock import patch import os import sys import shutil @@ -757,6 +758,62 @@ def test_multiprocess_reader(self): self.assertTrue(df_multi.equals(df_single)) self.assertEqual(meta_multi.number_rows, meta_single.number_rows) + def test_multiprocess_reader_sparse_columns(self): + if self.backend != "polars": + self.skipTest("Polars null column inference") + data = pd.DataFrame({ + "number": [None] * 4 + [1.5, 2.5, 3.5, 4.5], + "empty": [float("nan")] * 8, + "text": [""] * 4 + ["one", "two", "three", "four"], + "stamp": pd.to_datetime([None] * 4 + ["2025-01-01"] * 4), + "day": [None] * 4 + [date(2025, 1, 1)] * 4, + "clock": [None] * 4 + [datetime(2025, 1, 1, 12, 30).time()] * 4, + }) + with tempfile.TemporaryDirectory() as directory, patch( + "pyreadstat.pyreadstat.mp.Pool", mp.get_context("spawn").Pool, + ): + path = os.path.join(directory, "sparse.sav") + for reverse in (False, True): + with self.subTest(reverse=reverse): + frame = data.iloc[::-1] if reverse else data + pyreadstat.write_sav(frame, path, missing_ranges={"text": [""]}) + for selection in ({}, {"row_offset": 1, "row_limit": 6}, + {"row_offset": 0, "row_limit": 4}): + with self.subTest(selection=selection): + expected, meta = pyreadstat.read_sav( + path, output_format=self.backend, **selection, + ) + actual, actual_meta = pyreadstat.read_file_multiprocessing( + pyreadstat.read_sav, path, num_processes=2, + output_format=self.backend, **selection, + ) + self.assertTrue(actual.equals(expected)) + self.assertEqual(actual_meta.column_names, meta.column_names) + self.assertEqual(actual_meta.number_rows, len(data)) + if self.backend == "polars": + self.assertEqual(actual.schema, expected.schema) + + def test_multiprocess_reader_sparse_value_labels(self): + if self.backend != "polars": + self.skipTest("Polars null column inference") + data = pd.DataFrame({"value": [None] * 4 + [1.0] * 4}) + with tempfile.TemporaryDirectory() as directory, patch( + "pyreadstat.pyreadstat.mp.Pool", mp.get_context("spawn").Pool, + ): + path = os.path.join(directory, "labels.sav") + pyreadstat.write_sav(data, path, variable_value_labels={"value": {1: "one"}}) + for categorical in (False, True): + with self.subTest(categorical=categorical): + kwargs = dict(output_format=self.backend, apply_value_formats=True, + formats_as_category=categorical) + expected, _ = pyreadstat.read_sav(path, **kwargs) + actual, _ = pyreadstat.read_file_multiprocessing( + pyreadstat.read_sav, path, num_processes=2, **kwargs, + ) + self.assertTrue(actual.equals(expected)) + if self.backend == "polars": + self.assertEqual(actual.schema, expected.schema) + def test_chunk_reader_multiprocess(self): fpath = os.path.join(self.basic_data_folder, "sample_large.sav") reader = pyreadstat.read_file_in_chunks(pyreadstat.read_sav, fpath, chunksize= 50, multiprocess=True, output_format=self.backend)