Skip to content

Commit 6fd90af

Browse files
vini-fdavini-fda
authored
Fix row index miscalculation in ParquetLoader (#810)
* fix bug where parquet tables with non-uniform row group sizes get their index miscalculated * remove test with shuffle due to ruff lint * force sequential execution on tests sharing the global default cache dir * fix: handle Windows drive letters in get_parquet_indexer_cls * fix: run xdist_group marker hook before xdist processes nodeids * fix: isolate optimize() cache dir per test to avoid /tmp/chunks race
1 parent 4b8c43d commit 6fd90af

7 files changed

Lines changed: 142 additions & 12 deletions

File tree

‎.github/workflows/ci-testing.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ jobs:
6262
pytest tests \
6363
--ignore=tests/processing \
6464
--ignore=tests/raw \
65-
-n 2 --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
65+
-n 2 --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose
6666
6767
- name: Run processing tests sequentially
6868
run: |

‎src/litdata/streaming/item_loader.py‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -642,6 +642,7 @@ def setup(
642642
self._df: dict[int, Any] = {}
643643
self._chunk_row_groups: dict[int, Any] = {}
644644
self._chunk_row_group_item_read_count: dict[int, Any] = {}
645+
self._chunk_row_group_offsets: dict[int, list[int]] = {}
645646

646647
def generate_intervals(self) -> list[Interval]:
647648
intervals = []
@@ -712,18 +713,28 @@ def _get_item_with_low_memory(self, chunk_index: int, chunk_filepath: str, row_i
712713
Returns:
713714
Any: The dataframe row corresponding to the specified index.
714715
"""
716+
import bisect
717+
715718
import polars as pl
716719
import pyarrow.parquet as pq
717720

718721
# Load the Parquet file metadata if not already loaded
719722
if chunk_index not in self._df:
720-
self._df[chunk_index] = pq.ParquetFile(chunk_filepath)
721-
722-
# Determine the row group and the row index within the row group
723-
parquet_file = self._df[chunk_index]
724-
num_rows_per_row_group = parquet_file.metadata.row_group(0).num_rows
725-
row_group_index = row_index // num_rows_per_row_group
726-
row_index_within_group = row_index % num_rows_per_row_group
723+
parquet_file = pq.ParquetFile(chunk_filepath)
724+
self._df[chunk_index] = parquet_file
725+
# Precompute cumulative row offsets as a prefix-sum so lookup works for row groups of any size.
726+
offsets = [0]
727+
num_row_groups = parquet_file.metadata.num_row_groups
728+
for i in range(num_row_groups):
729+
num_rows = parquet_file.metadata.row_group(i).num_rows
730+
offsets.append(offsets[-1] + num_rows)
731+
self._chunk_row_group_offsets[chunk_index] = offsets
732+
733+
# Locate the row group containing row_index and the offset inside it.
734+
offsets = self._chunk_row_group_offsets[chunk_index]
735+
row_group_index = bisect.bisect_right(offsets, row_index) - 1
736+
row_index_within_group = row_index - offsets[row_group_index]
737+
row_group_size = offsets[row_group_index + 1] - offsets[row_group_index]
727738

728739
# Check if the row group is already loaded
729740
if chunk_index in self._chunk_row_groups and row_group_index in self._chunk_row_groups[chunk_index]:
@@ -746,7 +757,7 @@ def _get_item_with_low_memory(self, chunk_index: int, chunk_filepath: str, row_i
746757

747758
# Check if the row group has been fully read and release memory if necessary
748759
read_count = self._chunk_row_group_item_read_count[chunk_index][row_group_index]
749-
if read_count >= num_rows_per_row_group:
760+
if read_count >= row_group_size:
750761
# Release memory for the fully read row group
751762
del self._chunk_row_groups[chunk_index][row_group_index]
752763
del self._chunk_row_group_item_read_count[chunk_index][row_group_index]
@@ -797,6 +808,8 @@ def delete(self, chunk_index: int, chunk_filepath: str) -> None:
797808

798809
if chunk_index in self._chunk_row_group_item_read_count:
799810
del self._chunk_row_group_item_read_count[chunk_index]
811+
if chunk_index in self._chunk_row_group_offsets:
812+
del self._chunk_row_group_offsets[chunk_index]
800813
if os.path.exists(chunk_filepath):
801814
os.remove(chunk_filepath)
802815
logger.debug(
@@ -820,5 +833,8 @@ def close(self, chunk_index: int) -> None:
820833
if chunk_index in self._chunk_row_group_item_read_count:
821834
del self._chunk_row_group_item_read_count[chunk_index]
822835

836+
if chunk_index in self._chunk_row_group_offsets:
837+
del self._chunk_row_group_offsets[chunk_index]
838+
823839
def encode_data(self, data: list[bytes], sizes: list[int], flattened: list[Any]) -> Any:
824840
pass

‎src/litdata/utilities/parquet.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,12 @@ def get_parquet_indexer_cls(
326326

327327
obj = parse.urlparse(dir_path)
328328

329-
if obj.scheme in ("local", ""):
329+
# On Windows, paths like "C:\\Users\\..." parse with scheme='c'. A single-letter
330+
# scheme is never a valid URI scheme (RFC 3986 requires >=2 chars) so treat it
331+
# as a Windows drive letter and dispatch to LocalParquetDir.
332+
is_windows_drive = len(obj.scheme) == 1 and obj.scheme.isalpha()
333+
334+
if obj.scheme in ("local", "") or is_windows_drive:
330335
return LocalParquetDir(*args)
331336
if obj.scheme in ("gs", "s3"):
332337
return CloudParquetDir(*args)

‎tests/conftest.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,21 @@ def _thread_police():
162162
raise AssertionError(f"Test left zombie thread: {thread}")
163163

164164

165+
@pytest.hookimpl(tryfirst=True)
166+
def pytest_collection_modifyitems(items):
167+
"""Force tests sharing the global default cache dir onto the same xdist worker.
168+
169+
Why: ``clean_pq_index_cache`` shutil.rmtree's ``~/.lightning/chunks`` on setup/teardown.
170+
Under ``pytest -n>1`` two such tests on different workers race, deleting each other's
171+
chunks mid-iteration. Requires the runner to pass ``--dist=loadgroup``. Must run with
172+
``tryfirst=True`` so the marker is in place before xdist's own hook reads it and
173+
appends the ``@group`` nodeid suffix that LoadGroupScheduling uses for routing.
174+
"""
175+
for item in items:
176+
if "clean_pq_index_cache" in getattr(item, "fixturenames", ()):
177+
item.add_marker(pytest.mark.xdist_group("hf_default_cache"))
178+
179+
165180
# ==== fixtures for parquet ====
166181
@pytest.fixture
167182
def pq_data():

‎tests/streaming/test_dataloader.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -505,8 +505,10 @@ def getter(index: int):
505505

506506
@pytest.mark.skipif(sys.platform == "win32", reason="too slow")
507507
@pytest.mark.parametrize("num_workers", [1, 2])
508-
def test_dataloader_with_align_chunking(tmp_path, num_workers):
508+
def test_dataloader_with_align_chunking(tmp_path, num_workers, monkeypatch):
509509
output_dir = tmp_path / f"output_workers_{num_workers}"
510+
monkeypatch.setenv("DATA_OPTIMIZER_CACHE_FOLDER", str(tmp_path / "chunks"))
511+
monkeypatch.setenv("DATA_OPTIMIZER_DATA_CACHE_FOLDER", str(tmp_path / "data"))
510512

511513
optimize(
512514
fn=getter,

‎tests/streaming/test_dataset.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,8 +116,11 @@ def test_optimize_dataset(
116116
chunk_bytes,
117117
chunk_size,
118118
tmpdir,
119+
monkeypatch,
119120
):
120121
data_dir = str(tmpdir / "optimized")
122+
monkeypatch.setenv("DATA_OPTIMIZER_CACHE_FOLDER", str(tmpdir / "chunks"))
123+
monkeypatch.setenv("DATA_OPTIMIZER_DATA_CACHE_FOLDER", str(tmpdir / "data"))
121124

122125
optimize(
123126
fn=_simple_optimize_fn,

‎tests/streaming/test_item_loader.py‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
from litdata.constants import _NUMPY_DTYPES_MAPPING, _TORCH_DTYPES_MAPPING
88
from litdata.streaming import Cache, item_loader
99
from litdata.streaming.dataset import StreamingDataset
10-
from litdata.streaming.item_loader import PyTreeLoader, TokensLoader
10+
from litdata.streaming.item_loader import ParquetLoader, PyTreeLoader, TokensLoader
11+
from litdata.streaming.writer import index_parquet_dataset
1112

1213

1314
def test_serializer_setup():
@@ -89,3 +90,91 @@ def test_force_download(monkeypatch, tmpdir):
8990

9091
with pytest.raises(Exception, match="worked"):
9192
loader.load_item_from_chunk(0, 0, "chunk_filepath", 0, 1)
93+
94+
95+
def _write_parquet_with_row_groups(path, row_group_values):
96+
"""Write a parquet file where each element of row_group_values becomes its own row group."""
97+
import pyarrow as pa
98+
import pyarrow.parquet as pq
99+
100+
schema = pa.schema([("col", pa.int64())])
101+
with pq.ParquetWriter(path, schema) as writer:
102+
for values in row_group_values:
103+
writer.write_table(pa.table({"col": list(values)}, schema=schema))
104+
105+
106+
@pytest.mark.parametrize(
107+
"row_group_sizes",
108+
[
109+
[10, 5, 5], # regression: uneven groups, shrinking
110+
[3, 7, 2, 8], # uneven groups, varying
111+
[20], # single group
112+
[1, 1, 1, 1, 1], # many size-1 groups
113+
[5, 5, 5], # uniform control case
114+
],
115+
)
116+
@pytest.mark.parametrize("low_memory", [True, False])
117+
def test_parquet_loader_row_group_sizes(tmp_path, row_group_sizes, low_memory):
118+
"""ParquetLoader must correctly read every row regardless of row-group layout."""
119+
parquet_dir = tmp_path / "pq"
120+
parquet_dir.mkdir()
121+
122+
row_group_values = []
123+
expected = []
124+
125+
for value, size in enumerate(row_group_sizes):
126+
row_group_values.append([value] * size)
127+
expected.extend([value] * size)
128+
value += 1
129+
_write_parquet_with_row_groups(parquet_dir / "data.parquet", row_group_values)
130+
131+
index_parquet_dataset(str(parquet_dir))
132+
dataset = StreamingDataset(str(parquet_dir), item_loader=ParquetLoader(low_memory=low_memory))
133+
134+
assert len(dataset) == sum(row_group_sizes)
135+
actual = [dataset[i]["col"] for i in range(len(dataset))]
136+
assert actual == expected
137+
138+
139+
def test_parquet_loader_row_group_boundaries(tmp_path):
140+
"""First and last row of each group (the modulo edges in the old implementation)."""
141+
parquet_dir = tmp_path / "pq"
142+
parquet_dir.mkdir()
143+
144+
row_group_sizes = [10, 5, 5]
145+
_write_parquet_with_row_groups(
146+
parquet_dir / "data.parquet",
147+
[[v] * s for v, s in enumerate(row_group_sizes)],
148+
)
149+
150+
index_parquet_dataset(str(parquet_dir))
151+
dataset = StreamingDataset(str(parquet_dir), item_loader=ParquetLoader(low_memory=True))
152+
153+
boundaries = [0, 9, 10, 14, 15, 19]
154+
expected = [0, 0, 1, 1, 2, 2]
155+
for idx, exp in zip(boundaries, expected):
156+
assert dataset[idx]["col"] == exp
157+
158+
159+
def test_parquet_loader_cache_eviction_with_uneven_groups(tmp_path):
160+
"""After fully reading a row group, it must be evicted from the in-memory cache."""
161+
parquet_dir = tmp_path / "pq"
162+
parquet_dir.mkdir()
163+
164+
row_group_sizes = [10, 5, 5]
165+
_write_parquet_with_row_groups(
166+
parquet_dir / "data.parquet",
167+
[[v] * s for v, s in enumerate(row_group_sizes)],
168+
)
169+
170+
index_parquet_dataset(str(parquet_dir))
171+
loader = ParquetLoader(low_memory=True)
172+
dataset = StreamingDataset(str(parquet_dir), item_loader=loader)
173+
174+
# Iterate through the whole dataset sequentially.
175+
for i in range(len(dataset)):
176+
dataset[i]
177+
178+
# After a sequential pass every row group in the chunk should have been evicted.
179+
for chunk_index, groups in loader._chunk_row_groups.items():
180+
assert groups == {}, f"chunk {chunk_index} still holds row groups: {list(groups)}"

0 commit comments

Comments
 (0)