add more tests

Author: ramanishsingh

test/nodes/test_csv_dataloader.py

@@ -7,11 +7,9 @@
 import csv
 import os
 import tempfile
-from typing import Any, Dict, List, Union
 
 from parameterized import parameterized
 from torch.testing._internal.common_utils import TestCase
-from torchdata.nodes.base_node import BaseNode
 
 from torchdata.nodes.csv_reader import CSVReader
 
@@ -61,6 +59,7 @@ def test_basic_read_dict(self):
         self.assertEqual(len(results), len(self.test_data) - 1)
         self.assertEqual(results[0], {"name": "Alice", "age": "30", "city": "New York"})
         self.assertEqual(results[1]["city"], "London")
+        self.assertEqual(results[-1]["city"], "Bogota")
         node.close()
 
     def test_different_delimiters(self):
@@ -70,25 +69,27 @@ def test_different_delimiters(self):
 
         self.assertEqual(len(results), len(self.test_data) - 1)
         self.assertEqual(results[2]["city"], "Paris")
+        self.assertEqual(results[-1]["city"], "Bogota")
         node.close()
 
     def test_state_management(self):
         path = self._create_temp_csv()
         node = CSVReader(path, has_header=True, return_dict=True)
-
+        print(f"initial state: {node.state_dict()}")
         for _ in range(11):
-            next(node)
+            _ = next(node)
+            print(f"element = {_}, state: {node.state_dict()}")
 
         state = node.state_dict()
 
         node.reset(state)
-
         item = next(node)
+
         with self.assertRaises(StopIteration):
             next(node)
 
         self.assertEqual(item["name"], "Lily")
-        self.assertEqual(state[CSVReader.LINE_NUM_KEY], 11)
+        self.assertEqual(state[CSVReader.NUM_LINES_YIELDED], 11)
         node.close()
 
     @parameterized.expand([3, 5, 7])
@@ -98,6 +99,28 @@ def test_save_load_state(self, midpoint: int):
         run_test_save_load_state(self, node, midpoint)
         node.close()
 
+    def test_load_wrong_state(self):
+        path = self._create_temp_csv(header=True)
+        node = CSVReader(path, has_header=True)
+
+        state = node.state_dict()
+        state[CSVReader.HEADER_KEY] = None
+        with self.assertRaisesRegex(ValueError, "Check if has_header=True matches the state header=None"):
+            node.reset(state)
+
+        node.close()
+
+        node = CSVReader(path, has_header=False)
+        state = node.state_dict()
+        state[CSVReader.HEADER_KEY] = ["name", "age"]
+        with self.assertRaisesRegex(
+            ValueError,
+            r"Check if has_header=False matches the state header=\['name', 'age'\]",
+        ):
+            node.reset(state)
+
+        node.close()
+
     def test_empty_file(self):
         path = self._create_temp_csv()
         # Overwrite with empty file

torchdata/nodes/csv_reader.py

@@ -1,4 +1,5 @@
 import csv
+from itertools import islice
 from typing import Any, Dict, Iterator, List, Optional, Sequence, TextIO, Union
 
 from torchdata.nodes.base_node import BaseNode
@@ -13,7 +14,7 @@ class CSVReader(BaseNode[Union[List[str], Dict[str, str]]]):
         return_dict: Return rows as dictionaries (requires has_header=True)
     """
 
-    LINE_NUM_KEY = "line_num"
+    NUM_LINES_YIELDED = "num_lines_yielded"
     HEADER_KEY = "header"
 
     def __init__(
@@ -22,6 +23,7 @@ def __init__(
         has_header: bool = False,
         delimiter: str = ",",
         return_dict: bool = False,
+        encoding: str = "utf-8",
     ):
         super().__init__()
         self.file_path = file_path
@@ -30,64 +32,84 @@ def __init__(
         self.return_dict = return_dict
         if return_dict and not has_header:
             raise ValueError("return_dict=True requires has_header=True")
+        self.encoding = encoding
         self._file: Optional[TextIO] = None
         self._reader: Optional[Iterator[Union[List[str], Dict[str, str]]]] = None
         self._header: Optional[Sequence[str]] = None
-        self._line_num: int = 0
+        self._num_lines_yielded: int = 0
         self.reset()  # Initialize reader
 
     def reset(self, initial_state: Optional[Dict[str, Any]] = None):
-        super().reset(initial_state)
-
-        if self._file and not self._file.closed:
-            self._file.close()
+        super().reset()
+        self.close()
+
+        # Reopen the file and reset counters
+        self._file = open(self.file_path, encoding=self.encoding)
+        self._num_lines_yielded = 0
+        if initial_state is not None:
+            self._handle_initial_state(initial_state)
+        else:
+            self._initialize_reader()
 
-        self._file = open(self.file_path, newline="", encoding="utf-8")
-        self._line_num = 0
+    def _handle_initial_state(self, state: Dict[str, Any]):
+        """Restore reader state from checkpoint."""
+        # Validate header compatibility
+        if (not self.has_header and self.HEADER_KEY in state) or (self.has_header and state[self.HEADER_KEY] is None):
+            raise ValueError(f"Check if has_header={self.has_header} matches the state header={state[self.HEADER_KEY]}")
 
-        if initial_state:
-            self._header = initial_state.get(self.HEADER_KEY)
-            target_line_num = initial_state[self.LINE_NUM_KEY]
+        self._header = state.get(self.HEADER_KEY)
+        target_line_num = state[self.NUM_LINES_YIELDED]
+        assert self._file is not None
+        # Create appropriate reader
+        if self.return_dict:
 
-            if self.return_dict:
-                if self._header is None:
-                    raise ValueError("return_dict=True requires has_header=True")
-                self._reader = csv.DictReader(self._file, delimiter=self.delimiter, fieldnames=self._header)
-            else:
-                self._reader = csv.reader(self._file, delimiter=self.delimiter)
+            self._reader = csv.DictReader(self._file, delimiter=self.delimiter, fieldnames=self._header)
+        else:
+            self._reader = csv.reader(self._file, delimiter=self.delimiter)
+        # Skip header if needed (applies only when file has header)
+
+        assert isinstance(self._reader, Iterator)
+        if self.has_header:
+            try:
+                next(self._reader)  # Skip header line
+            except StopIteration:
+                pass  # Empty file
+        # Fast-forward to target line using efficient slicing
+        consumed = sum(1 for _ in islice(self._reader, target_line_num))
+        self._num_lines_yielded = consumed
+
+    def _initialize_reader(self):
+        """Create fresh reader without state."""
+        assert self._file is not None
+        if self.return_dict:
+            self._reader = csv.DictReader(self._file, delimiter=self.delimiter)
+            self._header = self._reader.fieldnames
+        else:
+            self._reader = csv.reader(self._file, delimiter=self.delimiter)
 
-            assert isinstance(self._reader, Iterator)
             if self.has_header:
-                next(self._reader)  # Skip header
-            for _ in range(target_line_num - self._line_num):
-                try:
-                    next(self._reader)
-                    self._line_num += 1
-                except StopIteration:
-                    break
-        else:
 
-            if self.return_dict:
-                self._reader = csv.DictReader(self._file, delimiter=self.delimiter)
-                self._header = self._reader.fieldnames
-            else:
-                self._reader = csv.reader(self._file, delimiter=self.delimiter)
-                if self.has_header:
+                try:
                     self._header = next(self._reader)
+                except StopIteration:
+                    self._header = None  # Handle empty file
 
     def next(self) -> Union[List[str], Dict[str, str]]:
         try:
             assert isinstance(self._reader, Iterator)
             row = next(self._reader)
-            self._line_num += 1
+            self._num_lines_yielded += 1
             return row
 
         except StopIteration:
             self.close()
             raise
 
     def get_state(self) -> Dict[str, Any]:
-        return {self.LINE_NUM_KEY: self._line_num, self.HEADER_KEY: self._header}
+        return {
+            self.NUM_LINES_YIELDED: self._num_lines_yielded,
+            self.HEADER_KEY: self._header,
+        }
 
     def close(self):
         if self._file and not self._file.closed: