Skip to content

Commit

Permalink
Add check for BQ_EXPORT_PARQUET_FILE_SIZE_MB value
Browse files Browse the repository at this point in the history
  • Loading branch information
KarolisKont committed Nov 7, 2023
1 parent 7374e78 commit db1aec6
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion sdk/python/feast/infra/offline_stores/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -575,8 +575,16 @@ def to_remote_storage(self) -> List[str]:

parquet_file_size_mb = os.getenv("BQ_EXPORT_PARQUET_FILE_SIZE_MB")
if parquet_file_size_mb is not None and parquet_file_size_mb.isdigit():
parquet_file_size_mb_int = int(parquet_file_size_mb)
if parquet_file_size_mb_int >= 1000:
raise ValueError(
"The value for the BQ_EXPORT_PARQUET_FILE_SIZE_MB environment variable cannot "
"exceed 1000; however, it was set to %s.",
parquet_file_size_mb_int,
)

table_size_in_mb = self.client.get_table(table).num_bytes / 1024 / 1024
number_of_files = max(1, table_size_in_mb // int(parquet_file_size_mb))
number_of_files = max(1, table_size_in_mb // parquet_file_size_mb_int)
destination_uris = [
f"{self._gcs_path}/{n:0>12}.parquet" for n in range(number_of_files)
]
Expand Down

0 comments on commit db1aec6

Please sign in to comment.