Skip to content

Commit 70016b8

Browse files
davidheryantodavidheryanto
committed
Add integration tests folder
1 parent 78cdcb3 commit 70016b8

15 files changed

Lines changed: 377 additions & 166 deletions

File tree

‎sdk/python/feast/sdk/utils/bq_util.py‎

Lines changed: 119 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -13,19 +13,25 @@
1313
# limitations under the License.
1414

1515
import os
16+
import tempfile
1617
import time
17-
18+
from datetime import datetime
19+
import pytz
20+
import fastavro
1821
import pandas as pd
22+
from google.cloud import bigquery
1923
from google.cloud.bigquery.client import Client as BQClient
2024
from google.cloud.bigquery.job import ExtractJobConfig, DestinationFormat
2125
from google.cloud.bigquery.table import Table
26+
from google.cloud.exceptions import NotFound
2227
from google.cloud.storage import Client as GCSClient
28+
from google.cloud import storage
2329

2430
from feast.sdk.utils.gs_utils import is_gs_path, split_gs_path, gcs_to_df
2531

2632

2733
def head(client, table, max_rows=10):
28-
'''Get the head of the table. Retrieves rows from the given table at
34+
"""Get the head of the table. Retrieves rows from the given table at
2935
minimum cost
3036
3137
Args:
@@ -37,12 +43,13 @@ def head(client, table, max_rows=10):
3743
3844
Returns:
3945
pandas.DataFrame: dataframe containing the head of rows
40-
'''
46+
"""
4147

4248
rows = client.list_rows(table, max_results=max_rows)
4349
rows = [x for x in rows]
4450
return pd.DataFrame(
45-
data=[list(x.values()) for x in rows], columns=list(rows[0].keys()))
51+
data=[list(x.values()) for x in rows], columns=list(rows[0].keys())
52+
)
4653

4754

4855
def get_table_name(feature_id, storage_spec):
@@ -71,6 +78,108 @@ def get_table_name(feature_id, storage_spec):
7178
return ".".join([project, dataset, table_name])
7279

7380

81+
def get_default_templocation(bigquery_client, project=None):
82+
if project is None:
83+
project = bigquery_client.project
84+
assert isinstance(project, str)
85+
assert len(project) > 0
86+
storage_client = storage.Client()
87+
default_bucket_name = f"feast-templocation-{project}"
88+
try:
89+
storage_client.get_bucket(default_bucket_name)
90+
except NotFound:
91+
print(
92+
f'Default bucket "{default_bucket_name}" not found. Attempting to create it.'
93+
)
94+
storage_client.create_bucket(bucket_name=default_bucket_name, project=project)
95+
return f"gs://{default_bucket_name}"
96+
97+
98+
def query_to_dataframe(
99+
query: str,
100+
bigquery_client: bigquery.Client = None,
101+
storage_client: storage.Client = None,
102+
project: str = None,
103+
templocation: str = None,
104+
) -> pd.DataFrame:
105+
"""
106+
Run a query job on BigQuery and return the result in Pandas DataFrame format
107+
108+
Args:
109+
query: BigQuery query e.g. "SELECT * FROM dataset.table"
110+
bigquery_client:
111+
storage_client:
112+
project: Google Cloud project id
113+
templocation: Google Cloud Storage location to store intermediate files, must start with "gs://"
114+
115+
Returns: Pandas DataFrame of the query result
116+
117+
"""
118+
if isinstance(templocation, str) and not templocation.startswith("gs://"):
119+
raise RuntimeError('templocation must start with "gs://"')
120+
121+
if bigquery_client is None:
122+
bigquery_client = bigquery.Client(project=project)
123+
124+
if project is None:
125+
project = bigquery_client.project
126+
127+
query_job = bigquery_client.query(query, project=project)
128+
query_job_state = ""
129+
130+
while not query_job.done():
131+
if query_job.state != query_job_state:
132+
print(f"Query status: {query_job.state}")
133+
query_job_state = query_job.state
134+
time.sleep(5)
135+
136+
if query_job.state != query_job_state:
137+
print(f"Query status: {query_job.state}")
138+
139+
if query_job.exception():
140+
raise query_job.exception()
141+
142+
if not templocation:
143+
templocation = get_default_templocation(bigquery_client, project=project)
144+
145+
if templocation.endswith("/"):
146+
templocation += templocation[:-1]
147+
148+
destination_uri = (
149+
f"{templocation}/bq-{datetime.now(pytz.utc).strftime('%Y%m%dT%H%M%SZ')}.avro"
150+
)
151+
extract_job_config = bigquery.job.ExtractJobConfig(destination_format="AVRO")
152+
extract_job = bigquery_client.extract_table(
153+
query_job.destination, destination_uri, job_config=extract_job_config
154+
)
155+
156+
while not extract_job.done():
157+
time.sleep(5)
158+
159+
if extract_job.exception():
160+
raise extract_job.exception()
161+
162+
if not storage_client:
163+
storage_client = storage.Client(project=project)
164+
165+
print("Reading query result into DataFrame")
166+
167+
bucket_name, blob_name = (
168+
destination_uri.split("/")[2],
169+
"/".join(destination_uri.split("/")[3:]),
170+
)
171+
bucket = storage_client.get_bucket(bucket_name)
172+
blob = bucket.get_blob(blob_name)
173+
downloaded_avro_filename = tempfile.NamedTemporaryFile().name
174+
blob.download_to_filename(downloaded_avro_filename)
175+
176+
with open(downloaded_avro_filename, "rb") as avro_file:
177+
avro_reader = fastavro.reader(avro_file)
178+
df = pd.DataFrame.from_records(avro_reader)
179+
180+
return df
181+
182+
74183
class TableDownloader:
75184
def __init__(self):
76185
self._bq = None
@@ -88,8 +197,7 @@ def bq(self):
88197
self._bq = BQClient()
89198
return self._bq
90199

91-
def download_table_as_file(self, table_id, dest, staging_location,
92-
file_type):
200+
def download_table_as_file(self, table_id, dest, staging_location, file_type):
93201
"""
94202
Download a bigquery table as file
95203
Args:
@@ -105,14 +213,13 @@ def download_table_as_file(self, table_id, dest, staging_location,
105213
if not is_gs_path(staging_location):
106214
raise ValueError("staging_uri must be a directory in GCS")
107215

108-
temp_file_name = 'temp_{}'.format(int(round(time.time() * 1000)))
216+
temp_file_name = "temp_{}".format(int(round(time.time() * 1000)))
109217
staging_file_path = os.path.join(staging_location, temp_file_name)
110218

111219
job_config = ExtractJobConfig()
112220
job_config.destination_format = file_type
113221
src_table = Table.from_string(table_id)
114-
job = self.bq.extract_table(
115-
src_table, staging_file_path, job_config=job_config)
222+
job = self.bq.extract_table(src_table, staging_file_path, job_config=job_config)
116223

117224
# await completion
118225
job.result()
@@ -137,15 +244,14 @@ def download_table_as_df(self, table_id, staging_location):
137244
if not is_gs_path(staging_location):
138245
raise ValueError("staging_uri must be a directory in GCS")
139246

140-
temp_file_name = 'temp_{}'.format(int(round(time.time() * 1000)))
247+
temp_file_name = "temp_{}".format(int(round(time.time() * 1000)))
141248
staging_file_path = os.path.join(staging_location, temp_file_name)
142249

143250
job_config = ExtractJobConfig()
144251
job_config.destination_format = DestinationFormat.CSV
145252
job = self.bq.extract_table(
146-
Table.from_string(table_id),
147-
staging_file_path,
148-
job_config=job_config)
253+
Table.from_string(table_id), staging_file_path, job_config=job_config
254+
)
149255

150256
# await completion
151257
job.result()

‎sdk/python/setup.py‎

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
# Copyright 2018 The Feast Authors
2-
#
2+
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
55
# You may obtain a copy of the License at
6-
#
6+
#
77
# https://www.apache.org/licenses/LICENSE-2.0
8-
#
8+
#
99
# Unless required by applicable law or agreed to in writing, software
1010
# distributed under the License is distributed on an "AS IS" BASIS,
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -16,23 +16,23 @@
1616
import os
1717
from setuptools import find_packages, setup, Command
1818

19-
NAME = 'Feast'
20-
DESCRIPTION = 'Python sdk for Feast'
21-
URL = 'https://github.com/gojek/feast'
22-
AUTHOR = 'Feast'
23-
REQUIRES_PYTHON = '>=3.6.0'
24-
VERSION = imp.load_source(
25-
'feast.version', os.path.join('feast', 'version.py')).VERSION
19+
NAME = "Feast"
20+
DESCRIPTION = "Python sdk for Feast"
21+
URL = "https://github.com/gojek/feast"
22+
AUTHOR = "Feast"
23+
REQUIRES_PYTHON = ">=3.6.0"
24+
VERSION = imp.load_source("feast.version", os.path.join("feast", "version.py")).VERSION
2625
REQUIRED = [
27-
'google-api-core>=1.7.0',
28-
'google-auth>=1.6.0',
29-
'google-cloud-bigquery>=1.8.0',
30-
'google-cloud-storage>=1.13.0',
31-
'googleapis-common-protos>=1.5.5',
32-
'grpcio>=1.16.1',
33-
'pandas',
34-
'protobuf>=3.0.0',
35-
'PyYAML',
26+
"google-api-core>=1.7.0",
27+
"google-auth>=1.6.0",
28+
"google-cloud-bigquery>=1.8.0",
29+
"google-cloud-storage>=1.13.0",
30+
"googleapis-common-protos>=1.5.5",
31+
"grpcio>=1.16.1",
32+
"pandas",
33+
"protobuf>=3.0.0",
34+
"PyYAML",
35+
"fastavro>=0.21.19"
3636
]
3737

3838
setup(
@@ -42,16 +42,16 @@
4242
author=AUTHOR,
4343
python_requires=REQUIRES_PYTHON,
4444
url=URL,
45-
packages=find_packages(exclude=('tests',)),
45+
packages=find_packages(exclude=("tests",)),
4646
install_requires=REQUIRED,
4747
include_package_data=True,
48-
license='Apache',
48+
license="Apache",
4949
classifiers=[
5050
# Trove classifiers
5151
# Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
52-
'License :: OSI Approved :: Apache Software License',
53-
'Programming Language :: Python',
54-
'Programming Language :: Python :: 3',
55-
'Programming Language :: Python :: 3.6',
56-
]
57-
)
52+
"License :: OSI Approved :: Apache Software License",
53+
"Programming Language :: Python",
54+
"Programming Language :: Python :: 3",
55+
"Programming Language :: Python :: 3.6",
56+
],
57+
)

‎sdk/python/tests/data/austin_bikeshare.bikeshare_stations.avro‎

7.47 KB
Binary file not shown.

0 commit comments

Comments
 (0)