BUG: read_sas IndexError for column text on SAS7BDAT amd pages (GH#60809)

@@ -194,6 +194,7 @@ I/O - Bug in :meth:`DataFrame.to_string`, :meth:`Series.to_string` and :meth:`DataFrame.to_html` ignoring ``na_rep`` for ``None``, ``pd.NA``, and ``NaT`` values in object-dtype columns and an explicit ``na_rep="NaN"`` for any missing value, and in :meth:`DataFrame.to_xml` raising for ``na_rep`` with nullable, pyarrow-backed, or categorical columns containing missing values (:issue:`54872`) - Bug in the repr of a :class:`Series` or :class:`DataFrame` with ``longdouble`` (e.g. ``float128``) dtype showing values outside the ``float64`` range as ``0`` or ``inf`` (:issue:`17809`) - Bug in :class:`ExcelWriter` with ``engine="openpyxl"`` ignoring ``date_format`` and ``datetime_format`` (:issue:`44284`, :issue:`59882`)+- Bug in :func:`read_sas` raising ``IndexError`` or reading the wrong format or label for SAS7BDAT files with columns renamed or relabeled after the file was written (:issue:`60809`) -  Period

@@ -334,6 +334,12 @@ def __init__(         self.column_names: list[str | bytes] = []         self.column_formats: list[str | bytes] = []         self.columns: list[_Column] = []+        # (text subheader index, offset, length) into column_names_raw, resolved+        # once all of the column text has been read+        self._column_name_refs: list[tuple[int, int, int]] = []+        self._column_format_refs: list[+            tuple[tuple[int, int, int], tuple[int, int, int]]+        ] = []          # (offset, length) of each data subheader on the current page, filled         # in by collect_page_subheaders. Sized in _get_properties, once the@@ -386,7 +392,7 @@ def _metadata_signature(self) -> tuple[int, int, int, int, bytes, int, int]:             self.column_count,             self.compression,             len(self._column_data_offsets),-            len(self.column_names),+            len(self._column_name_refs),         )      def _validate_column_data_ranges(self) -> None:@@ -614,6 +620,7 @@ def _parse_metadata(self) -> None:             if len(self._cached_page) != self._page_length:                 raise ValueError("Failed to read a meta data page from the SAS file.")             done = self._process_page_meta()+        self._resolve_column_text()          self._column_convert_types: list[str | None] = []         for j in range(self.column_count):@@ -807,10 +814,7 @@ def _process_columnname_subheader(self, offset: int, length: int) -> None:                 col_name_offset, const.column_name_offset_length             )             col_len = self._read_uint(col_name_length, const.column_name_length_length)--            name_raw = self.column_names_raw[idx]-            cname = name_raw[col_offset : col_offset + col_len]-            self.column_names.append(self._convert_header_text(cname))+            self._column_name_refs.append((idx, col_offset, col_len))      def _process_columnattributes_subheader(self, offset: int, length: int) -> None:         int_len = self._int_length@@ -855,11 +859,9 @@ def _process_format_subheader(self, offset: int, length: int) -> None:         col_label_offset = offset + const.column_label_offset_offset + 3 * int_len         col_label_len = offset + const.column_label_length_offset + 3 * int_len -        x = self._read_uint(+        format_idx = self._read_uint(             text_subheader_format, const.column_format_text_subheader_index_length         )-        format_idx = min(x, len(self.column_names_raw) - 1)-         format_start = self._read_uint(             col_format_offset, const.column_format_offset_length         )@@ -868,34 +870,127 @@ def _process_format_subheader(self, offset: int, length: int) -> None:         label_idx = self._read_uint(             text_subheader_label, const.column_label_text_subheader_index_length         )-        label_idx = min(label_idx, len(self.column_names_raw) - 1)-         label_start = self._read_uint(             col_label_offset, const.column_label_offset_length         )         label_len = self._read_uint(col_label_len, const.column_label_length_length) -        label_names = self.column_names_raw[label_idx]-        column_label = self._convert_header_text(-            label_names[label_start : label_start + label_len]-        )-        format_names = self.column_names_raw[format_idx]-        column_format = self._convert_header_text(-            format_names[format_start : format_start + format_len]-        )-        current_column_number = len(self.columns)--        col = _Column(-            current_column_number,-            self.column_names[current_column_number],-            column_label,-            column_format,-            self._column_types[current_column_number],-            self._column_data_lengths[current_column_number],+        self._column_format_refs.append(+            (+                (format_idx, format_start, format_len),+                (label_idx, label_start, label_len),+            )         ) -        self.column_formats.append(column_format)-        self.columns.append(col)+    def _resolve_column_text(self) -> None:+        # Refs index the column text subheaders in file order, and an amd page+        #  after the data can hold some of that text. GH#60809+        n_text = len(self.column_names_raw)+        if any(idx >= n_text for idx, _, _ in self._column_name_refs) or any(+            idx >= n_text and length > 0+            for refs in self._column_format_refs+            for idx, _, length in refs+        ):+            self._read_amd_column_text()+            n_text = len(self.column_names_raw)++        for idx, start, length in self._column_name_refs:+            if idx >= n_text:+                raise ValueError(+                    f"Column name refers to text subheader {idx}, which the file "+                    "does not contain; the file is corrupt"+                )+            name_raw = self.column_names_raw[idx]+            self.column_names.append(+                self._convert_header_text(name_raw[start : start + length])+            )++        for format_ref, label_ref in self._column_format_refs:+            format_idx, format_start, format_len = format_ref+            label_idx, label_start, label_len = label_ref+            label_names = self.column_names_raw[min(label_idx, n_text - 1)]+            column_label = self._convert_header_text(+                label_names[label_start : label_start + label_len]+            )+            format_names = self.column_names_raw[min(format_idx, n_text - 1)]+            column_format = self._convert_header_text(+                format_names[format_start : format_start + format_len]+            )+            current_column_number = len(self.columns)++            col = _Column(+                current_column_number,+                self.column_names[current_column_number],+                column_label,+                column_format,+                self._column_types[current_column_number],+                self._column_data_lengths[current_column_number],+            )++            self.column_formats.append(column_format)+            self.columns.append(col)++    def _read_amd_column_text(self) -> None:+        # Scan the pages after the metadata for amd pages and collect only their+        #  column text, as ReadStat does. Unusable pages and subheaders are skipped,+        #  leaving the ref unresolved. GH#60809+        if not self._path_or_buf.seekable():+            # e.g. zstd, which cannot seek back to the data after the scan+            self._path_or_buf = io.BytesIO(self._path_or_buf.read())+        cached_page = self._cached_page+        file_pos = self._path_or_buf.tell()+        bit_offset = self._page_bit_offset+        int_len = self._int_length+        page_length = self._page_length+        pointer_length = self._subheader_pointer_length+        head_len = bit_offset + const.subheader_pointers_offset+        page_start = file_pos+        while True:+            self._path_or_buf.seek(page_start)+            self._cached_page = self._path_or_buf.read(head_len)+            if len(self._cached_page) < head_len:+                break+            page_type = (+                self._read_uint(+                    bit_offset + const.page_type_offset, const.page_type_length+                )+                & const.page_type_mask2+            )+            if page_type == const.page_amd_type:+                self._cached_page += self._path_or_buf.read(page_length - head_len)+                if len(self._cached_page) < page_length:+                    break+                subheader_count = self._read_uint(+                    bit_offset + const.subheader_count_offset,+                    const.subheader_count_length,+                )+                if head_len + subheader_count * pointer_length > page_length:+                    subheader_count = 0+                for i in range(subheader_count):+                    pointer = head_len + i * pointer_length+                    offset = self._read_uint(pointer, int_len)+                    length = self._read_uint(pointer + int_len, int_len)+                    compression = self._read_uint(pointer + 2 * int_len, 1)+                    if (+                        length < int_len + const.text_block_size_length+                        or offset + length > page_length+                        or compression == const.truncated_subheader_id+                    ):+                        continue+                    signature = self._read_bytes(offset, int_len)+                    if (+                        const.subheader_signature_to_index.get(signature)+                        != const.SASIndex.column_text_index+                    ):+                        continue+                    text_block_size = self._read_uint(+                        offset + int_len, const.text_block_size_length+                    )+                    if offset + int_len + text_block_size <= page_length:+                        self._process_columntext_subheader(offset, length)+            page_start += page_length+        self._cached_page = cached_page+        self._path_or_buf.seek(file_pos)      # Same order as const.SASIndex. Plain functions, not bound methods: those     # would make the reader a reference cycle, so its file would stay open until

GitHub publishes no diff for this file. View it on GitHub ↗

GitHub publishes no diff for this file. View it on GitHub ↗

@@ -1343,3 +1343,129 @@ def test_utf8_translation_table_rejects_context_dependent(encoding):     # once a continuation byte follows; raw_unicode_escape instead maps every     # byte one to one yet still reads b"\\u0041" as "A".     assert _utf8_translation_table(encoding) is None+++def test_amd_page(datapath):+    # GH#60809 a column renamed after the file was written keeps its new name on+    #  an amd page after the data; SAS rewriting the file moves it inline+    df = pd.read_sas(+        datapath("io", "sas", "data", "amd_page.sas7bdat.bz2"), encoding="infer"+    )+    expected = pd.read_sas(+        datapath("io", "sas", "data", "amd_page_resaved.sas7bdat.bz2"),+        encoding="infer",+    )+    assert df.columns[11] == "TotalNoise"+    tm.assert_frame_equal(df, expected)+++def _with_amd_page(datapath, fname, field_offset, text):+    # fname with an amd page appended holding one column text subheader, and the+    #  ref at field_offset (from the end of the header) pointed into it+    path = datapath("io", "sas", "data", fname)+    with SAS7BDATReader(path) as reader:+        header_length = reader.header_length+        page_length = reader._page_length+        bit_offset = reader._page_bit_offset+        int_len = reader._int_length+    with open(path, "rb") as fd:+        data = bytearray(fd.read())+    struct.pack_into("<HHH", data, header_length + field_offset, 1, 8, len(text))++    amd_page = bytearray(page_length)+    struct.pack_into("<HHH", amd_page, bit_offset, const.page_amd_type, 1, 1)+    block = struct.pack("<H", 8 + len(text)) + bytes(6) + text+    subheader = b"\xfd" + b"\xff" * (int_len - 1) + block+    position = page_length // 2+    amd_page[position : position + len(subheader)] = subheader+    struct.pack_into(+        "<qqBB" if int_len == 8 else "<iiBB",+        amd_page,+        bit_offset + const.subheader_pointers_offset,+        position,+        len(subheader),+        0,+        1,+    )+    return bytes(data) + bytes(amd_page)+++@pytest.mark.parametrize(+    "field_offset, text, attr, expected",+    [+        # datetimecol's name ref, then its format and label refs+        (63960 + 8 + 16, b"renamed", "name", "renamed"),+        (63718 + 22 + 24, b"DATETIME", "format", "DATETIME"),+        (63718 + 28 + 24, b"new label", "label", "new label"),+    ],+)+def test_amd_page_column_text(datapath, field_offset, text, attr, expected):+    # GH#60809 column text on an amd page was not found or read from the wrong+    #  text; a wrong format turned a datetime column into float+    data = _with_amd_page(datapath, "dates_null.sas7bdat", field_offset, text)+    with SAS7BDATReader(io.BytesIO(data), encoding="infer") as reader:+        assert getattr(reader.columns[1], attr) == expected+        df = reader.read()+    assert df["datecol"].dtype.kind == "M"+    assert df.iloc[:, 1].dtype.kind == "M"+++def test_column_name_text_missing_raises(datapath):+    # GH#60809 a name ref past the column text with no amd page to supply it+    data = _with_amd_page(datapath, "dates_null.sas7bdat", 63960 + 8 + 16, b"renamed")+    with pytest.raises(ValueError, match="refers to text subheader 1"):+        pd.read_sas(io.BytesIO(data[:-65536]), format="sas7bdat", encoding="infer")+++def test_amd_page_rows_after_scan(datapath):+    # GH#60809 cars has rows on the pages after its metadata, so they are read+    #  after the amd scan has moved through the file+    data = _with_amd_page(datapath, "cars.sas7bdat", 3060 + 4 + 32, b"renamed")+    expected = pd.read_sas(datapath("io", "sas", "data", "cars.sas7bdat"))+    expected = expected.rename(columns={"WGT": "renamed"})++    result = pd.read_sas(io.BytesIO(data), format="sas7bdat")+    tm.assert_frame_equal(result, expected)++    with pd.read_sas(io.BytesIO(data), format="sas7bdat", chunksize=100) as reader:+        result = pd.concat(reader)+    tm.assert_frame_equal(result, expected)+++def test_amd_page_zstd(datapath, tmp_path):+    # GH#60809 a zstd stream cannot seek back to the rows after the amd scan+    zstandard = pytest.importorskip("zstandard")+    data = _with_amd_page(datapath, "cars.sas7bdat", 3060 + 4 + 32, b"renamed")+    path = tmp_path / "cars.sas7bdat.zst"+    path.write_bytes(zstandard.ZstdCompressor().compress(data))+    result = pd.read_sas(path)+    expected = pd.read_sas(io.BytesIO(data), format="sas7bdat")+    tm.assert_frame_equal(result, expected)+++@pytest.mark.parametrize("corruption", ["truncated", "pointer past page"])+def test_amd_page_unusable(datapath, corruption):+    # GH#60809 an unusable amd page is skipped: a format ref past the column+    #  text falls back to clamping, as before, and a name ref raises+    def corrupt(data):+        if corruption == "truncated":+            return data[: -65536 // 2]+        data = bytearray(data)+        pointer = (+            len(data)+            - 65536+            + const.page_bit_offset_x64+            + const.subheader_pointers_offset+        )+        struct.pack_into("<q", data, pointer, 65536)+        return bytes(data)++    fname = "dates_null.sas7bdat"+    data = corrupt(_with_amd_page(datapath, fname, 63718 + 22 + 24, b"DATETIME"))+    with SAS7BDATReader(io.BytesIO(data), encoding="infer") as reader:+        # the bytes at the ref's offset in the text block that is present+        assert reader.columns[1].format == "\x00\x00\x00\x00    "++    data = corrupt(_with_amd_page(datapath, fname, 63960 + 8 + 16, b"renamed"))+    with pytest.raises(ValueError, match="refers to text subheader 1"):+        pd.read_sas(io.BytesIO(data), format="sas7bdat", encoding="infer")