Merge pull request #770 from pyathena-dev/fix/759-binary-compliance

Restore SQLAlchemy binary compliance and preserve CSV binary NULLs

Author: laughingman7743

docs/null_handling.md

@@ -60,6 +60,17 @@ which correctly interprets unquoted empty values as NULL, while `S3FSCursor` use
 `AthenaCSVReader` that respects CSV quoting rules.
 ```
 
+## Binary Values
+
+The string comparison above does not apply to `VARBINARY` columns.
+With the default CSV settings and converters, pandas and Arrow cursors distinguish SQL NULL from empty binary
+values when reading CSV results: `fetchone()`, `fetchmany()`, and `fetchall()` return `None`
+for NULL and `b''` for an empty binary value.
+This also applies to their asynchronous variants and pandas chunked reads.
+Pandas DataFrames preserve the same values.
+Arrow Tables retain the CSV hexadecimal strings, with NULL represented as an Arrow null;
+fetch methods convert the hexadecimal strings to Python bytes.
+
 ## Default Cursor (API-based)
 
 The default `Cursor` and `DictCursor` fetch results directly from the Athena API,

pyathena/__init__.py

@@ -63,6 +63,7 @@ def __hash__(self):
 Date: type[datetime.date] = datetime.date
 Time: type[datetime.time] = datetime.time
 Timestamp: type[datetime.datetime] = datetime.datetime
+Binary: type[bytes] = bytes
 
 
 @overload

pyathena/aio/sqlalchemy/base.py

@@ -171,6 +171,8 @@ class AsyncAdapt_pyathena_dbapi:
     """
 
     paramstyle = "pyformat"
+    Binary = pyathena.Binary
+    BINARY = pyathena.BINARY
 
     # DBAPI exception hierarchy
     Error = Error

pyathena/arrow/result_set.py

@@ -276,6 +276,7 @@ def _read_csv(self) -> Table:
         ):
             return pa.Table.from_pydict({})
         length = self._get_content_length()
+        binary_columns = {d[0] for d in self.description or [] if d[1] == "varbinary"}
         if length and self.output_location.endswith(".txt"):
             description = self.description if self.description else []
             column_names = [d[0] for d in description]
@@ -296,6 +297,7 @@ def _read_csv(self) -> Table:
             parse_opts = csv.ParseOptions(
                 delimiter=",",
                 quote_char='"',
+                ignore_empty_lines=not binary_columns,
                 double_quote=True,
                 escape_char=False,
             )
@@ -304,16 +306,25 @@ def _read_csv(self) -> Table:
 
         bucket, key = parse_output_location(self.output_location)
         try:
-            return csv.read_csv(
+            table = csv.read_csv(
                 self._fs.open_input_stream(f"{bucket}/{key}"),
                 read_options=read_opts,
                 parse_options=parse_opts,
                 convert_options=csv.ConvertOptions(
+                    strings_can_be_null=bool(binary_columns),
                     quoted_strings_can_be_null=False,
                     timestamp_parsers=self.timestamp_parsers,
                     column_types=self.column_types,
                 ),
             )
+            if binary_columns:
+                for index, field in enumerate(table.schema):
+                    if field.name not in binary_columns and (
+                        pa.types.is_string(field.type) or pa.types.is_binary(field.type)
+                    ):
+                        # Preserve the existing CSV behavior for non-binary Athena columns.
+                        table = table.set_column(index, field, table.column(index).fill_null(""))
+            return table
         except Exception as e:
             _logger.exception(f"Failed to read {bucket}/{key}.")
             raise OperationalError(*e.args) from e

pyathena/formatter.py

@@ -251,6 +251,12 @@ def _format_str(formatter: Formatter, escaper: Callable[[str], str], val: Any) -
     return escaper(val)
 
 
+def _format_binary(
+    formatter: Formatter, escaper: Callable[[str], str], val: bytes | bytearray | memoryview
+) -> str:
+    return f"X'{val.hex()}'"
+
+
 def _format_seq(formatter: Formatter, escaper: Callable[[str], str], val: Any) -> Any:
     results = []
     for v in val:
@@ -291,6 +297,9 @@ def _format_decimal(formatter: Formatter, escaper: Callable[[str], str], val: An
     Decimal: _format_decimal,
     bool: _format_bool,
     str: _format_str,
+    bytes: _format_binary,
+    bytearray: _format_binary,
+    memoryview: _format_binary,
     list: _format_seq,
     set: _format_seq,
     tuple: _format_seq,
@@ -307,6 +316,7 @@ class DefaultParameterFormatter(Formatter):
     Supported types:
         - None: Converts to SQL NULL
         - Strings: Properly escaped and quoted
+        - Binary data: bytes, bytearray, memoryview as hexadecimal literals
         - Numbers: int, float, Decimal
         - Dates and times: date, datetime, time
         - Booleans: Converted to SQL boolean literals

pyathena/pandas/cursor.py

@@ -294,9 +294,10 @@ def iter_chunks(self) -> Generator[DataFrame, None, None]:
 
         import gc
 
-        for chunk_count, chunk in enumerate(result_set.iter_chunks(), 1):
-            yield chunk
+        with result_set.iter_chunks() as chunks:
+            for chunk_count, chunk in enumerate(chunks, 1):
+                yield chunk
 
-            # Suggest garbage collection every 10 chunks for large datasets
-            if chunk_count % 10 == 0:
-                gc.collect()
+                # Suggest garbage collection every 10 chunks for large datasets
+                if chunk_count % 10 == 0:
+                    gc.collect()

pyathena/pandas/reader.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+import re
+from io import RawIOBase
+from typing import Any
+
+from pyathena.s3fs.reader import AthenaCSVReader
+
+_BINARY_NULL = "__PYATHENA_BINARY_NULL__"
+_CSV_FIELD = re.compile(r'(?:^|,)(?P<value>"[^"]*(?:""[^"]*)*"|[^,]*)')
+
+
+class BinaryCSVReader(RawIOBase):
+    """Preserve binary NULL fields before pandas discards CSV quoting information.
+
+    The marker cannot occur in Athena's hexadecimal encoding of binary values.
+    Only unquoted empty binary fields are rewritten; all other CSV text is
+    passed through unchanged. Records are streamed to support chunked reads.
+    """
+
+    def __init__(self, stream: Any, binary_columns: set[int]) -> None:
+        super().__init__()
+        self._reader = AthenaCSVReader(stream)
+        self._binary_columns = binary_columns
+        self._header = True
+        self._buffer = b""
+
+    def readable(self) -> bool:
+        return True
+
+    def readinto(self, buffer: Any) -> int:
+        if self.closed:
+            raise ValueError("I/O operation on closed file.")
+        if not len(buffer):
+            return 0
+        if not self._buffer:
+            try:
+                record = self._reader._read_record()
+            except StopIteration:
+                self._reader.close()
+                return 0
+            if self._header:
+                self._header = False
+            else:
+                parts: list[str] = []
+                start = 0
+                for index, field in enumerate(_CSV_FIELD.finditer(record.rstrip("\r\n"))):
+                    if index in self._binary_columns and not field.group("value"):
+                        pos = field.start("value")
+                        parts.extend((record[start:pos], _BINARY_NULL))
+                        start = pos
+                parts.append(record[start:])
+                record = "".join(parts)
+            self._buffer = record.encode("utf-8")
+        size = min(len(buffer), len(self._buffer))
+        buffer[:size] = self._buffer[:size]
+        self._buffer = self._buffer[size:]
+        return size
+
+    def close(self) -> None:
+        try:
+            self._reader.close()
+        finally:
+            super().close()