Repository: innobi/pantab
Pull request #243 (2024-01-26 17:46:33)

Ability to return more than just pandas


Patch:
diff --git a/pantab/_reader.py b/pantab/_reader.py
index 36db20b..b3fa8f5 100644
--- a/pantab/_reader.py
+++ b/pantab/_reader.py
@@ -1,7 +1,7 @@
 import pathlib
 import shutil
 import tempfile
-from typing import Optional, Union
+from typing import Union
 
 import pandas as pd
 import pyarrow as pa
@@ -12,22 +12,43 @@ import pantab.src.pantab as libpantab  # type: ignore
 TableType = Union[str, tab_api.Name, tab_api.TableName]
 
 
+def frame_from_hyper_query(
+    source: Union[str, pathlib.Path],
+    query: str,
+    *,
+    return_type="pandas",
+):
+    """See api.rst for documentation."""
+    # Call native library to read tuples from result set
+    capsule = libpantab.read_from_hyper_query(str(source), query)
+    stream = pa.RecordBatchReader._import_from_c_capsule(capsule)
+    tbl = stream.read_all()
+
+    if return_type == "pyarrow":
+        return tbl
+
+    df = tbl.to_pandas(types_mapper=pd.ArrowDtype)
+    return df
+
+
 def frame_from_hyper(
     source: Union[str, pathlib.Path],
     *,
     table: TableType,
-) -> pd.DataFrame:
+    return_type="pandas",
+):
     """See api.rst for documentation"""
     if isinstance(table, (str, tab_api.Name)) or not table.schema_name:
         table = tab_api.TableName("public", table)
 
     query = f"SELECT * FROM {table}"
-    return frame_from_hyper_query(source, query)
+    return frame_from_hyper_query(source, query, return_type=return_type)
 
 
 def frames_from_hyper(
     source: Union[str, pathlib.Path],
-) -> dict[tab_api.TableName, pd.DataFrame]:
+    return_type="pandas",
+):
     """See api.rst for documentation."""
     result: dict[TableType, pd.DataFrame] = {}
 
@@ -45,22 +66,7 @@ def frames_from_hyper(
         result[table] = frame_from_hyper(
             source=source,
             table=table,
+            return_type=return_type,
         )
 
     return result
-
-
-def frame_from_hyper_query(
-    source: Union[str, pathlib.Path],
-    query: str,
-    *,
-    hyper_process: Optional[tab_api.HyperProcess] = None,
-) -> pd.DataFrame:
-    """See api.rst for documentation."""
-    # Call native library to read tuples from result set
-    capsule = libpantab.read_from_hyper_query(str(source), query)
-    stream = pa.RecordBatchReader._import_from_c_capsule(capsule)
-    tbl = stream.read_all()
-    df = tbl.to_pandas(types_mapper=pd.ArrowDtype)
-
-    return df


Test patch:
diff --git a/pantab/tests/__init__.py b/pantab/tests/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/pantab/tests/conftest.py b/pantab/tests/conftest.py
index b42c21a..729c591 100644
--- a/pantab/tests/conftest.py
+++ b/pantab/tests/conftest.py
@@ -3,29 +3,58 @@ import pathlib
 
 import numpy as np
 import pandas as pd
+import pandas.testing as tm
 import pyarrow as pa
 import pytest
 import tableauhyperapi as tab_api
 
 
 def basic_arrow_table():
+    schema = pa.schema(
+        [
+            ("int16", pa.int16()),
+            ("int32", pa.int32()),
+            ("int64", pa.int64()),
+            ("Int16", pa.int16()),
+            ("Int32", pa.int32()),
+            ("Int64", pa.int64()),
+            ("float32", pa.float32()),
+            ("float64", pa.float64()),
+            ("Float32", pa.float32()),
+            ("Float64", pa.float64()),
+            ("bool", pa.bool_()),
+            ("boolean", pa.bool_()),
+            ("date32", pa.date32()),
+            ("datetime64", pa.timestamp("us")),
+            ("datetime64_utc", pa.timestamp("us", "utc")),
+            ("object", pa.large_string()),
+            ("string", pa.string()),
+            ("int16_limits", pa.int16()),
+            ("int32_limits", pa.int32()),
+            ("int64_limits", pa.int64()),
+            ("float32_limits", pa.float32()),
+            ("float64_limits", pa.float64()),
+            ("non-ascii", pa.utf8()),
+            ("binary", pa.binary()),
+            ("time64us", pa.time64("us")),
+        ]
+    )
     tbl = pa.Table.from_arrays(
         [
-            pa.array([1, 6, 0], type=pa.int16()),
-            pa.array([2, 7, 0], type=pa.int32()),
-            pa.array([3, 8, 0], type=pa.int64()),
-            pa.array([1, None, None], type=pa.int16()),
-            pa.array([2, None, None], type=pa.int32()),
-            pa.array([3, None, None], type=pa.int64()),
-            pa.array([4, 9.0, None], type=pa.float32()),
-            pa.array([5, 10.0, None], type=pa.float64()),
-            pa.array([1.0, 1.0, None], type=pa.float32()),
-            pa.array([2.0, 2.0, None], type=pa.float64()),
-            pa.array([True, False, False], type=pa.bool_()),
-            pa.array([True, False, None], type=pa.bool_()),
+            pa.array([1, 6, 0]),
+            pa.array([2, 7, 0]),
+            pa.array([3, 8, 0]),
+            pa.array([1, None, None]),
+            pa.array([2, None, None]),
+            pa.array([3, None, None]),
+            pa.array([4, 9.0, None]),
+            pa.array([5, 10.0, None]),
+            pa.array([1.0, 1.0, None]),
+            pa.array([2.0, 2.0, None]),
+            pa.array([True, False, False]),
+            pa.array([True, False, None]),
             pa.array(
                 [datetime.date(2024, 1, 1), datetime.date(2024, 1, 1), None],
-                type=pa.date32(),
             ),
             pa.array(
                 [
@@ -33,7 +62,6 @@ def basic_arrow_table():
                     datetime.datetime(2019, 1, 1, 0, 0, 0),
                     None,
                 ],
-                type=pa.timestamp("us"),
             ),
             pa.array(
                 [
@@ -41,48 +69,21 @@ def basic_arrow_table():
                     datetime.datetime(2019, 1, 1, 0, 0, 0),
                     None,
                 ],
-                type=pa.timestamp("us", "utc"),
             ),
-            pa.array(["foo", "bar", None], type=pa.large_string()),
-            pa.array(["foo", "bar", None], type=pa.string()),
-            pa.array([-(2**15), 2**15 - 1, 0], type=pa.int16()),
-            pa.array([-(2**31), 2**31 - 1, 0], type=pa.int32()),
-            pa.array([-(2**63), 2**63 - 1, 0], type=pa.int64()),
-            pa.array([-(2**24), 2**24 - 1, None], type=pa.float32()),
-            pa.array([-(2**53), 2**53 - 1, None], type=pa.float64()),
+            pa.array(["foo", "bar", None]),
+            pa.array(["foo", "bar", None]),
+            pa.array([-(2**15), 2**15 - 1, 0]),
+            pa.array([-(2**31), 2**31 - 1, 0]),
+            pa.array([-(2**63), 2**63 - 1, 0]),
+            pa.array([-(2**24), 2**24 - 1, None]),
+            pa.array([-(2**53), 2**53 - 1, None]),
             pa.array(
-                ["\xef\xff\xdc\xde\xee", "\xfa\xfb\xdd\xaf\xaa", None], type=pa.utf8()
+                ["\xef\xff\xdc\xde\xee", "\xfa\xfb\xdd\xaf\xaa", None],
             ),
-            pa.array([b"\xde\xad\xbe\xef", b"\xff\xee", None], type=pa.binary()),
-            pa.array([234, 42, None], type=pa.time64("us")),
-        ],
-        names=[
-            "int16",
-            "int32",
-            "int64",
-            "Int16",
-            "Int32",
-            "Int64",
-            "float32",
-            "float64",
-            "Float32",
-            "Float64",
-            "bool",
-            "boolean",
-            "date32",
-            "datetime64",
-            "datetime64_utc",
-            "object",
-            "string",
-            "int16_limits",
-            "int32_limits",
-            "int64_limits",
-            "float32_limits",
-            "float64_limits",
-            "non-ascii",
-            "binary",
-            "time64us",
+            pa.array([b"\xde\xad\xbe\xef", b"\xff\xee", None]),
+            pa.array([234, 42, None]),
         ],
+        schema=schema,
     )
 
     return tbl
@@ -239,8 +240,42 @@ def frame(request):
     return request.param()
 
 
-@pytest.fixture
-def roundtripped():
+def roundtripped_pyarrow():
+    schema = pa.schema(
+        [
+            ("int16", pa.int16()),
+            ("int32", pa.int32()),
+            ("int64", pa.int64()),
+            ("Int16", pa.int16()),
+            ("Int32", pa.int32()),
+            ("Int64", pa.int64()),
+            ("float32", pa.float64()),
+            ("float64", pa.float64()),
+            ("Float32", pa.float64()),
+            ("Float64", pa.float64()),
+            ("bool", pa.bool_()),
+            ("boolean", pa.bool_()),
+            ("date32", pa.date32()),
+            ("datetime64", pa.timestamp("us")),
+            ("datetime64_utc", pa.timestamp("us", "UTC")),
+            ("object", pa.large_string()),
+            ("string", pa.large_string()),
+            ("int16_limits", pa.int16()),
+            ("int32_limits", pa.int32()),
+            ("int64_limits", pa.int64()),
+            ("float32_limits", pa.float64()),
+            ("float64_limits", pa.float64()),
+            ("non-ascii", pa.large_string()),
+            ("binary", pa.large_binary()),
+            ("time64us", pa.time64("us")),
+        ]
+    )
+    tbl = basic_arrow_table()
+
+    return tbl.cast(schema)
+
+
+def roundtripped_pandas():
     """Roundtripped DataFrames should use arrow dtypes by default"""
     df = basic_dataframe()
     df = df.astype(
@@ -259,7 +294,6 @@ def roundtripped():
             "boolean": "boolean[pyarrow]",
             "datetime64": "timestamp[us][pyarrow]",
             "datetime64_utc": "timestamp[us, UTC][pyarrow]",
-            # "timedelta64": "timedelta64[ns]",
             "object": "large_string[pyarrow]",
             "int16_limits": "int16[pyarrow]",
             "int32_limits": "int32[pyarrow]",
@@ -275,6 +309,17 @@ def roundtripped():
     return df
 
 
+@pytest.fixture(
+    params=[
+        ("pandas", roundtripped_pandas),
+        ("pyarrow", roundtripped_pyarrow),
+    ]
+)
+def roundtripped(request):
+    result_obj = request.param[1]()
+    return (request.param[0], result_obj)
+
+
 @pytest.fixture
 def tmp_hyper(tmp_path):
     """A temporary file name to write / read a Hyper extract from."""
@@ -305,3 +350,82 @@ def table_name(request):
 def datapath():
     """Location of data files in test folder."""
     return pathlib.Path(__file__).parent / "data"
+
+
+class Compat:
+    @staticmethod
+    def assert_frame_equal(result, expected):
+        assert isinstance(result, type(expected))
+        if isinstance(result, pd.DataFrame):
+            tm.assert_frame_equal(result, expected)
+            return
+        elif isinstance(result, pa.Table):
+            assert result.equals(expected, check_metadata=True)
+            return
+        else:
+            raise NotImplementedError("assert_frame_equal not implemented for type")
+
+    @staticmethod
+    def concat_frames(frame1, frame2):
+        assert isinstance(frame1, type(frame2))
+        if isinstance(frame1, pd.DataFrame):
+            return pd.concat([frame1, frame2]).reset_index(drop=True)
+        elif isinstance(frame1, pa.Table):
+            return pa.concat_tables([frame1, frame2])
+        else:
+            raise NotImplementedError("concat_frames not implemented for type")
+
+    @staticmethod
+    def empty_like(frame):
+        if isinstance(frame, pd.DataFrame):
+            return frame.iloc[:0]
+        elif isinstance(frame, pa.Table):
+            return frame.schema.empty_table()
+        else:
+            raise NotImplementedError("empty_like not implemented for type")
+
+    @staticmethod
+    def drop_columns(frame, columns):
+        if isinstance(frame, pd.DataFrame):
+            return frame.drop(columns=columns)
+        elif isinstance(frame, pa.Table):
+            return frame.drop_columns(columns)
+        else:
+            raise NotImplementedError("drop_columns not implemented for type")
+
+    @staticmethod
+    def select_columns(frame, columns):
+        if isinstance(frame, pd.DataFrame):
+            return frame[columns]
+        elif isinstance(frame, pa.Table):
+            return frame.select(columns)
+        else:
+            raise NotImplementedError("select_columns not implemented for type")
+
+    @staticmethod
+    def cast_column_to_type(frame, column, type_):
+        if isinstance(frame, pd.DataFrame):
+            frame[column] = frame[column].astype(type_)
+            return frame
+        elif isinstance(frame, pa.Table):
+            schema = pa.schema([pa.field(column, type_)])
+            return frame.cast(schema)
+        else:
+            raise NotImplementedError("cast_column_to_type not implemented for type")
+
+    @staticmethod
+    def add_non_writeable_column(frame):
+        if isinstance(frame, pd.DataFrame):
+            frame["should_fail"] = pd.Series([list((1, 2))])
+            return frame
+        elif isinstance(frame, pa.Table):
+            new_column = pa.array([[1, 2], None, None])
+            frame = frame.append_column("should_fail", new_column)
+            return frame
+        else:
+            raise NotImplementedError("test not implemented for object")
+
+
+@pytest.fixture()
+def compat():
+    return Compat
diff --git a/pantab/tests/test_reader.py b/pantab/tests/test_reader.py
index 009410e..6824160 100644
--- a/pantab/tests/test_reader.py
+++ b/pantab/tests/test_reader.py
@@ -2,7 +2,6 @@ import json
 
 import pandas as pd
 import pandas.testing as tm
-import pytest
 import tableauhyperapi as tab_api
 
 import pantab
@@ -64,24 +63,6 @@ def test_read_query(frame, tmp_hyper):
     tm.assert_frame_equal(result, expected)
 
 
-def test_empty_read_query(frame, roundtripped, tmp_hyper):
-    """
-    red-green for empty query results
-    """
-    # sql cols need to base case insensitive & unique
-    table_name = "test"
-    pantab.frame_to_hyper(frame, tmp_hyper, table=table_name)
-    query = f"SELECT * FROM {table_name} limit 0"
-
-    if not isinstance(frame, pd.DataFrame):
-        pytest.skip("Need to implement this test properly for pyarrow")
-    expected = pd.DataFrame(columns=frame.columns)
-    expected = expected.astype(roundtripped.dtypes)
-
-    result = pantab.frame_from_hyper_query(tmp_hyper, query)
-    tm.assert_frame_equal(result, expected)
-
-
 def test_read_varchar(tmp_hyper):
     column_name = "VARCHAR Column"
     table_name = tab_api.TableName("public", "table")
diff --git a/pantab/tests/test_roundtrip.py b/pantab/tests/test_roundtrip.py
index 62f513b..aabac47 100644
--- a/pantab/tests/test_roundtrip.py
+++ b/pantab/tests/test_roundtrip.py
@@ -1,25 +1,27 @@
-import pandas as pd
-import pandas.testing as tm
-import pyarrow as pa
 from tableauhyperapi import TableName
 
 import pantab
 
 
-def test_basic(frame, roundtripped, tmp_hyper, table_name, table_mode):
+def test_basic(frame, roundtripped, tmp_hyper, table_name, table_mode, compat):
     # Write twice; depending on mode this should either overwrite or duplicate entries
     pantab.frame_to_hyper(frame, tmp_hyper, table=table_name, table_mode=table_mode)
     pantab.frame_to_hyper(frame, tmp_hyper, table=table_name, table_mode=table_mode)
-    result = pantab.frame_from_hyper(tmp_hyper, table=table_name)
 
-    expected = roundtripped
+    return_type, expected = roundtripped
+    result = pantab.frame_from_hyper(
+        tmp_hyper, table=table_name, return_type=return_type
+    )
+
     if table_mode == "a":
-        expected = pd.concat([expected, expected]).reset_index(drop=True)
+        expected = compat.concat_frames(expected, expected)
 
-    tm.assert_frame_equal(result, expected)
+    compat.assert_frame_equal(result, expected)
 
 
-def test_multiple_tables(frame, roundtripped, tmp_hyper, table_name, table_mode):
+def test_multiple_tables(
+    frame, roundtripped, tmp_hyper, table_name, table_mode, compat
+):
     # Write twice; depending on mode this should either overwrite or duplicate entries
     pantab.frames_to_hyper(
         {table_name: frame, "table2": frame}, tmp_hyper, table_mode=table_mode
@@ -27,11 +29,12 @@ def test_multiple_tables(frame, roundtripped, tmp_hyper, table_name, table_mode)
     pantab.frames_to_hyper(
         {table_name: frame, "table2": frame}, tmp_hyper, table_mode=table_mode
     )
-    result = pantab.frames_from_hyper(tmp_hyper)
 
-    expected = roundtripped
+    return_type, expected = roundtripped
+    result = pantab.frames_from_hyper(tmp_hyper, return_type=return_type)
+
     if table_mode == "a":
-        expected = pd.concat([expected, expected]).reset_index(drop=True)
+        expected = compat.concat_frames(expected, expected)
 
     # some test trickery here
     if not isinstance(table_name, TableName) or table_name.schema_name is None:
@@ -39,26 +42,23 @@ def test_multiple_tables(frame, roundtripped, tmp_hyper, table_name, table_mode)
 
     assert set(result.keys()) == set((table_name, TableName("public", "table2")))
     for val in result.values():
-        tm.assert_frame_equal(val, expected)
+        compat.assert_frame_equal(val, expected)
 
 
-def test_empty_roundtrip(frame, roundtripped, tmp_hyper, table_name, table_mode):
+def test_empty_roundtrip(
+    frame, roundtripped, tmp_hyper, table_name, table_mode, compat
+):
     # object case is by definition vague, so lets punt that for now
-    if isinstance(frame, pd.DataFrame):
-        frame = frame.drop(columns=["object"])
-    elif isinstance(frame, pa.Table):
-        frame = frame.drop_columns(["object"])
-    else:
-        raise NotImplementedError("test not implemented for object")
-
-    if isinstance(frame, pa.Table):
-        empty = frame.schema.empty_table()
-    else:
-        empty = frame.iloc[:0]
+    frame = compat.drop_columns(frame, ["object"])
+    empty = compat.empty_like(frame)
     pantab.frame_to_hyper(empty, tmp_hyper, table=table_name, table_mode=table_mode)
     pantab.frame_to_hyper(empty, tmp_hyper, table=table_name, table_mode=table_mode)
-    result = pantab.frame_from_hyper(tmp_hyper, table=table_name)
 
-    expected = roundtripped.iloc[:0]
-    expected = expected.drop(columns=["object"])
-    tm.assert_frame_equal(result, expected)
+    return_type, expected = roundtripped
+    result = pantab.frame_from_hyper(
+        tmp_hyper, table=table_name, return_type=return_type
+    )
+
+    expected = compat.drop_columns(expected, ["object"])
+    expected = compat.empty_like(expected)
+    compat.assert_frame_equal(result, expected)
diff --git a/pantab/tests/test_writer.py b/pantab/tests/test_writer.py
index 48f9b98..0ad11dd 100644
--- a/pantab/tests/test_writer.py
+++ b/pantab/tests/test_writer.py
@@ -2,7 +2,6 @@ import re
 from datetime import datetime, timezone
 
 import pandas as pd
-import pyarrow as pa
 import pytest
 from tableauhyperapi import Connection, CreateMode, HyperProcess, Telemetry
 
@@ -27,42 +26,29 @@ def test_bad_table_mode_raises(frame, tmp_hyper):
     "new_dtype,hyper_type_name", [("int64", "BIGINT"), ("float", "DOUBLE PRECISION")]
 )
 def test_append_mode_raises_column_dtype_mismatch(
-    new_dtype, hyper_type_name, frame, tmp_hyper, table_name
+    new_dtype, hyper_type_name, frame, tmp_hyper, table_name, compat
 ):
-    if isinstance(frame, pd.DataFrame):
-        frame = frame[["int16"]].copy()
-    else:
-        frame = frame.select(["int16"])
+    frame = compat.select_columns(frame, ["int16"])
     pantab.frame_to_hyper(frame, tmp_hyper, table=table_name)
 
-    if isinstance(frame, pd.DataFrame):
-        frame["int16"] = frame["int16"].astype(new_dtype)
-    elif isinstance(frame, pa.Table):
-        schema = pa.schema([pa.field("int16", new_dtype)])
-        frame = frame.cast(schema)
-    else:
-        raise NotImplementedError("test not implemented for object")
-
+    frame = compat.cast_column_to_type(frame, "int16", new_dtype)
     msg = f"Column type mismatch at index 0; new: {hyper_type_name} old: SMALLINT"
     with pytest.raises(ValueError, match=msg):
         pantab.frame_to_hyper(frame, tmp_hyper, table=table_name, table_mode="a")
 
 
-def test_append_mode_raises_ncolumns_mismatch(frame, tmp_hyper, table_name):
+def test_append_mode_raises_ncolumns_mismatch(frame, tmp_hyper, table_name, compat):
     pantab.frame_to_hyper(frame, tmp_hyper, table=table_name)
 
-    if isinstance(frame, pd.DataFrame):
-        frame = frame.drop(columns=["int16"])
-    elif isinstance(frame, pa.Table):
-        frame = frame.drop_columns(["int16"])
-    else:
-        raise NotImplementedError("test not implemented for object")
+    frame = compat.drop_columns(frame, ["int16"])
     msg = "Number of columns"
     with pytest.raises(ValueError, match=msg):
         pantab.frame_to_hyper(frame, tmp_hyper, table=table_name, table_mode="a")
 
 
-def test_failed_write_doesnt_overwrite_file(frame, tmp_hyper, monkeypatch, table_mode):
+def test_failed_write_doesnt_overwrite_file(
+    frame, tmp_hyper, monkeypatch, table_mode, compat
+):
     pantab.frame_to_hyper(
         frame,
         tmp_hyper,
@@ -72,14 +58,7 @@ def test_failed_write_doesnt_overwrite_file(frame, tmp_hyper, monkeypatch, table
     last_modified = tmp_hyper.stat().st_mtime
 
     # Pick a dtype we know will fail
-    if isinstance(frame, pd.DataFrame):
-        frame["should_fail"] = pd.Series([list((1, 2))])
-    elif isinstance(frame, pa.Table):
-        new_column = pa.array([[1, 2], None, None])
-        frame = frame.append_column("should_fail", new_column)
-    else:
-        raise NotImplementedError("test not implemented for object")
-
+    frame = compat.add_non_writeable_column(frame)
     # Try out our write methods
     msg = "Unsupported Arrow type"
     with pytest.raises(ValueError, match=msg):