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..8de356ff4 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,31 @@ 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. + +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..13d7443b6 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,20 @@ 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. +It parses result files with Python's `csv` module and returns `None` for both NULL and quoted empty fields. ```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 +209,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 +221,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) +``` + +`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 + +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. + The cursor converts each value by its column type and treats `None` as NULL. +- Implement `close()` to close the underlying stream. + +The reader is used only for result files; managed query results are read through the Athena API. + ### 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..53da4da4d 100644 --- a/pyathena/s3fs/reader.py +++ b/pyathena/s3fs/reader.py @@ -10,27 +10,49 @@ 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. - 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. + Reader classes accept a text stream and a delimiter. They yield sequences + 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.""" + ... - Use this reader when you need backward compatibility with the behavior - where empty strings are treated the same as NULL values. + 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 | None]]): + """CSV reader that reads both NULL and empty string as None. + + 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 = DefaultCSVReader(StringIO(',"",text')) + >>> 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 @@ -48,16 +70,16 @@ def __init__(self, file_obj: Any, delimiter: str = ",") -> None: self._reader = csv.reader(file_obj, delimiter=delimiter) @override - def __iter__(self) -> DefaultCSVReader: + 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. @@ -65,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.""" @@ -77,7 +99,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 +128,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: diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index 0bc268fc9..b080c7375 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,18 @@ 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 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..625f09628 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,20 @@ class TestAioS3FSCursor: + @pytest.mark.parametrize( + ("aio_s3fs_cursor", "expected_empty"), + [ + ({}, ""), + ({"cursor_kwargs": {"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 +108,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..06d082683 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,20 @@ class TestAsyncS3FSCursor: + @pytest.mark.parametrize( + ("async_s3fs_cursor", "expected_empty"), + [ + ({}, ""), + ({"cursor_kwargs": {"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 +203,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..d5f56e677 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,8 @@ def test_empty_string_with_athena_reader(self): @pytest.mark.parametrize( ("csv_reader", "expected_empty"), [ - (DefaultCSVReader, None), # DefaultCSVReader: empty string becomes None + # EmptyStringAsNullCSVReader: empty string becomes None + (EmptyStringAsNullCSVReader, None), (AthenaCSVReader, ""), # AthenaCSVReader: empty string is preserved ], ) @@ -455,10 +456,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 +504,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 +570,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..e6d193fde 100644 --- a/tests/pyathena/s3fs/test_reader.py +++ b/tests/pyathena/s3fs/test_reader.py @@ -8,75 +8,74 @@ 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.""" - # 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 = DefaultCSVReader(data, delimiter=",") + 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') - 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 == [] 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 = DefaultCSVReader(data, delimiter=",") + reader = EmptyStringAsNullCSVReader(data, delimiter=",") rows = list(reader) - assert rows == [[""]] + assert rows == [[None]] 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 +84,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..b07cc033b --- /dev/null +++ b/tests/pyathena/s3fs/test_result_set.py @@ -0,0 +1,151 @@ +# 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 + + +class _TupleCSVReader: + """An independent implementation that returns tuples rather than lists.""" + + 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() + + +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 + + +class TestAthenaS3FSResultSet: + @pytest.mark.parametrize( + ("csv_reader", "expected_empty"), + [ + (None, ""), + (AthenaCSVReader, ""), + (EmptyStringAsNullCSVReader, None), + (_TupleCSVReader, ""), + ], + ) + @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), + (_TupleCSVReader, ""), + ], + ) + 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,)]