Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion google/cloud/bigquery/_pandas_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,7 @@ def types_mapper(arrow_data_type):
# recognized by coverage, hence the pragma. See Issue: #2132
elif (
range_timestamp_dtype is not None
and arrow_data_type.equals( # pragma: NO COVER
and arrow_data_type.equals(
range_timestamp_dtype.pyarrow_dtype
)
):
Expand Down
46 changes: 46 additions & 0 deletions tests/unit/test__pandas_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
from google.cloud.bigquery import _versions_helpers
from google.cloud.bigquery import schema
from google.cloud.bigquery._pandas_helpers import determine_requested_streams
# from google.cloud.bigquery.dtypes import RangeTIMESTAMPDtype # Commented out as it causes ImportError and test uses a mock

pyarrow = _versions_helpers.PYARROW_VERSIONS.try_import()

Expand Down Expand Up @@ -1294,6 +1295,51 @@ def test_dataframe_to_bq_schema_returns_schema_with_pandas_gbq(
assert got is not None


@pytest.mark.skipif(isinstance(pyarrow, mock.Mock), reason="Requires `pyarrow`")
@pytest.mark.skipif(pandas is None, reason="Requires `pandas`") # default_types_mapper is pandas-specific
def test_default_types_mapper_range_timestamp(module_under_test):
"""Test default_types_mapper with a mock RangeTIMESTAMPDtype."""

# Define a duck-typed class that mimics RangeTIMESTAMPDtype's structure
# for the purpose of this test, as the actual class import is problematic.
class MockRangeTSDtype:
def __init__(self):
self.pyarrow_dtype = pyarrow.struct(
[
("start", pyarrow.timestamp("us", tz="UTC")),
("end", pyarrow.timestamp("us", tz="UTC")),
]
)

def __eq__(self, other):
if not isinstance(other, MockRangeTSDtype):
return NotImplemented
return self.pyarrow_dtype == other.pyarrow_dtype

# PyArrow dtype that the mapper will receive
range_timestamp_pyarrow_dtype = pyarrow.struct(
[
("start", pyarrow.timestamp("us", tz="UTC")),
("end", pyarrow.timestamp("us", tz="UTC")),
]
)

# Instantiate the mock RangeTIMESTAMPDtype
mock_range_ts_dtype = MockRangeTSDtype()

# Get the mapper function from the module_under_test fixture
mapper = module_under_test.default_types_mapper(
range_timestamp_dtype=mock_range_ts_dtype
)

# Call the mapper with the created PyArrow struct
result_dtype = mapper(range_timestamp_pyarrow_dtype)

# Assert the result
assert isinstance(result_dtype, MockRangeTSDtype)
assert result_dtype == mock_range_ts_dtype


@pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
def test_dataframe_to_bq_schema_w_named_index(module_under_test, monkeypatch):
monkeypatch.setattr(module_under_test, "pandas_gbq", None)
Expand Down
Loading