diff --git a/airbyte/datasets/_sql.py b/airbyte/datasets/_sql.py index c2065c800..fe4f99c86 100644 --- a/airbyte/datasets/_sql.py +++ b/airbyte/datasets/_sql.py @@ -6,6 +6,9 @@ import warnings from typing import TYPE_CHECKING, Any, Literal, cast +import pandas as pd +import pyarrow as pa +import pyarrow.dataset as ds from overrides import overrides from sqlalchemy import and_, func, select, text @@ -23,7 +26,6 @@ if TYPE_CHECKING: from collections.abc import Iterator - from pandas import DataFrame from pyarrow.dataset import Dataset from sqlalchemy import Table from sqlalchemy.sql import ClauseElement @@ -105,15 +107,23 @@ def __len__(self) -> int: return cast("int", self._length) - def to_pandas(self) -> DataFrame: - return self._cache.get_pandas_dataframe(self._stream_name) + def to_pandas(self) -> pd.DataFrame: + return pd.read_sql_query(self._query_statement, self._cache.get_sql_engine()) def to_arrow( self, *, max_chunk_size: int = DEFAULT_ARROW_MAX_CHUNK_SIZE, ) -> Dataset: - return self._cache.get_arrow_dataset(self._stream_name, max_chunk_size=max_chunk_size) + pandas_chunks = pd.read_sql_query( + self._query_statement, + self._cache.get_sql_engine(), + chunksize=max_chunk_size, + ) + arrow_tables = [pa.Table.from_pandas(pandas_chunk) for pandas_chunk in pandas_chunks] + # Filtering can expose an all-null first chunk. Reconcile chunk schemas rather + # than forcing later values into the first chunk's inferred null types. + return ds.dataset(pa.concat_tables(arrow_tables, promote_options="permissive")) def with_filter(self, *filter_expressions: ClauseElement | str) -> SQLDataset: """Filter the dataset by a set of column values. @@ -180,7 +190,7 @@ def __init__( ) @overrides - def to_pandas(self) -> DataFrame: + def to_pandas(self) -> pd.DataFrame: """Return the underlying dataset data as a pandas DataFrame.""" return self._cache.get_pandas_dataframe(self._stream_name) diff --git a/tests/integration_tests/test_source_test_fixture.py b/tests/integration_tests/test_source_test_fixture.py index 824e282a3..e47010242 100644 --- a/tests/integration_tests/test_source_test_fixture.py +++ b/tests/integration_tests/test_source_test_fixture.py @@ -720,6 +720,48 @@ def test_cached_dataset_filter() -> None: ) +@pytest.mark.parametrize("conversion", ["pandas", "arrow"]) +@pytest.mark.parametrize( + ("filter_expression", "expected_values"), + [("column2 = 1", [1]), ("column2 < 0", [])], + ids=["matching", "empty"], +) +def test_filtered_dataset_conversions( + conversion: str, + filter_expression: str, + expected_values: list[int], +) -> None: + """Conversions must retain the same filter as SQL iteration, including empty results.""" + source = ab.get_source("source-test", config={"apiKey": "test"}) + source.select_all_streams() + + with ab.new_local_cache() as cache: + result = source.read(cache) + original = result["stream1"] + filtered = original.with_filter(filter_expression) + assert [record["column2"] for record in filtered] == expected_values + + if conversion == "pandas": + converted = filtered.to_pandas() + assert converted["column2"].tolist() == expected_values + assert set(converted.columns) == { + "column1", + "column2", + "sometimes_object", + *AB_INTERNAL_COLUMNS, + } + else: + converted = filtered.to_arrow(max_chunk_size=1) + assert converted.to_table().column("column2").to_pylist() == expected_values + assert set(converted.schema.names) == { + "column1", + "column2", + "sometimes_object", + *AB_INTERNAL_COLUMNS, + } + assert all(batch.num_rows <= 1 for batch in converted.to_batches()) + + def test_lazy_dataset_from_source( expected_test_stream_data: dict[str, list[dict[str, str | int]]], ) -> None: diff --git a/tests/unit_tests/test_sql_datasets.py b/tests/unit_tests/test_sql_datasets.py new file mode 100644 index 000000000..0aa607259 --- /dev/null +++ b/tests/unit_tests/test_sql_datasets.py @@ -0,0 +1,28 @@ +# Copyright (c) 2023 Airbyte, Inc., all rights reserved. +"""SQL dataset conversion regressions using a real local cache.""" + +from pathlib import Path + +from sqlalchemy import column, select, text + +from airbyte.caches import DuckDBCache +from airbyte.datasets import CachedDataset, SQLDataset + + +def test_sql_dataset_arrow_promotes_nullable_chunks(tmp_path: Path) -> None: + """A filtered first batch of nulls must not fix the schema to null types.""" + with DuckDBCache(db_path=tmp_path / "cache.duckdb", cache_dir=tmp_path) as cache: + cache.execute_sql([ + "CREATE TABLE main.sample (id INTEGER, s VARCHAR, n INTEGER)", + "INSERT INTO main.sample VALUES (1, 'first', 10), (2, NULL, NULL), (3, 'third', 30)", + ]) + original = CachedDataset(cache, "sample", stream_configuration=False) + assert original.to_arrow(max_chunk_size=1).count_rows() == 3 + + query = select("*").select_from(text("main.sample")).where(column("id") >= 2) + filtered = SQLDataset(cache, "sample", query, stream_configuration=False) + expected = [{"id": 2, "s": None, "n": None}, {"id": 3, "s": "third", "n": 30}] + assert list(filtered) == expected + converted = filtered.to_arrow(max_chunk_size=1) + assert converted.to_table().to_pylist() == expected + assert [batch.num_rows for batch in converted.to_batches()] == [1, 1]