From 146455fe6c59b4e0b87153b426b527d3afd10fd4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 7 Oct 2026 09:28:51 +0900 Subject: [PATCH 1/4] Rename the opt-in S3FS CSV reader and define its protocol --- docs/api/s3fs.rst | 5 +- docs/null_handling.md | 19 ++- docs/s3fs.md | 44 ++++-- pyathena/aio/s3fs/cursor.py | 4 +- pyathena/s3fs/async_cursor.py | 5 +- pyathena/s3fs/cursor.py | 11 +- pyathena/s3fs/reader.py | 63 ++++++-- pyathena/s3fs/result_set.py | 53 +++---- tests/pyathena/aio/s3fs/test_cursor.py | 19 ++- tests/pyathena/s3fs/test_async_cursor.py | 19 ++- tests/pyathena/s3fs/test_cursor.py | 33 +++-- tests/pyathena/s3fs/test_reader.py | 32 ++--- tests/pyathena/s3fs/test_result_set.py | 175 +++++++++++++++++++++++ 13 files changed, 379 insertions(+), 103 deletions(-) create mode 100644 tests/pyathena/s3fs/test_result_set.py diff --git a/docs/api/s3fs.rst b/docs/api/s3fs.rst index 519f1aa86..0ac181d52 100644 --- a/docs/api/s3fs.rst +++ b/docs/api/s3fs.rst @@ -30,10 +30,13 @@ S3FS CSV Readers S3FSCursor supports pluggable CSV reader implementations to control how NULL values and empty strings are handled when parsing Athena's CSV output. +.. autoclass:: pyathena.s3fs.reader.CSVReader + :members: + .. autoclass:: pyathena.s3fs.reader.AthenaCSVReader :members: -.. autoclass:: pyathena.s3fs.reader.DefaultCSVReader +.. autoclass:: pyathena.s3fs.reader.EmptyStringAsNullCSVReader :members: S3FS Data Converters diff --git a/docs/null_handling.md b/docs/null_handling.md index 7c745abb7..4e2daa353 100644 --- a/docs/null_handling.md +++ b/docs/null_handling.md @@ -327,7 +327,7 @@ df = cursor.execute("SELECT * FROM your_table").as_polars() `S3FSCursor` supports two CSV readers with different NULL handling behaviors: - **AthenaCSVReader** (default): Properly distinguishes NULL from empty strings by respecting CSV quoting rules -- **DefaultCSVReader**: For backward compatibility; both NULL and empty strings become `None` +- **EmptyStringAsNullCSVReader**: Opt-in behavior; both NULL and empty strings become `None` in cursor results ### AthenaCSVReader (Default) @@ -361,26 +361,33 @@ for row in cursor: # id=3, value='hello', is_none=False ``` -### DefaultCSVReader (Backward Compatibility) +### EmptyStringAsNullCSVReader -If you need backward compatibility where both NULL and empty strings are treated as `None`, -you can explicitly specify `DefaultCSVReader`: +Specify `EmptyStringAsNullCSVReader` to treat both NULL and empty strings as `None` when reading query result files. +The result set applies the reader's `empty_strings_as_null` flag before type conversion. +When iterated directly, this reader returns empty strings for both cases. + +PyAthena 4.0 renames `DefaultCSVReader` to `EmptyStringAsNullCSVReader`; the old name is removed. +Update existing imports and `csv_reader` arguments to use the new name, as shown below. ```python from pyathena import connect from pyathena.s3fs.cursor import S3FSCursor -from pyathena.s3fs.reader import DefaultCSVReader +from pyathena.s3fs.reader import EmptyStringAsNullCSVReader cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", region_name="us-west-2", cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": DefaultCSVReader}).cursor() + cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}).cursor() cursor.execute("SELECT '' AS empty_col, NULL AS null_col") row = cursor.fetchone() print(row) # (None, None) - both become None ``` +For custom reader requirements, see {ref}`s3fs-cursor`. +The default and opt-in behavior also apply to `AsyncS3FSCursor` and `AioS3FSCursor`, which use the same result set. + ```{note} `S3FSCursor` does not support the `unload` option. However, since `AthenaCSVReader` (the default) can already distinguish NULL from empty strings when reading CSV files, this is not a limitation for NULL handling. diff --git a/docs/s3fs.md b/docs/s3fs.md index 31e44d73b..81bb8f62a 100644 --- a/docs/s3fs.md +++ b/docs/s3fs.md @@ -164,7 +164,7 @@ S3FSCursor supports pluggable CSV reader implementations to control how NULL val are handled. Two readers are provided: - `AthenaCSVReader` (default): Custom parser that distinguishes between NULL and empty string -- `DefaultCSVReader`: Uses Python's built-in `csv` module (treats both NULL and empty string as empty string) +- `EmptyStringAsNullCSVReader`: Uses Python's built-in `csv` module; both NULL and empty string become `None` in cursor results **Default behavior (AthenaCSVReader):** @@ -184,21 +184,22 @@ row = cursor.fetchone() print(row) # (None, '') - NULL is None, empty string is '' ``` -**Switching to Python's built-in csv module (DefaultCSVReader):** +**Treating empty strings as NULL:** -If you prefer to use Python's built-in `csv` module, you can switch to `DefaultCSVReader`. -Note that this reader cannot distinguish between NULL and empty string - both become empty strings -in the parsed result, which are then converted to `None` by the type converter. +Use `EmptyStringAsNullCSVReader` when empty strings should be treated as NULL. +Its parser returns empty strings for both NULL and quoted empty fields. +The result set reads the reader's `empty_strings_as_null` property and replaces empty strings with `None` before applying type conversion. +Iterating over this reader directly still returns empty strings. ```python from pyathena import connect from pyathena.s3fs.cursor import S3FSCursor -from pyathena.s3fs.reader import DefaultCSVReader +from pyathena.s3fs.reader import EmptyStringAsNullCSVReader cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", region_name="us-west-2", cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": DefaultCSVReader}).cursor() + cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}).cursor() cursor.execute("SELECT NULL AS null_col, '' AS empty_col") row = cursor.fetchone() @@ -210,7 +211,7 @@ print(row) # (None, None) - Both NULL and empty string become None | Reader | Implementation | NULL value | Empty string | | --- | --- | --- | --- | | AthenaCSVReader (default) | Custom parser | None | '' (empty string) | -| DefaultCSVReader | Python csv module | None | None | +| EmptyStringAsNullCSVReader | Python csv module | None | None | **Why the difference?** @@ -222,6 +223,33 @@ Athena's CSV output format distinguishes between NULL values and empty strings: Python's standard `csv` module parses both cases as empty strings, losing this distinction. The `AthenaCSVReader` implements a custom parser that preserves the difference. +### CSV reader rename in PyAthena 4.0 + +PyAthena 4.0 renames `DefaultCSVReader` to `EmptyStringAsNullCSVReader` and removes the old name. +Update imports and `csv_reader` arguments to use the new name: + +```python +from pyathena.s3fs.reader import EmptyStringAsNullCSVReader + +cursor = connection.cursor(S3FSCursor, csv_reader=EmptyStringAsNullCSVReader) +``` + +The opt-in behavior is preserved, and `AthenaCSVReader` remains the default. + +### Custom CSV readers + +Custom reader classes must satisfy the `pyathena.s3fs.reader.CSVReader` protocol; inheriting from it is optional. +Pass the class through `csv_reader` in the cursor constructor or an individual `execute()` call. + +- Accept a text stream as `file_obj` and a `delimiter` keyword argument in the constructor. +- Implement `__iter__()` and `__next__()` to yield sequences of strings or `None`, and raise `StopIteration` at the end of the stream. +- Implement `close()` to close the underlying stream. +- Expose `empty_strings_as_null`: set it to `True` to convert empty strings to `None` before type conversion, or `False` to preserve them. + +The built-in readers expose this flag as a read-only property, which subclasses can override. +The result set uses the instance's flag rather than checking the reader class's identity. +This flag applies when reading result files; managed query results and other API fallback rows retain their existing conversion behavior. + ### Limitations S3FSCursor has some limitations compared to ArrowCursor or PandasCursor: diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 37d2a81de..4b16389d6 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -74,7 +74,9 @@ def __init__( result_reuse_minutes: Maximum age of a reused query result in minutes. csv_reader: CSV reader class for parsing the result files. If None, ``AthenaCSVReader`` is used, which distinguishes NULL from empty - strings. ``DefaultCSVReader`` reads both as empty strings. + strings. ``EmptyStringAsNullCSVReader`` treats both as NULL in cursor + results. Custom readers must satisfy the + :class:`~pyathena.s3fs.reader.CSVReader` protocol. **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, passed to the parent ``__init__``. """ diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 5b2f18272..aa3d91216 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -86,8 +86,9 @@ def __init__( csv_reader: CSV reader class to use for parsing results. Use AthenaCSVReader (default) to distinguish between NULL (unquoted empty) and empty string (quoted empty ""). - Use DefaultCSVReader for backward compatibility where empty - strings are treated as NULL. + Use EmptyStringAsNullCSVReader to treat empty strings as NULL. + Custom readers must satisfy the + :class:`~pyathena.s3fs.reader.CSVReader` protocol. **kwargs: Additional connection parameters. Example: diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index 8ec1f7204..5cb40221a 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -82,17 +82,18 @@ def __init__( csv_reader: CSV reader class to use for parsing results. Use AthenaCSVReader (default) to distinguish between NULL (unquoted empty) and empty string (quoted empty ""). - Use DefaultCSVReader for backward compatibility where empty - strings are treated as NULL. + Use EmptyStringAsNullCSVReader to treat empty strings as NULL. + Custom readers must satisfy the + :class:`~pyathena.s3fs.reader.CSVReader` protocol. **kwargs: Additional connection parameters. Example: >>> cursor = connection.cursor(S3FSCursor) >>> cursor.execute("SELECT * FROM my_table") >>> - >>> # Use DefaultCSVReader for backward compatibility - >>> from pyathena.s3fs.reader import DefaultCSVReader - >>> cursor = connection.cursor(S3FSCursor, csv_reader=DefaultCSVReader) + >>> # Treat empty strings as NULL + >>> from pyathena.s3fs.reader import EmptyStringAsNullCSVReader + >>> cursor = connection.cursor(S3FSCursor, csv_reader=EmptyStringAsNullCSVReader) """ super().__init__( s3_staging_dir=s3_staging_dir, diff --git a/pyathena/s3fs/reader.py b/pyathena/s3fs/reader.py index 098857401..061bfd629 100644 --- a/pyathena/s3fs/reader.py +++ b/pyathena/s3fs/reader.py @@ -10,25 +10,56 @@ from __future__ import annotations import csv -from collections.abc import Iterator -from typing import Any +from collections.abc import Iterator, Sequence +from typing import Any, Protocol from pyathena.util import override -class DefaultCSVReader(Iterator[list[str]]): - """CSV reader using Python's standard csv module. +class CSVReader(Protocol): + """Contract for readers of Athena CSV result files. + + Reader classes accept a text stream and a delimiter. They yield sequences + containing strings or None, and close the underlying stream when closed. + Implementations do not need to inherit from this protocol. + """ + + def __init__(self, file_obj: Any, delimiter: str = ",") -> None: + """Initialize the reader from a text stream and field delimiter.""" + ... + + @property + def empty_strings_as_null(self) -> bool: + """Whether the result set converts empty strings to None before conversion.""" + ... + + def __iter__(self) -> Iterator[Sequence[str | None]]: + """Iterate over rows in the result file.""" + ... + + def __next__(self) -> Sequence[str | None]: + """Read the next row, raising StopIteration at the end of the stream.""" + ... + + def close(self) -> None: + """Close the underlying stream.""" + ... + + +class EmptyStringAsNullCSVReader(Iterator[list[str]]): + """CSV reader for treating empty strings as NULL in S3FS cursor results. This reader wraps Python's standard csv.reader and treats empty fields as empty strings. It does not distinguish between NULL and empty strings in Athena's CSV output - both become empty strings. - Use this reader when you need backward compatibility with the behavior - where empty strings are treated the same as NULL values. + When used by an S3FS cursor, the result set converts these empty strings + to None before applying type conversion. Iterating over the reader directly + still returns strings from Python's standard csv.reader. Example: >>> from io import StringIO - >>> reader = DefaultCSVReader(StringIO(',"",text')) + >>> reader = EmptyStringAsNullCSVReader(StringIO(',"",text')) >>> list(reader) [['', '', 'text']] # Both NULL and empty string become '' @@ -47,8 +78,13 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._file: Any | None = file_obj self._reader = csv.reader(file_obj, delimiter=delimiter) + @property + def empty_strings_as_null(self) -> bool: + """Whether the result set converts empty strings to None before conversion.""" + return True + @override - def __iter__(self) -> DefaultCSVReader: + def __iter__(self) -> EmptyStringAsNullCSVReader: """Iterate over rows in the CSV file.""" return self @@ -77,7 +113,7 @@ def close(self) -> None: self._file.close() self._file = None - def __enter__(self) -> DefaultCSVReader: + def __enter__(self) -> EmptyStringAsNullCSVReader: """Enter context manager.""" return self @@ -106,8 +142,8 @@ class AthenaCSVReader(Iterator[list[str | None]]): [[None, '', 'text']] # NULL and empty string are distinguished Note: - Use DefaultCSVReader if you need backward compatibility where both - NULL and empty string are treated as empty string. + Use EmptyStringAsNullCSVReader to treat both NULL and empty string as None + in S3FS cursor results. """ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: @@ -120,6 +156,11 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._file: Any | None = file_obj self._delimiter = delimiter + @property + def empty_strings_as_null(self) -> bool: + """Whether the result set converts empty strings to None before conversion.""" + return False + @override def __iter__(self) -> AthenaCSVReader: """Iterate over rows in the CSV file.""" diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 0bc268fc9..58379748b 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -20,13 +20,13 @@ from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet -from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader +from pyathena.s3fs.reader import AthenaCSVReader, CSVReader from pyathena.util import RetryConfig, override, parse_output_location if TYPE_CHECKING: from pyathena.connection import Connection -CSVReaderType = type[DefaultCSVReader] | type[AthenaCSVReader] +CSVReaderType = type[CSVReader] _logger = logging.getLogger(__name__) @@ -88,7 +88,8 @@ def __init__( block_size: The default block size in bytes for the filesystem. If not set, ``DEFAULT_BLOCK_SIZE`` is used. csv_reader: The CSV reader class for the results. If None, - ``AthenaCSVReader`` is used. + ``AthenaCSVReader`` is used. Custom readers must satisfy the + :class:`~pyathena.s3fs.reader.CSVReader` protocol. filesystem_class: The filesystem class for reading the results. If None, PyAthena's ``S3FileSystem`` is used. result_set_type_hints: Athena type signatures for complex-type columns, @@ -115,7 +116,7 @@ def __init__( self._csv_reader_class: CSVReaderType = csv_reader or AthenaCSVReader self._filesystem_class: type[AbstractFileSystem] = filesystem_class or S3FileSystem self._fs = self._create_s3_file_system() - self._csv_reader: Any | None = None + self._csv_reader: CSVReader | None = None if self.state == AthenaQueryExecution.STATE_SUCCEEDED and self.output_location: self._init_csv_reader() @@ -199,37 +200,21 @@ def _fetch(self) -> None: except StopIteration: break - # Convert row values using converters - # AthenaCSVReader returns None for NULL values directly, - # DefaultCSVReader returns empty string which needs conversion - if self._csv_reader_class is DefaultCSVReader: - if col_hints: - converted_row = tuple( - self._converter.convert( - col_type, value if value != "" else None, type_hint=hint - ) - if hint - else self._converter.convert(col_type, value if value != "" else None) - for col_type, value, hint in zip(col_types, row, col_hints, strict=False) - ) - else: - converted_row = tuple( - self._converter.convert(col_type, value if value != "" else None) - for col_type, value in zip(col_types, row, strict=False) - ) + if self._csv_reader.empty_strings_as_null: + row = [None if value == "" else value for value in row] + + if col_hints: + converted_row = tuple( + self._converter.convert(col_type, value, type_hint=hint) + if hint + else self._converter.convert(col_type, value) + for col_type, value, hint in zip(col_types, row, col_hints, strict=False) + ) else: - if col_hints: - converted_row = tuple( - self._converter.convert(col_type, value, type_hint=hint) - if hint - else self._converter.convert(col_type, value) - for col_type, value, hint in zip(col_types, row, col_hints, strict=False) - ) - else: - converted_row = tuple( - self._converter.convert(col_type, value) - for col_type, value in zip(col_types, row, strict=False) - ) + converted_row = tuple( + self._converter.convert(col_type, value) + for col_type, value in zip(col_types, row, strict=False) + ) self._rows.append(converted_row) rows_fetched += 1 diff --git a/tests/pyathena/aio/s3fs/test_cursor.py b/tests/pyathena/aio/s3fs/test_cursor.py index 42f29ab1c..02ed3f174 100644 --- a/tests/pyathena/aio/s3fs/test_cursor.py +++ b/tests/pyathena/aio/s3fs/test_cursor.py @@ -12,7 +12,7 @@ from pyathena.aio.s3fs.cursor import AioS3FSCursor from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution -from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader +from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig from tests import ENV @@ -20,6 +20,21 @@ class TestAioS3FSCursor: + @pytest.mark.parametrize( + ("aio_s3fs_cursor", "expected_empty"), + [ + ({}, ""), + ({"csv_reader": AthenaCSVReader}, ""), + ({"csv_reader": EmptyStringAsNullCSVReader}, None), + ], + indirect=["aio_s3fs_cursor"], + ) + async def test_null_vs_empty_string(self, aio_s3fs_cursor, expected_empty): + await aio_s3fs_cursor.execute( + "SELECT NULL AS null_col, '' AS empty_col, 'text' AS value_col" + ) + assert await aio_s3fs_cursor.fetchall() == [(None, expected_empty, "text")] + async def test_fetchone(self, aio_s3fs_cursor): await aio_s3fs_cursor.execute("SELECT * FROM one_row") assert aio_s3fs_cursor.rownumber == 0 @@ -94,7 +109,7 @@ async def test_read_options(self, execute_kwargs): No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"csv_reader": DefaultCSVReader} + cursor_kwargs = {"csv_reader": EmptyStringAsNullCSVReader} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = AioS3FSCursor( connection=MagicMock(), diff --git a/tests/pyathena/s3fs/test_async_cursor.py b/tests/pyathena/s3fs/test_async_cursor.py index ba28c93b6..02ffb3f44 100644 --- a/tests/pyathena/s3fs/test_async_cursor.py +++ b/tests/pyathena/s3fs/test_async_cursor.py @@ -17,7 +17,7 @@ from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.s3fs.async_cursor import AsyncS3FSCursor -from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader +from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig from tests import ENV @@ -25,6 +25,21 @@ class TestAsyncS3FSCursor: + @pytest.mark.parametrize( + ("async_s3fs_cursor", "expected_empty"), + [ + ({}, ""), + ({"csv_reader": AthenaCSVReader}, ""), + ({"csv_reader": EmptyStringAsNullCSVReader}, None), + ], + indirect=["async_s3fs_cursor"], + ) + def test_null_vs_empty_string(self, async_s3fs_cursor, expected_empty): + _, future = async_s3fs_cursor.execute( + "SELECT NULL AS null_col, '' AS empty_col, 'text' AS value_col" + ) + assert future.result().fetchall() == [(None, expected_empty, "text")] + def test_fetchone(self, async_s3fs_cursor): query_id, future = async_s3fs_cursor.execute("SELECT * FROM one_row") result_set = future.result() @@ -189,7 +204,7 @@ def test_read_options(self, execute_kwargs): No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"csv_reader": DefaultCSVReader} + cursor_kwargs = {"csv_reader": EmptyStringAsNullCSVReader} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) with ( AsyncS3FSCursor( diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 4b643d5b0..43d8fabf7 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -14,7 +14,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.cursor import S3FSCursor -from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader +from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig from tests import ENV @@ -333,7 +333,7 @@ def test_contain_tab_character(self, s3fs_cursor): @pytest.mark.parametrize( "csv_reader_class", - [DefaultCSVReader, AthenaCSVReader], + [EmptyStringAsNullCSVReader, AthenaCSVReader], ) def test_basic_query_with_reader(self, csv_reader_class): """Both readers should work for basic queries.""" @@ -353,7 +353,7 @@ def test_basic_query_with_reader(self, csv_reader_class): @pytest.mark.parametrize( "csv_reader_class", - [DefaultCSVReader, AthenaCSVReader], + [EmptyStringAsNullCSVReader, AthenaCSVReader], ) def test_multiple_columns_with_reader(self, csv_reader_class): """Test multiple columns work with both readers.""" @@ -378,14 +378,14 @@ def test_multiple_columns_with_reader(self, csv_reader_class): result = cursor.fetchone() assert result == (1, "text", 1.5) - def test_null_with_default_reader(self): - """DefaultCSVReader: NULL is returned as None.""" + def test_null_with_empty_string_as_null_reader(self): + """EmptyStringAsNullCSVReader: NULL is returned as None.""" with ( contextlib.closing( connect( schema_name=ENV.schema, cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": DefaultCSVReader}, + cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}, ) ) as conn, conn.cursor() as cursor, @@ -410,21 +410,21 @@ def test_null_with_athena_reader(self): result = cursor.fetchone() assert result == (None,) - def test_empty_string_with_default_reader(self): - """DefaultCSVReader: Empty string becomes None (loses distinction).""" + def test_empty_string_with_empty_string_as_null_reader(self): + """EmptyStringAsNullCSVReader: Empty string becomes None (loses distinction).""" with ( contextlib.closing( connect( schema_name=ENV.schema, cursor_class=S3FSCursor, - cursor_kwargs={"csv_reader": DefaultCSVReader}, + cursor_kwargs={"csv_reader": EmptyStringAsNullCSVReader}, ) ) as conn, conn.cursor() as cursor, ): cursor.execute("SELECT '' AS empty_col") result = cursor.fetchone() - # DefaultCSVReader treats empty string same as NULL + # EmptyStringAsNullCSVReader treats empty string same as NULL assert result == (None,) def test_empty_string_with_athena_reader(self): @@ -447,7 +447,10 @@ def test_empty_string_with_athena_reader(self): @pytest.mark.parametrize( ("csv_reader", "expected_empty"), [ - (DefaultCSVReader, None), # DefaultCSVReader: empty string becomes None + ( + EmptyStringAsNullCSVReader, + None, + ), # EmptyStringAsNullCSVReader: empty string becomes None (AthenaCSVReader, ""), # AthenaCSVReader: empty string is preserved ], ) @@ -455,10 +458,10 @@ def test_null_vs_empty_string(self, csv_reader, expected_empty): """ Test NULL vs empty string handling with different CSV readers. - DefaultCSVReader: Both NULL and empty string become None. + EmptyStringAsNullCSVReader: Both NULL and empty string become None. AthenaCSVReader: NULL is None, empty string is preserved as ''. - See docs/null_handling.rst for details. + See docs/null_handling.md for details. """ with ( contextlib.closing( @@ -503,7 +506,7 @@ def test_mixed_values_with_athena_reader(self): @pytest.mark.parametrize( "csv_reader_class", - [DefaultCSVReader, AthenaCSVReader], + [EmptyStringAsNullCSVReader, AthenaCSVReader], ) def test_quoted_string_with_comma(self, csv_reader_class): """Both readers should handle strings containing commas.""" @@ -569,7 +572,7 @@ def test_read_options(self, execute_kwargs): No AWS calls; the query and its result set are mocked. """ - cursor_kwargs = {"csv_reader": DefaultCSVReader} + cursor_kwargs = {"csv_reader": EmptyStringAsNullCSVReader} query_execution = MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED) cursor = S3FSCursor( connection=MagicMock(), diff --git a/tests/pyathena/s3fs/test_reader.py b/tests/pyathena/s3fs/test_reader.py index e2f46858b..762abeb49 100644 --- a/tests/pyathena/s3fs/test_reader.py +++ b/tests/pyathena/s3fs/test_reader.py @@ -8,54 +8,54 @@ from collections.abc import Iterator from io import StringIO -from pyathena.s3fs.reader import AthenaCSVReader, DefaultCSVReader +from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader -class TestDefaultCSVReader: - """Tests for DefaultCSVReader using Python's standard csv module.""" +class TestEmptyStringAsNullCSVReader: + """Tests for EmptyStringAsNullCSVReader using Python's standard csv module.""" def test_basic_parsing(self): data = StringIO("a,b,c\n1,2,3\n") - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [["a", "b", "c"], ["1", "2", "3"]] def test_tab_delimiter(self): data = StringIO("a\tb\tc\n1\t2\t3\n") - reader = DefaultCSVReader(data, delimiter="\t") + reader = EmptyStringAsNullCSVReader(data, delimiter="\t") rows = list(reader) assert rows == [["a", "b", "c"], ["1", "2", "3"]] def test_empty_field_returns_empty_string(self): - """DefaultCSVReader returns empty string for both NULL and empty string.""" + """EmptyStringAsNullCSVReader returns empty string for both NULL and empty string.""" # Both ,, (NULL) and ,"", (empty string) become '' data = StringIO('a,,b,"",c\n') - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) # Standard csv.reader treats both as empty strings assert rows == [["a", "", "b", "", "c"]] def test_quoted_field_with_comma(self): data = StringIO('"a,b",c\n') - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [["a,b", "c"]] def test_quoted_field_with_newline(self): data = StringIO('"line1\nline2",b\n') - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [["line1\nline2", "b"]] def test_escaped_quote(self): data = StringIO('"a""b",c\n') - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [['a"b', "c"]] def test_empty_file(self): data = StringIO("") - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [] @@ -63,20 +63,20 @@ def test_empty_line(self): """Empty line returns single empty string (e.g., SELECT NULL).""" # Python's csv.reader returns [] for empty lines; we normalize to [''] data = StringIO("\n") - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) assert rows == [[""]] def test_implements_iterator_protocol(self): - """DefaultCSVReader implements collections.abc.Iterator.""" + """EmptyStringAsNullCSVReader implements collections.abc.Iterator.""" data = StringIO("a,b\n") - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") assert isinstance(reader, Iterator) def test_close(self): """close() releases file resources.""" data = StringIO("a,b\n") - reader = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") reader.close() # After close, iteration should stop immediately rows = list(reader) @@ -85,7 +85,7 @@ def test_close(self): def test_context_manager(self): """Reader can be used as context manager.""" data = StringIO("a,b\n1,2\n") - with DefaultCSVReader(data, delimiter=",") as reader: + with EmptyStringAsNullCSVReader(data, delimiter=",") as reader: rows = list(reader) assert rows == [["a", "b"], ["1", "2"]] # After exiting context, reader should be closed diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py new file mode 100644 index 000000000..4d6abe662 --- /dev/null +++ b/tests/pyathena/s3fs/test_result_set.py @@ -0,0 +1,175 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +from collections.abc import Iterator +from io import BytesIO +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from pyathena.model import AthenaQueryExecution +from pyathena.s3fs.converter import DefaultS3FSTypeConverter +from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader +from pyathena.s3fs.result_set import AthenaS3FSResultSet +from pyathena.util import RetryConfig, override + + +class _InheritedEmptyStringAsNullCSVReader(EmptyStringAsNullCSVReader): + pass + + +class _NullifyingAthenaCSVReader(AthenaCSVReader): + @property + @override + def empty_strings_as_null(self) -> bool: + return True + + +class _TupleCSVReader: + """An independent implementation that returns tuples rather than lists.""" + + empty_strings_as_null = False + + def __init__(self, file_obj: Any, delimiter: str = ",") -> None: + self._reader = AthenaCSVReader(file_obj, delimiter=delimiter) + + def __iter__(self) -> Iterator[tuple[str | None, ...]]: + return self + + def __next__(self) -> tuple[str | None, ...]: + return tuple(next(self._reader)) + + def close(self) -> None: + self._reader.close() + + +class _NullifyingTupleCSVReader(_TupleCSVReader): + empty_strings_as_null = True + + +def _result_set(data, csv_reader, suffix=".csv", column_types=None, type_hints=None): + stream = BytesIO(data.encode()) + filesystem = MagicMock() + filesystem._open.return_value = stream + connection = MagicMock() + connection.s3_client.head_object.return_value = {"ContentLength": len(data.encode())} + response = { + "ResultSet": { + "ResultSetMetadata": { + "ColumnInfo": [ + {"Name": name, "Type": col_type} + for name, col_type in zip( + ("null_col", "empty_col", "value_col"), + column_types or ("varchar", "varchar", "varchar"), + strict=True, + ) + ] + }, + "Rows": [], + } + } + with patch.object(AthenaS3FSResultSet, "_get_query_results", return_value=response): + result_set = AthenaS3FSResultSet( + connection=connection, + converter=DefaultS3FSTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location=f"s3://bucket/query{suffix}", + substatement_type="SELECT", + ), + arraysize=1, + retry_config=RetryConfig(), + csv_reader=csv_reader, + filesystem_class=MagicMock(return_value=filesystem), + result_set_type_hints=type_hints, + ) + return result_set, stream + + +@pytest.mark.parametrize( + ("csv_reader", "expected_empty"), + [ + (None, ""), + (AthenaCSVReader, ""), + (EmptyStringAsNullCSVReader, None), + (_InheritedEmptyStringAsNullCSVReader, None), + (_NullifyingAthenaCSVReader, None), + (_TupleCSVReader, ""), + (_NullifyingTupleCSVReader, None), + ], +) +@pytest.mark.parametrize( + ("suffix", "data"), + [ + (".csv", 'null_col,empty_col,value_col\n,"",text\nnext,"",last\n'), + (".txt", '\t""\ttext\nnext\t""\tlast\n'), + ], +) +def test_reader_null_conversion(csv_reader, expected_empty, suffix, data): + result_set, stream = _result_set(data, csv_reader, suffix=suffix) + with result_set: + assert result_set.fetchmany() == [(None, expected_empty, "text")] + assert result_set.fetchall() == [("next", expected_empty, "last")] + assert result_set.fetchone() is None + assert result_set.rownumber == 2 + assert stream.closed + + +@pytest.mark.parametrize( + ("csv_reader", "expected_empty"), + [ + (AthenaCSVReader, ""), + (EmptyStringAsNullCSVReader, None), + (_InheritedEmptyStringAsNullCSVReader, None), + (_NullifyingAthenaCSVReader, None), + (_TupleCSVReader, ""), + (_NullifyingTupleCSVReader, None), + ], +) +def test_reader_null_conversion_with_type_hints(csv_reader, expected_empty): + result_set, stream = _result_set( + 'null_col,empty_col,value_col\n,"","[1, 2]"\n,"",\n', + csv_reader, + column_types=("varchar", "varchar", "array"), + type_hints={"value_col": "array(integer)"}, + ) + with result_set: + assert result_set.fetchall() == [ + (None, expected_empty, [1, 2]), + (None, expected_empty, None), + ] + assert stream.closed + + +@pytest.mark.parametrize("csv_reader", [AthenaCSVReader, EmptyStringAsNullCSVReader]) +def test_managed_results_preserve_api_values(csv_reader): + response = { + "ResultSet": { + "ResultSetMetadata": {"ColumnInfo": [{"Name": "value", "Type": "varchar"}]}, + "Rows": [ + {"Data": [{"VarCharValue": "value"}]}, + {"Data": [{"VarCharValue": ""}]}, + {"Data": [{}]}, + ], + } + } + with ( + patch.object(AthenaS3FSResultSet, "_get_query_results", return_value=response), + AthenaS3FSResultSet( + connection=MagicMock(), + converter=DefaultS3FSTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, output_location=None + ), + arraysize=1, + retry_config=RetryConfig(), + csv_reader=csv_reader, + filesystem_class=MagicMock(), + ) as result_set, + ): + assert result_set.fetchall() == [("",), (None,)] From 774219e113d55e6bc4e54c290e7351e7dcfac433 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 7 Oct 2026 09:37:05 +0900 Subject: [PATCH 2/4] Pass reader options through integration cursor fixtures --- docs/s3fs.md | 1 + tests/pyathena/aio/s3fs/test_cursor.py | 4 ++-- tests/pyathena/s3fs/test_async_cursor.py | 4 ++-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/s3fs.md b/docs/s3fs.md index 81bb8f62a..e92242682 100644 --- a/docs/s3fs.md +++ b/docs/s3fs.md @@ -235,6 +235,7 @@ cursor = connection.cursor(S3FSCursor, csv_reader=EmptyStringAsNullCSVReader) ``` The opt-in behavior is preserved, and `AthenaCSVReader` remains the default. +Custom readers must now expose `empty_strings_as_null`; use `False` to retain empty strings, or `True` to treat them as NULL in result files. ### Custom CSV readers diff --git a/tests/pyathena/aio/s3fs/test_cursor.py b/tests/pyathena/aio/s3fs/test_cursor.py index 02ed3f174..7f0b0fbb6 100644 --- a/tests/pyathena/aio/s3fs/test_cursor.py +++ b/tests/pyathena/aio/s3fs/test_cursor.py @@ -24,8 +24,8 @@ class TestAioS3FSCursor: ("aio_s3fs_cursor", "expected_empty"), [ ({}, ""), - ({"csv_reader": AthenaCSVReader}, ""), - ({"csv_reader": EmptyStringAsNullCSVReader}, None), + ({"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, ""), + ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), ], indirect=["aio_s3fs_cursor"], ) diff --git a/tests/pyathena/s3fs/test_async_cursor.py b/tests/pyathena/s3fs/test_async_cursor.py index 02ffb3f44..0926b0f9b 100644 --- a/tests/pyathena/s3fs/test_async_cursor.py +++ b/tests/pyathena/s3fs/test_async_cursor.py @@ -29,8 +29,8 @@ class TestAsyncS3FSCursor: ("async_s3fs_cursor", "expected_empty"), [ ({}, ""), - ({"csv_reader": AthenaCSVReader}, ""), - ({"csv_reader": EmptyStringAsNullCSVReader}, None), + ({"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, ""), + ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), ], indirect=["async_s3fs_cursor"], ) From 010be3f740c587c9f86f5f8743d0388fa52f19ae Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 23:08:44 +0900 Subject: [PATCH 3/4] Group S3FS result-set tests by object behavior --- tests/pyathena/s3fs/test_result_set.py | 159 ++++++++++++------------- 1 file changed, 79 insertions(+), 80 deletions(-) diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 4d6abe662..d3ffda87e 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -91,85 +91,84 @@ def _result_set(data, csv_reader, suffix=".csv", column_types=None, type_hints=N return result_set, stream -@pytest.mark.parametrize( - ("csv_reader", "expected_empty"), - [ - (None, ""), - (AthenaCSVReader, ""), - (EmptyStringAsNullCSVReader, None), - (_InheritedEmptyStringAsNullCSVReader, None), - (_NullifyingAthenaCSVReader, None), - (_TupleCSVReader, ""), - (_NullifyingTupleCSVReader, None), - ], -) -@pytest.mark.parametrize( - ("suffix", "data"), - [ - (".csv", 'null_col,empty_col,value_col\n,"",text\nnext,"",last\n'), - (".txt", '\t""\ttext\nnext\t""\tlast\n'), - ], -) -def test_reader_null_conversion(csv_reader, expected_empty, suffix, data): - result_set, stream = _result_set(data, csv_reader, suffix=suffix) - with result_set: - assert result_set.fetchmany() == [(None, expected_empty, "text")] - assert result_set.fetchall() == [("next", expected_empty, "last")] - assert result_set.fetchone() is None - assert result_set.rownumber == 2 - assert stream.closed - - -@pytest.mark.parametrize( - ("csv_reader", "expected_empty"), - [ - (AthenaCSVReader, ""), - (EmptyStringAsNullCSVReader, None), - (_InheritedEmptyStringAsNullCSVReader, None), - (_NullifyingAthenaCSVReader, None), - (_TupleCSVReader, ""), - (_NullifyingTupleCSVReader, None), - ], -) -def test_reader_null_conversion_with_type_hints(csv_reader, expected_empty): - result_set, stream = _result_set( - 'null_col,empty_col,value_col\n,"","[1, 2]"\n,"",\n', - csv_reader, - column_types=("varchar", "varchar", "array"), - type_hints={"value_col": "array(integer)"}, +class TestAthenaS3FSResultSet: + @pytest.mark.parametrize( + ("csv_reader", "expected_empty"), + [ + (None, ""), + (AthenaCSVReader, ""), + (EmptyStringAsNullCSVReader, None), + (_InheritedEmptyStringAsNullCSVReader, None), + (_NullifyingAthenaCSVReader, None), + (_TupleCSVReader, ""), + (_NullifyingTupleCSVReader, None), + ], ) - with result_set: - assert result_set.fetchall() == [ - (None, expected_empty, [1, 2]), - (None, expected_empty, None), - ] - assert stream.closed - - -@pytest.mark.parametrize("csv_reader", [AthenaCSVReader, EmptyStringAsNullCSVReader]) -def test_managed_results_preserve_api_values(csv_reader): - response = { - "ResultSet": { - "ResultSetMetadata": {"ColumnInfo": [{"Name": "value", "Type": "varchar"}]}, - "Rows": [ - {"Data": [{"VarCharValue": "value"}]}, - {"Data": [{"VarCharValue": ""}]}, - {"Data": [{}]}, - ], + @pytest.mark.parametrize( + ("suffix", "data"), + [ + (".csv", 'null_col,empty_col,value_col\n,"",text\nnext,"",last\n'), + (".txt", '\t""\ttext\nnext\t""\tlast\n'), + ], + ) + def test_reader_null_conversion(self, csv_reader, expected_empty, suffix, data): + result_set, stream = _result_set(data, csv_reader, suffix=suffix) + with result_set: + assert result_set.fetchmany() == [(None, expected_empty, "text")] + assert result_set.fetchall() == [("next", expected_empty, "last")] + assert result_set.fetchone() is None + assert result_set.rownumber == 2 + assert stream.closed + + @pytest.mark.parametrize( + ("csv_reader", "expected_empty"), + [ + (AthenaCSVReader, ""), + (EmptyStringAsNullCSVReader, None), + (_InheritedEmptyStringAsNullCSVReader, None), + (_NullifyingAthenaCSVReader, None), + (_TupleCSVReader, ""), + (_NullifyingTupleCSVReader, None), + ], + ) + def test_reader_null_conversion_with_type_hints(self, csv_reader, expected_empty): + result_set, stream = _result_set( + 'null_col,empty_col,value_col\n,"","[1, 2]"\n,"",\n', + csv_reader, + column_types=("varchar", "varchar", "array"), + type_hints={"value_col": "array(integer)"}, + ) + with result_set: + assert result_set.fetchall() == [ + (None, expected_empty, [1, 2]), + (None, expected_empty, None), + ] + assert stream.closed + + @pytest.mark.parametrize("csv_reader", [AthenaCSVReader, EmptyStringAsNullCSVReader]) + def test_managed_results_preserve_api_values(self, csv_reader): + response = { + "ResultSet": { + "ResultSetMetadata": {"ColumnInfo": [{"Name": "value", "Type": "varchar"}]}, + "Rows": [ + {"Data": [{"VarCharValue": "value"}]}, + {"Data": [{"VarCharValue": ""}]}, + {"Data": [{}]}, + ], + } } - } - with ( - patch.object(AthenaS3FSResultSet, "_get_query_results", return_value=response), - AthenaS3FSResultSet( - connection=MagicMock(), - converter=DefaultS3FSTypeConverter(), - query_execution=MagicMock( - state=AthenaQueryExecution.STATE_SUCCEEDED, output_location=None - ), - arraysize=1, - retry_config=RetryConfig(), - csv_reader=csv_reader, - filesystem_class=MagicMock(), - ) as result_set, - ): - assert result_set.fetchall() == [("",), (None,)] + with ( + patch.object(AthenaS3FSResultSet, "_get_query_results", return_value=response), + AthenaS3FSResultSet( + connection=MagicMock(), + converter=DefaultS3FSTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, output_location=None + ), + arraysize=1, + retry_config=RetryConfig(), + csv_reader=csv_reader, + filesystem_class=MagicMock(), + ) as result_set, + ): + assert result_set.fetchall() == [("",), (None,)] From f41b8a07bd860ec94496f0c4c2eaafdf047c193e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 00:53:29 +0900 Subject: [PATCH 4/4] Read empty fields as None in EmptyStringAsNullCSVReader Move the empty-string-to-NULL mapping from the result set into the opt-in reader, so the reader returns what its name says and the result set needs no per-reader flag. Drop empty_strings_as_null from both readers and the CSVReader protocol; custom readers that worked before keep working without a new attribute. Remove the redundant explicit AthenaCSVReader case from the async and aio integration tests, which duplicated the default case. Co-Authored-By: Claude Opus 5.5 --- docs/null_handling.md | 2 -- docs/s3fs.md | 15 ++++---- pyathena/s3fs/reader.py | 45 +++++++----------------- pyathena/s3fs/result_set.py | 3 -- tests/pyathena/aio/s3fs/test_cursor.py | 1 - tests/pyathena/s3fs/test_async_cursor.py | 1 - tests/pyathena/s3fs/test_cursor.py | 6 ++-- tests/pyathena/s3fs/test_reader.py | 15 ++++---- tests/pyathena/s3fs/test_result_set.py | 25 +------------ 9 files changed, 29 insertions(+), 84 deletions(-) diff --git a/docs/null_handling.md b/docs/null_handling.md index 4e2daa353..8de356ff4 100644 --- a/docs/null_handling.md +++ b/docs/null_handling.md @@ -364,8 +364,6 @@ for row in cursor: ### EmptyStringAsNullCSVReader Specify `EmptyStringAsNullCSVReader` to treat both NULL and empty strings as `None` when reading query result files. -The result set applies the reader's `empty_strings_as_null` flag before type conversion. -When iterated directly, this reader returns empty strings for both cases. PyAthena 4.0 renames `DefaultCSVReader` to `EmptyStringAsNullCSVReader`; the old name is removed. Update existing imports and `csv_reader` arguments to use the new name, as shown below. diff --git a/docs/s3fs.md b/docs/s3fs.md index e92242682..13d7443b6 100644 --- a/docs/s3fs.md +++ b/docs/s3fs.md @@ -187,9 +187,7 @@ print(row) # (None, '') - NULL is None, empty string is '' **Treating empty strings as NULL:** Use `EmptyStringAsNullCSVReader` when empty strings should be treated as NULL. -Its parser returns empty strings for both NULL and quoted empty fields. -The result set reads the reader's `empty_strings_as_null` property and replaces empty strings with `None` before applying type conversion. -Iterating over this reader directly still returns empty strings. +It parses result files with Python's `csv` module and returns `None` for both NULL and quoted empty fields. ```python from pyathena import connect @@ -234,8 +232,9 @@ from pyathena.s3fs.reader import EmptyStringAsNullCSVReader cursor = connection.cursor(S3FSCursor, csv_reader=EmptyStringAsNullCSVReader) ``` -The opt-in behavior is preserved, and `AthenaCSVReader` remains the default. -Custom readers must now expose `empty_strings_as_null`; use `False` to retain empty strings, or `True` to treat them as NULL in result files. +`AthenaCSVReader` remains the default. +Iterating over `EmptyStringAsNullCSVReader` directly now returns `None` for empty fields instead of empty strings. +Cursor results are unchanged. ### Custom CSV readers @@ -244,12 +243,10 @@ Pass the class through `csv_reader` in the cursor constructor or an individual ` - Accept a text stream as `file_obj` and a `delimiter` keyword argument in the constructor. - Implement `__iter__()` and `__next__()` to yield sequences of strings or `None`, and raise `StopIteration` at the end of the stream. + The cursor converts each value by its column type and treats `None` as NULL. - Implement `close()` to close the underlying stream. -- Expose `empty_strings_as_null`: set it to `True` to convert empty strings to `None` before type conversion, or `False` to preserve them. -The built-in readers expose this flag as a read-only property, which subclasses can override. -The result set uses the instance's flag rather than checking the reader class's identity. -This flag applies when reading result files; managed query results and other API fallback rows retain their existing conversion behavior. +The reader is used only for result files; managed query results are read through the Athena API. ### Limitations diff --git a/pyathena/s3fs/reader.py b/pyathena/s3fs/reader.py index 061bfd629..53da4da4d 100644 --- a/pyathena/s3fs/reader.py +++ b/pyathena/s3fs/reader.py @@ -20,19 +20,14 @@ class CSVReader(Protocol): """Contract for readers of Athena CSV result files. Reader classes accept a text stream and a delimiter. They yield sequences - containing strings or None, and close the underlying stream when closed. - Implementations do not need to inherit from this protocol. + of field values, with None for NULL, and close the underlying stream when + closed. Implementations do not need to inherit from this protocol. """ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: """Initialize the reader from a text stream and field delimiter.""" ... - @property - def empty_strings_as_null(self) -> bool: - """Whether the result set converts empty strings to None before conversion.""" - ... - def __iter__(self) -> Iterator[Sequence[str | None]]: """Iterate over rows in the result file.""" ... @@ -46,22 +41,18 @@ def close(self) -> None: ... -class EmptyStringAsNullCSVReader(Iterator[list[str]]): - """CSV reader for treating empty strings as NULL in S3FS cursor results. +class EmptyStringAsNullCSVReader(Iterator[list[str | None]]): + """CSV reader that reads both NULL and empty string as None. - This reader wraps Python's standard csv.reader and treats empty fields - as empty strings. It does not distinguish between NULL and empty strings - in Athena's CSV output - both become empty strings. - - When used by an S3FS cursor, the result set converts these empty strings - to None before applying type conversion. Iterating over the reader directly - still returns strings from Python's standard csv.reader. + This reader wraps Python's standard csv.reader, which does not distinguish + between NULL and empty strings in Athena's CSV output. It returns None for + every empty field. Example: >>> from io import StringIO >>> reader = EmptyStringAsNullCSVReader(StringIO(',"",text')) >>> list(reader) - [['', '', 'text']] # Both NULL and empty string become '' + [[None, None, 'text']] # Both NULL and empty string become None Note: The default reader for S3FSCursor is AthenaCSVReader, which @@ -78,22 +69,17 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._file: Any | None = file_obj self._reader = csv.reader(file_obj, delimiter=delimiter) - @property - def empty_strings_as_null(self) -> bool: - """Whether the result set converts empty strings to None before conversion.""" - return True - @override def __iter__(self) -> EmptyStringAsNullCSVReader: """Iterate over rows in the CSV file.""" return self @override - def __next__(self) -> list[str]: + def __next__(self) -> list[str | None]: """Read and parse the next line. Returns: - List of field values as strings. + List of field values, with None for empty fields. Raises: StopIteration: When end of file is reached or reader is closed. @@ -101,11 +87,11 @@ def __next__(self) -> list[str]: if self._file is None: raise StopIteration row = next(self._reader) - # Python's csv.reader returns [] for empty lines; normalize to [''] + # Python's csv.reader returns [] for empty lines; normalize to [None] # to represent a single empty field (consistent with single-value handling) if not row: - return [""] - return row + return [None] + return [None if value == "" else value for value in row] def close(self) -> None: """Close the underlying file object.""" @@ -156,11 +142,6 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._file: Any | None = file_obj self._delimiter = delimiter - @property - def empty_strings_as_null(self) -> bool: - """Whether the result set converts empty strings to None before conversion.""" - return False - @override def __iter__(self) -> AthenaCSVReader: """Iterate over rows in the CSV file.""" diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 58379748b..b080c7375 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -200,9 +200,6 @@ def _fetch(self) -> None: except StopIteration: break - if self._csv_reader.empty_strings_as_null: - row = [None if value == "" else value for value in row] - if col_hints: converted_row = tuple( self._converter.convert(col_type, value, type_hint=hint) diff --git a/tests/pyathena/aio/s3fs/test_cursor.py b/tests/pyathena/aio/s3fs/test_cursor.py index 7f0b0fbb6..625f09628 100644 --- a/tests/pyathena/aio/s3fs/test_cursor.py +++ b/tests/pyathena/aio/s3fs/test_cursor.py @@ -24,7 +24,6 @@ class TestAioS3FSCursor: ("aio_s3fs_cursor", "expected_empty"), [ ({}, ""), - ({"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, ""), ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), ], indirect=["aio_s3fs_cursor"], diff --git a/tests/pyathena/s3fs/test_async_cursor.py b/tests/pyathena/s3fs/test_async_cursor.py index 0926b0f9b..06d082683 100644 --- a/tests/pyathena/s3fs/test_async_cursor.py +++ b/tests/pyathena/s3fs/test_async_cursor.py @@ -29,7 +29,6 @@ class TestAsyncS3FSCursor: ("async_s3fs_cursor", "expected_empty"), [ ({}, ""), - ({"cursor_kwargs": {"csv_reader": AthenaCSVReader}}, ""), ({"cursor_kwargs": {"csv_reader": EmptyStringAsNullCSVReader}}, None), ], indirect=["async_s3fs_cursor"], diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 43d8fabf7..d5f56e677 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -447,10 +447,8 @@ def test_empty_string_with_athena_reader(self): @pytest.mark.parametrize( ("csv_reader", "expected_empty"), [ - ( - EmptyStringAsNullCSVReader, - None, - ), # EmptyStringAsNullCSVReader: empty string becomes None + # EmptyStringAsNullCSVReader: empty string becomes None + (EmptyStringAsNullCSVReader, None), (AthenaCSVReader, ""), # AthenaCSVReader: empty string is preserved ], ) diff --git a/tests/pyathena/s3fs/test_reader.py b/tests/pyathena/s3fs/test_reader.py index 762abeb49..e6d193fde 100644 --- a/tests/pyathena/s3fs/test_reader.py +++ b/tests/pyathena/s3fs/test_reader.py @@ -26,14 +26,13 @@ def test_tab_delimiter(self): rows = list(reader) assert rows == [["a", "b", "c"], ["1", "2", "3"]] - def test_empty_field_returns_empty_string(self): - """EmptyStringAsNullCSVReader returns empty string for both NULL and empty string.""" - # Both ,, (NULL) and ,"", (empty string) become '' + def test_empty_field_returns_none(self): + """EmptyStringAsNullCSVReader returns None for both NULL and empty string.""" + # Both ,, (NULL) and ,"", (empty string) become None data = StringIO('a,,b,"",c\n') reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) - # Standard csv.reader treats both as empty strings - assert rows == [["a", "", "b", "", "c"]] + assert rows == [["a", None, "b", None, "c"]] def test_quoted_field_with_comma(self): data = StringIO('"a,b",c\n') @@ -60,12 +59,12 @@ def test_empty_file(self): assert rows == [] def test_empty_line(self): - """Empty line returns single empty string (e.g., SELECT NULL).""" - # Python's csv.reader returns [] for empty lines; we normalize to [''] + """Empty line returns single None (e.g., SELECT NULL).""" + # Python's csv.reader returns [] for empty lines; we normalize to [None] data = StringIO("\n") reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) - assert rows == [[""]] + assert rows == [[None]] def test_implements_iterator_protocol(self): """EmptyStringAsNullCSVReader implements collections.abc.Iterator.""" diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index d3ffda87e..b07cc033b 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -16,25 +16,12 @@ from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet -from pyathena.util import RetryConfig, override - - -class _InheritedEmptyStringAsNullCSVReader(EmptyStringAsNullCSVReader): - pass - - -class _NullifyingAthenaCSVReader(AthenaCSVReader): - @property - @override - def empty_strings_as_null(self) -> bool: - return True +from pyathena.util import RetryConfig class _TupleCSVReader: """An independent implementation that returns tuples rather than lists.""" - empty_strings_as_null = False - def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._reader = AthenaCSVReader(file_obj, delimiter=delimiter) @@ -48,10 +35,6 @@ def close(self) -> None: self._reader.close() -class _NullifyingTupleCSVReader(_TupleCSVReader): - empty_strings_as_null = True - - def _result_set(data, csv_reader, suffix=".csv", column_types=None, type_hints=None): stream = BytesIO(data.encode()) filesystem = MagicMock() @@ -98,10 +81,7 @@ class TestAthenaS3FSResultSet: (None, ""), (AthenaCSVReader, ""), (EmptyStringAsNullCSVReader, None), - (_InheritedEmptyStringAsNullCSVReader, None), - (_NullifyingAthenaCSVReader, None), (_TupleCSVReader, ""), - (_NullifyingTupleCSVReader, None), ], ) @pytest.mark.parametrize( @@ -125,10 +105,7 @@ def test_reader_null_conversion(self, csv_reader, expected_empty, suffix, data): [ (AthenaCSVReader, ""), (EmptyStringAsNullCSVReader, None), - (_InheritedEmptyStringAsNullCSVReader, None), - (_NullifyingAthenaCSVReader, None), (_TupleCSVReader, ""), - (_NullifyingTupleCSVReader, None), ], ) def test_reader_null_conversion_with_type_hints(self, csv_reader, expected_empty):