fix: preserve SQLDataset filters in pandas and Arrow conversions

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)
 

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:

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]