Feat: Support pandas in BigQuery cache #597
Conversation
📝 WalkthroughWalkthroughThis pull request introduces modifications to the caching and data reading mechanisms in the Airbyte project. The changes primarily focus on enhancing the data retrieval process for different cache types, with a specific emphasis on BigQuery integration. The modifications include adding a new private method for reading SQL tables into Pandas DataFrames, updating import statements, and adding a new dependency to support BigQuery data handling. Changes
Possibly related PRs
Suggested Labels
Suggested Reviewers
Hey there! 👋 I noticed some interesting changes in the caching and data retrieval mechanisms. The new The addition of ✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (6)
airbyte/caches/bigquery.py (3)
55-60: Consider adding a docstring.Would it help to include a docstring here explaining
_read_to_pandas_dataframeusage, the newchunksizeparameter, and any exceptions that might arise? Wdyt?
61-64: Validate or document popped keyword arguments.We're discarding
"con"and"schema". Is there a reason to ignore them silently versus logging a warning or raising an error to inform the caller? Wdyt?
82-83: Confirm returning unaltered DataFrame.We’re returning the DataFrame as-is. Any chance we'd want a copy or some form of read-only structure for safety? Wdyt?
tests/integration_tests/cloud/test_cloud_sql_reads.py (1)
72-72: Explore partial validation of DataFrame contents.We assert
pandas_df.shape == (100, 20), which is great for row-column checks. Might it be useful to validate a few columns or cells to ensure data integrity? Wdyt?tests/integration_tests/test_all_cache_types.py (1)
161-165: Double-check boundary conditions for chunk batching.We're splitting the Arrow dataset into chunks of size 10 and expecting 20 batches for 200 rows. Do we want a test verifying off-by-one or partial chunks if the row count isn't divisible by 10? Wdyt?
airbyte/caches/base.py (1)
184-191: Consider adding a docstring to explain the method's purpose?The method looks good, but adding a docstring would help explain its role as an extension point for different cache implementations, wdyt?
def _read_to_pandas_dataframe( self, table_name: str, con: Engine, **kwargs, ) -> pd.DataFrame: + """Read a SQL table into a pandas DataFrame. + + This method serves as an extension point for different cache implementations + to customize how they read data into pandas DataFrames. + + Args: + table_name: Name of the table to read + con: SQLAlchemy engine connection + **kwargs: Additional arguments passed to the underlying read implementation + + Returns: + pd.DataFrame: The table data as a pandas DataFrame + """ return pd.read_sql_table(table_name, con=con, **kwargs)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
airbyte/caches/base.py(5 hunks)airbyte/caches/bigquery.py(2 hunks)pyproject.toml(1 hunks)tests/integration_tests/cloud/test_cloud_sql_reads.py(2 hunks)tests/integration_tests/test_all_cache_types.py(1 hunks)
🧰 Additional context used
📓 Learnings (1)
airbyte/caches/bigquery.py (1)
Learnt from: aaronsteers
PR: airbytehq/PyAirbyte#281
File: airbyte/caches/bigquery.py:40-43
Timestamp: 2024-11-10T16:30:14.198Z
Learning: The `BigQueryCache.get_arrow_dataset` method should have a docstring that correctly states the reason for the `NotImplementedError` as BigQuery not supporting `to_arrow`, instead of incorrectly mentioning `pd.read_sql_table`.
🔇 Additional comments (9)
airbyte/caches/bigquery.py (4)
22-23: Ensure required version constraints for new libraries.We've introduced
pandasandpandas_gbq. Are you confident that ourpyproject.tomlor similar config file suitably pins or restricts their versions? This helps avoid unforeseen compatibility issues. Wdyt?
37-38: Check Python version for type hinting.Importing
Iteratorfromcollections.abcis a good practice. However, the|union syntax inchunksize: int | Nonerequires Python 3.10+. Is that guaranteed by our environment orpyproject.toml? Wdyt?
65-71: Handle potentialread_gbqexceptions.If
pandas_gbq.read_gbqfails (e.g., invalid credentials or query errors), do we want to catch and handle those exceptions more gracefully, perhaps with logging? Wdyt?
74-77: Confirm the returned format frompandas_gbq.
pandas_gbq.read_gbqtypically returns aDataFrame. Is there a scenario where it doesn't, and do we need to handle it differently? Wdyt?tests/integration_tests/cloud/test_cloud_sql_reads.py (1)
175-177: Consider verifying DataFrame columns.We again confirm
(100, 20). Would you like to confirm column names or types here, ensuring consistent schemas across jobs? Wdyt?airbyte/caches/base.py (3)
31-31: LGTM! Good type safety practice.The explicit import of
Enginefrom sqlalchemy.engine enhances type safety for the new pandas integration.
71-73: LGTM! Nice formatting.The multi-line method signature improves readability.
206-210: LGTM! Nice refactoring.Good job extracting the common pandas reading logic into
_read_to_pandas_dataframe. The chunked reading approach inget_arrow_datasetis particularly nice for handling large tables efficiently.Also applies to: 223-223
pyproject.toml (1)
32-32: LGTM! Good version constraint.The addition of
pandas-gbqwith>=0.26.1constraint is appropriate for the BigQuery pandas integration. The minimum version is recent enough to ensure good feature support while allowing for future updates.
| # Return chunks as iterator if chunksize is provided | ||
| if chunksize is not None: | ||
| return (result[i : i + chunksize] for i in range(0, len(result), chunksize)) | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Revisit chunking performance.
For very large tables, returning chunked slices of the DataFrame might still be memory-intensive as the entire DataFrame is loaded first. Would you consider a chunked read directly from pandas_gbq instead, if available? Wdyt?
There was a problem hiding this comment.
pandas_gbq doesn't support it :(
jscheel
force-pushed
the
js/add-pandas-support-to-bigquery-cache
branch
from
7ca8d05 to
56513e8
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
airbyte/caches/bigquery.py (2)
61-63: Consider handling unused kwargs differently.Instead of popping unused kwargs, would it be clearer to explicitly define the expected kwargs and pass only those to pandas_gbq.read_gbq? This could make the API more explicit and prevent silent parameter drops. Wdyt?
55-60: Add type hints for kwargs.Would it be helpful to add type hints for the expected kwargs? This could make it clearer what parameters are supported. Maybe something like:
- def _read_to_pandas_dataframe( - self, - table_name: str, - chunksize: int | None = None, - **kwargs, - ) -> pd.DataFrame | Iterator[pd.DataFrame]: + def _read_to_pandas_dataframe( + self, + table_name: str, + chunksize: int | None = None, + **kwargs: Any, # or more specific types based on pandas_gbq.read_gbq parameters + ) -> pd.DataFrame | Iterator[pd.DataFrame]:Wdyt?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
airbyte/caches/base.py(5 hunks)airbyte/caches/bigquery.py(2 hunks)pyproject.toml(2 hunks)tests/integration_tests/cloud/test_cloud_sql_reads.py(2 hunks)tests/integration_tests/test_all_cache_types.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- tests/integration_tests/cloud/test_cloud_sql_reads.py
- tests/integration_tests/test_all_cache_types.py
- airbyte/caches/base.py
- pyproject.toml
🔇 Additional comments (2)
airbyte/caches/bigquery.py (2)
78-81: Consider memory-efficient chunking implementation.I noticed this was mentioned in a previous review. The current implementation still loads the entire DataFrame into memory before chunking. Would you consider using pandas_gbq's native chunking capabilities if available? This could help with memory usage for large tables. Wdyt?
22-23: Verify pandas_gbq version compatibility.The pandas_gbq integration looks good! However, should we specify a minimum version requirement for pandas_gbq to ensure compatibility? This could help prevent potential issues with older versions. Wdyt?
| # Read the table using pandas_gbq | ||
| credentials = Credentials.from_service_account_file(self.credentials_path) | ||
| result = pandas_gbq.read_gbq( | ||
| f"{self.project_name}.{self.dataset_name}.{table_name}", | ||
| project_id=self.project_name, | ||
| credentials=credentials, | ||
| **kwargs, |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add error handling for credentials loading.
The credentials loading could fail for various reasons (file not found, invalid credentials, etc.). Should we add some error handling here? Maybe something like:
- credentials = Credentials.from_service_account_file(self.credentials_path)
+ try:
+ credentials = Credentials.from_service_account_file(self.credentials_path)
+ except FileNotFoundError as e:
+ raise ValueError(f"Credentials file not found at {self.credentials_path}") from e
+ except Exception as e:
+ raise ValueError(f"Failed to load credentials: {str(e)}") from eWdyt?
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Read the table using pandas_gbq | |
| credentials = Credentials.from_service_account_file(self.credentials_path) | |
| result = pandas_gbq.read_gbq( | |
| f"{self.project_name}.{self.dataset_name}.{table_name}", | |
| project_id=self.project_name, | |
| credentials=credentials, | |
| **kwargs, | |
| # Read the table using pandas_gbq | |
| try: | |
| credentials = Credentials.from_service_account_file(self.credentials_path) | |
| except FileNotFoundError as e: | |
| raise ValueError(f"Credentials file not found at {self.credentials_path}") from e | |
| except Exception as e: | |
| raise ValueError(f"Failed to load credentials: {str(e)}") from e | |
| result = pandas_gbq.read_gbq( | |
| f"{self.project_name}.{self.dataset_name}.{table_name}", | |
| project_id=self.project_name, | |
| credentials=credentials, | |
| **kwargs, |
jscheel
force-pushed
the
js/add-pandas-support-to-bigquery-cache
branch
from
56513e8 to
dc810cb
Compare
jscheel
changed the title
|
/fix-pr
|
|
/test-pr
|
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
airbyte/caches/bigquery.py (1)
61-64: Handling of Unused Keyword ArgumentsThe use of
kwargs.pop("con", None)and similar pops is practical but might be less explicit. Have you considered defining and documenting which keyword arguments are supported? This could make the code easier to maintain. What do you think?
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
poetry.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
airbyte/caches/base.py(5 hunks)airbyte/caches/bigquery.py(2 hunks)pyproject.toml(2 hunks)tests/integration_tests/cloud/test_cloud_sql_reads.py(2 hunks)tests/integration_tests/test_all_cache_types.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- tests/integration_tests/cloud/test_cloud_sql_reads.py
- tests/integration_tests/test_all_cache_types.py
- pyproject.toml
- airbyte/caches/base.py
🧰 Additional context used
🪛 GitHub Actions: Run Linters
airbyte/caches/bigquery.py
[error] 55-55: Function is missing a type annotation for one or more arguments [no-untyped-def]
[error] 55-55: Signature of "_read_to_pandas_dataframe" incompatible with supertype "CacheBase" [override]
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: Pytest (No Creds)
- GitHub Check: Pytest (Fast)
🔇 Additional comments (3)
airbyte/caches/bigquery.py (3)
66-66: Enhance Credentials Loading with Error HandlingDirectly invoking
Credentials.from_service_account_file(self.credentials_path)could lead to uninformative errors if the credentials file is missing or invalid. Previously, a similar suggestion was made. Would you consider wrapping this in a try/except block to provide clearer error messages? For example:- credentials = Credentials.from_service_account_file(self.credentials_path) + try: + credentials = Credentials.from_service_account_file(self.credentials_path) + except FileNotFoundError as e: + raise ValueError(f"Credentials file not found at {self.credentials_path}") from e + except Exception as e: + raise ValueError(f"Failed to load credentials: {str(e)}") from ewdyt?
78-81: Verify the Chunking MechanismThe generator expression for splitting the DataFrame into chunks is concise and clear. However, have you considered potential edge cases (e.g., when the DataFrame is empty or when
chunksizeexceeds the DataFrame's length)? It might be useful to add a quick review of these scenarios to ensure robustness. wdyt?
55-60: Signature and Typing for_read_to_pandas_dataframeThe new method signature lays the groundwork for pandas-based table reading, but the linter highlights two concerns:
- Missing type annotation for
**kwargs: Would you consider annotating**kwargs(for example, using**kwargs: Anyafter importingAnyfromtyping) to improve clarity and satisfy the type checker?- Override signature mismatch with
CacheBase: The pipeline error indicates that the signature here is incompatible with whatCacheBaseexpects. Could you verify if the parent interface should be updated or if adjustments are needed in this override?For example, you might update the signature as follows:
+from typing import Any ... -def _read_to_pandas_dataframe( - self, - table_name: str, - chunksize: int | None = None, - **kwargs, -) -> pd.DataFrame | Iterator[pd.DataFrame]: +def _read_to_pandas_dataframe( + self, + table_name: str, + chunksize: int | None = None, + **kwargs: Any, +) -> pd.DataFrame | Iterator[pd.DataFrame]:wdyt?
🧰 Tools
🪛 GitHub Actions: Run Linters
[error] 55-55: Function is missing a type annotation for one or more arguments [no-untyped-def]
[error] 55-55: Signature of "_read_to_pandas_dataframe" incompatible with supertype "CacheBase" [override]
Uses the pandas_gpb library to support fetching BigQuery tables to pandas DataFrames.
Summary by CodeRabbit
New Features
Dependencies
pandas-gbqpackage to project dependencies.Improvements
Bug Fixes