Keep columns with the same name, and key CSV column types by the reader's column labels

pyathena/arrow/result_set.py

@@ -12,7 +12,7 @@
 
 from pyathena import OperationalError
 from pyathena.arrow.util import to_column_info
-from pyathena.converter import Converter, _text_value_converter, _to_default
+from pyathena.converter import _TEXT_VALUE_TYPES, Converter, _text_value_converter, _to_default
 from pyathena.error import ProgrammingError
 from pyathena.model import AthenaQueryExecution
 from pyathena.result_set import AthenaResultSet
@@ -254,23 +254,23 @@ def _fetch(self) -> None:
         except StopIteration:
             return
         else:
-            dict_rows = rows.to_pydict()
-            converters = (
-                self.converters
-                if self._convert_rows
-                else self._text_value_converters(self.converters)
-            )
-            if converters:
-                column_names = dict_rows.keys()
+            # Read the columns and their converters by position; to_pydict() and the
+            # converters property keep one column per name.
+            columns = [column.to_pylist() for column in rows.columns]
+            description = self.description if self.description else []
+            converters = [
+                self._converter.get(d[1])
+                if self._convert_rows or d[1] in _TEXT_VALUE_TYPES
+                else _to_default
+                for d in description
+            ]
+            if any(convert is not _to_default for convert in converters):
                 processed_rows = [
-                    tuple(
-                        converters.get(k, _to_default)(v)
-                        for k, v in zip(column_names, row, strict=False)
-                    )
-                    for row in zip(*dict_rows.values(), strict=False)
+                    tuple(convert(v) for convert, v in zip(converters, row, strict=False))
+                    for row in zip(*columns, strict=False)
                 ]
             else:
-                processed_rows = list(zip(*dict_rows.values(), strict=False))
+                processed_rows = list(zip(*columns, strict=False))
             self._rows.extend(processed_rows)
 
     @override
@@ -403,8 +403,8 @@ def _as_arrow_from_api(self, converter: Converter | None = None) -> Table:
         if not rows:
             return pa.Table.from_pydict({})
         description = self.description if self.description else []
-        columns = [d[0] for d in description]
-        return pa.table(self._rows_to_columnar(rows, columns))
+        columns = [list(column) for column in zip(*rows, strict=True)]
+        return pa.table(columns, names=[d[0] for d in description])
 
     def as_arrow(self) -> Table:
         """Return the query results as an Apache Arrow Table.
@@ -447,4 +447,4 @@ def close(self) -> None:
 
         super().close()
         self._table = pa.Table.from_pydict({})
-        self._batches = []
+        self._batches = iter([])

pyathena/pandas/result_set.py

@@ -411,9 +411,10 @@ def __init__(
         # The converters that pandas.read_csv() applies, keyed by column name.
         self._csv_converters: dict[Any, Callable[[str | None], Any]] = {}
 
-        # Cache time column names for efficient _trunc_date processing
+        # Cache time column names for efficient _trunc_date processing. _read_csv()
+        # replaces them with the labels of the columns that it reads.
         description = self.description if self.description else []
-        self._time_columns: list[str] = [d[0] for d in description if d[1] == "time"]
+        self._time_columns: list[Any] = [d[0] for d in description if d[1] == "time"]
 
         import pandas as pd
 
@@ -475,6 +476,12 @@ def _get_csv_engine(
                 effective_chunksize is None
                 and self._quoting == 1
                 and not self.converters
+                # The pyarrow engine does not rename columns with the same name, and
+                # the column labels cannot be resolved for it (see
+                # _get_csv_column_labels()).
+                and not self._needs_csv_column_name_resolution(
+                    [d[0] for d in self.description or []]
+                )
                 and (file_size_bytes is None or file_size_bytes >= self.PYARROW_MIN_FILE_SIZE_BYTES)
             )
             if is_compatible:
@@ -588,6 +595,23 @@ def parse_dates(self) -> list[Any | None]:
         description = self.description if self.description else []
         return [d[0] for d in description if d[1] in self._PARSE_DATES]
 
+    def _get_column_names(self) -> list[Any]:
+        """Get the names of the result columns in the DataFrame.
+
+        Columns with the same name are renamed as pandas renames them when it reads
+        the header of a CSV file, such as ``x`` and ``x.1``.
+
+        Returns:
+            List of column names.
+        """
+        import pandas as pd
+
+        description = self.description if self.description else []
+        names = [d[0] for d in description]
+        if len(set(names)) == len(names):
+            return names
+        return self._resolve_csv_column_names(names, {}, pd.read_csv)[0]
+
     def _finish_csv_frame(self, df: DataFrame) -> DataFrame:
         """Finish a DataFrame read from the CSV result file.
 
@@ -641,8 +665,8 @@ def fetchone(
             return None
         else:
             self._rownumber = row[0] + 1
-            description = self.description if self.description else []
-            return tuple([row[1][d[0]] for d in description])
+            # By position, so that columns with the same name keep their own values.
+            return tuple(row[1].values())
 
     def _read_csv(self) -> TextFileReader | DataFrame:
         import pandas as pd
@@ -676,11 +700,14 @@ def _read_csv(self) -> TextFileReader | DataFrame:
 
         csv_engine = self._get_csv_engine(length, effective_chunksize)
         read_csv_kwargs = self._get_csv_read_options(csv_engine, effective_chunksize)
+        labels = self._get_csv_column_labels(csv_engine, read_csv_kwargs)
+        if labels is not None:
+            self._key_csv_columns_by_labels(read_csv_kwargs, labels)
 
         try:
             with ExitStack() as stack:
                 source: str | IOBase = self.output_location
-                binary_columns = self._configure_binary_csv_read(read_csv_kwargs, pd.read_csv)
+                binary_columns = self._configure_binary_csv_read(read_csv_kwargs, labels)
                 self._csv_converters = read_csv_kwargs.get("converters") or {}
                 if binary_columns:
                     # Given storage_options, even None, open the file through fsspec
@@ -721,7 +748,7 @@ def _get_csv_read_options(self, csv_engine: str, chunksize: int | None) -> dict[
         if self.output_location and self.output_location.endswith(".txt"):
             sep = "\t"
             header = None
-            names = [d[0] for d in self.description or []]
+            names = self._get_column_names()
         else:
             sep = ","
             header = 0
@@ -796,12 +823,17 @@ def _resolve_csv_column_names(
             )
         return column_names, selected_names
 
-    def _can_preserve_binary_csv_nulls(self, read_csv_kwargs: dict[str, Any]) -> bool:
-        """Whether CSV settings support distinguishing binary NULL from empty values."""
+    def _is_standard_csv_parsing(self, read_csv_kwargs: dict[str, Any]) -> bool:
+        """Whether the options read a CSV result file as Athena writes it.
+
+        Args:
+            read_csv_kwargs: The options for ``pandas.read_csv()``.
+
+        Returns:
+            True if pandas reads the header and quoted fields of the file as written.
+        """
         return not (
-            "varbinary" not in self._converter.mappings
-            or "converters" in self._kwargs
-            or not self.output_location
+            not self.output_location
             or not self.output_location.endswith(".csv")
             or read_csv_kwargs.get("header") != 0
             or read_csv_kwargs.get("skiprows") is not None
@@ -810,6 +842,14 @@ def _can_preserve_binary_csv_nulls(self, read_csv_kwargs: dict[str, Any]) -> boo
             or read_csv_kwargs.get("quotechar", '"') != '"'
         )
 
+    def _can_preserve_binary_csv_nulls(self, read_csv_kwargs: dict[str, Any]) -> bool:
+        """Whether CSV settings support distinguishing binary NULL from empty values."""
+        return (
+            "varbinary" in self._converter.mappings
+            and "converters" not in self._kwargs
+            and self._is_standard_csv_parsing(read_csv_kwargs)
+        )
+
     def _needs_csv_column_name_resolution(self, column_names: list[Any]) -> bool:
         """Whether pandas must resolve column names instead of using Athena metadata."""
         return (
@@ -829,38 +869,98 @@ def _needs_csv_column_name_resolution(self, column_names: list[Any]) -> bool:
             )
         )
 
+    def _get_csv_column_labels(
+        self, csv_engine: str, read_csv_kwargs: dict[str, Any]
+    ) -> list[Any] | None:
+        """Get the labels that pandas gives the result columns when it reads the CSV file.
+
+        Args:
+            csv_engine: The CSV engine that reads the file.
+            read_csv_kwargs: The options for ``pandas.read_csv()``.
+
+        Returns:
+            The label of each result column in the description order, with None for a
+            column that ``usecols`` leaves out. None if the options do not read the file
+            as Athena writes it, or if the labels do not match the result columns one
+            to one, such as with fewer ``names``.
+        """
+        import pandas as pd
+
+        # The pyarrow engine runs only when no labels need resolving (see
+        # _get_csv_engine()), and does not support reading only the header.
+        if csv_engine == "pyarrow" or not self._is_standard_csv_parsing(read_csv_kwargs):
+            return None
+        column_names = [d[0] for d in self.description or []]
+        if not self._needs_csv_column_name_resolution(column_names):
+            return column_names
+        labels, selected_labels = self._resolve_csv_column_names(
+            column_names, read_csv_kwargs, pd.read_csv
+        )
+        if len(labels) != len(column_names):
+            return None
+        return [label if label in selected_labels else None for label in labels]
+
+    def _key_csv_columns_by_labels(
+        self, read_csv_kwargs: dict[str, Any], labels: list[Any]
+    ) -> None:
+        """Key the column options of ``pandas.read_csv()`` by the labels of the columns.
+
+        Columns with the same name keep their own converters and date parsing, and
+        options that rename or select columns get the types of the columns they
+        read. The ``dtype``, ``converters``, and ``parse_dates`` given to
+        ``execute()`` are kept as they are.
+
+        Args:
+            read_csv_kwargs: The options for ``pandas.read_csv()``, updated in place.
+            labels: The labels from ``_get_csv_column_labels()``.
+        """
+        description = self.description or []
+        columns = [
+            (label, d) for label, d in zip(labels, description, strict=True) if label is not None
+        ]
+        if "dtype" not in self._kwargs:
+            read_csv_kwargs["dtype"] = {
+                label: dtype
+                for label, d in columns
+                if (dtype := self._converter.get_dtype(d[1], d[4], d[5])) is not None
+            }
+        if "converters" not in self._kwargs:
+            read_csv_kwargs["converters"] = {
+                label: self._get_csv_converter(d[1])
+                for label, d in columns
+                if d[1] in self._converter.mappings
+            }
+        if "parse_dates" not in self._kwargs:
+            read_csv_kwargs["parse_dates"] = [
+                label for label, d in columns if d[1] in self._PARSE_DATES
+            ]
+        self._time_columns = [label for label, d in columns if d[1] == "time"]
+
     def _configure_binary_csv_read(
-        self, read_csv_kwargs: dict[str, Any], read_csv: Callable[..., DataFrame]
+        self, read_csv_kwargs: dict[str, Any], labels: list[Any] | None
     ) -> set[int]:
-        """Wrap binary converters and return column positions needing NULL preservation."""
-        if not self._can_preserve_binary_csv_nulls(read_csv_kwargs):
-            return set()
+        """Wrap binary converters and return column positions needing NULL preservation.
 
-        description = self.description or []
-        binary_columns = {i for i, d in enumerate(description) if d[1] == "varbinary"}
-        if not binary_columns:
-            return set()
+        Args:
+            read_csv_kwargs: The options for ``pandas.read_csv()``, whose converters
+                are keyed by the column labels.
+            labels: The labels from ``_get_csv_column_labels()``.
 
-        column_names = [d[0] for d in description]
-        converters = read_csv_kwargs["converters"]
-        if self._needs_csv_column_name_resolution(column_names):
-            column_names, selected_names = self._resolve_csv_column_names(
-                column_names, read_csv_kwargs, read_csv
-            )
-            if len(column_names) != len(description):
-                return set()
-            converters = {
-                name: self._get_csv_converter(d[1])
-                for name, d in zip(column_names, description, strict=True)
-                if d[1] in self._converter.mappings and name in selected_names
-            }
-            binary_columns = {i for i in binary_columns if column_names[i] in selected_names}
+        Returns:
+            The positions of the binary columns whose NULL fields the stream preserves.
+        """
+        if labels is None or not self._can_preserve_binary_csv_nulls(read_csv_kwargs):
+            return set()
 
+        description = self.description or []
+        binary_columns = {
+            i for i, d in enumerate(description) if d[1] == "varbinary" and labels[i] is not None
+        }
         if binary_columns:
+            converters = read_csv_kwargs["converters"]
             for index in binary_columns:
-                name = column_names[index]
-                converters[name] = partial(_convert_binary_csv, converters[name])
-            read_csv_kwargs["converters"] = converters
+                label = labels[index]
+                converters[label] = partial(_convert_binary_csv, converters[label])
         return binary_columns
 
     def _open_binary_csv_stream(
@@ -957,24 +1057,20 @@ def _as_pandas_from_api(self, converter: Converter | None = None) -> DataFrame:
         if not rows:
             return pd.DataFrame()
         description = self.description if self.description else []
-        columns = [d[0] for d in description]
-        columnar = self._rows_to_columnar(rows, columns)
+        # Positional, so that columns with the same name keep their own values.
+        columns = [list(column) for column in zip(*rows, strict=True)]
         # Integer columns get the dtype that the CSV result file reads them with,
         # and json columns with NULL stay objects as there, so that NULL does not
         # make their values floats.
-        dtypes: dict[str, Any] = {}
-        for d in description:
+        data: dict[Any, Any] = {}
+        for name, values, d in zip(self._get_column_names(), columns, description, strict=True):
+            dtype = None
             if d[1] in self._INTEGER_TYPES:
-                if (dtype := self._converter.get_dtype(d[1], d[4], d[5])) is not None:
-                    dtypes[d[0]] = dtype
-            elif d[1] == "json" and None in columnar[d[0]]:
-                dtypes[d[0]] = object
-        return pd.DataFrame(
-            {
-                name: values if name not in dtypes else pd.array(values, dtype=dtypes[name])
-                for name, values in columnar.items()
-            }
-        )
+                dtype = self._converter.get_dtype(d[1], d[4], d[5])
+            elif d[1] == "json" and None in values:
+                dtype = object
+            data[name] = values if dtype is None else pd.array(values, dtype=dtype)
+        return pd.DataFrame(data)
 
     def as_pandas(self) -> PandasDataFrameIterator | DataFrame:
         """Return the query results as a DataFrame or an iterator of DataFrame chunks.