Fix ZstdError import handling for Python 3.14+ compatibility (Lightning-AI#767)

- Import ZstdError from compression.zstd for Python >= 3.14
- Import ZstdError from zstd for older Python versions
- Add test to verify ZstdError is properly caught for corrupted data

Fixes Lightning-AI#766

Author: bhimrazy

src/litdata/raw/indexer.py

@@ -149,24 +149,28 @@ def _load_index_file(self, index_path: str) -> Optional[list[FileMetadata]]:
         """Loads and decodes an index file."""
         if _PYTHON_GREATER_EQUAL_3_14:
             from compression import zstd
+            from compression.zstd import ZstdError
         else:
             import zstd
+            from zstd import Error as ZstdError
 
         try:
             with open(index_path, "rb") as f:
                 compressed_data = f.read()
             metadata = json.loads(zstd.decompress(compressed_data).decode("utf-8"))
             return [FileMetadata.from_dict(file_data) for file_data in metadata["files"]]
-        except (FileNotFoundError, json.JSONDecodeError, zstd.ZstdError, KeyError) as e:
+        except (FileNotFoundError, json.JSONDecodeError, ZstdError, KeyError) as e:
             logger.warning(f"Failed to load index from local cache at `{index_path}`: {e}. ")
             return None
 
     def _save_index_file(self, index_path: str, files: list[FileMetadata], source: str) -> None:
         """Encodes and saves an index file."""
         if _PYTHON_GREATER_EQUAL_3_14:
             from compression import zstd
+            from compression.zstd import ZstdError
         else:
             import zstd
+            from zstd import Error as ZstdError
 
         try:
             metadata = {
@@ -176,7 +180,7 @@ def _save_index_file(self, index_path: str, files: list[FileMetadata], source: s
             }
             with open(index_path, "wb") as f:
                 f.write(zstd.compress(json.dumps(metadata).encode("utf-8")))
-        except (OSError, zstd.ZstdError) as e:
+        except (OSError, ZstdError) as e:
             logger.warning(f"Error caching index to {index_path}: {e}")
 
     def _download_from_cloud(

tests/raw/test_indexer.py

@@ -315,3 +315,26 @@ def test_recompute_index_excludes_index_file(tmp_path):
     assert len(files) == 2
     for f in files:
         assert _INDEX_FILENAME not in f.path
+
+
+def test_load_index_file_handles_corrupted_zstd(tmp_path):
+    """Test that _load_index_file catches ZstdError for corrupted data."""
+    if _PYTHON_GREATER_EQUAL_3_14:
+        import compression.zstd as zstd
+        from compression.zstd import ZstdError
+    else:
+        import zstd
+        from zstd import Error as ZstdError
+
+    index_path = tmp_path / _INDEX_FILENAME
+
+    with open(index_path, "wb") as f:
+        f.write(b"dummy data")
+
+    # Verify corrupted data raises ZstdError
+    with pytest.raises(ZstdError):
+        zstd.decompress(b"dummy data")
+
+    # Verify the indexer catches this and returns None
+    indexer = FileIndexer()
+    assert indexer._load_index_file(str(index_path)) is None