Exemplo n.º 1
0
def test_to_arrow_no_pyarrow_raises_import_error(mut, class_under_test,
                                                 mock_client, monkeypatch):
    monkeypatch.setattr(mut, "pyarrow", None)
    arrow_schema = _bq_to_arrow_schema(SCALAR_COLUMNS)
    read_session = _generate_arrow_read_session(arrow_schema)
    arrow_batches = _bq_to_arrow_batches(SCALAR_BLOCKS, arrow_schema)
    reader = class_under_test(arrow_batches, mock_client,
                              bigquery_storage_v1beta1.types.StreamPosition(),
                              {})

    with pytest.raises(ImportError):
        reader.to_arrow(read_session)

    with pytest.raises(ImportError):
        reader.rows(read_session).to_arrow()

    with pytest.raises(ImportError):
        next(reader.rows(read_session).pages).to_arrow()
Exemplo n.º 2
0
def test_to_arrow_w_scalars_arrow(class_under_test):
    arrow_schema = _bq_to_arrow_schema(SCALAR_COLUMNS)
    read_session = _generate_arrow_read_session(arrow_schema)
    arrow_batches = _bq_to_arrow_batches(SCALAR_BLOCKS, arrow_schema)
    reader = class_under_test(arrow_batches, mock_client,
                              bigquery_storage_v1beta1.types.StreamPosition(),
                              {})
    actual_table = reader.to_arrow(read_session)
    expected_table = pyarrow.Table.from_batches(
        _bq_to_arrow_batch_objects(SCALAR_BLOCKS, arrow_schema))
    assert actual_table == expected_table