Implement query_foundry_sql(return_type='polars')#115
Conversation
There was a problem hiding this comment.
Pull request overview
This PR implements support for returning Polars DataFrames directly from query_foundry_sql() methods across the codebase, eliminating the need for intermediate Arrow Table conversion.
Changes:
- Added
"polars"as a valid return type for SQL queries in type definitions and method signatures - Implemented native Polars DataFrame conversion in
FoundrySqlServerClient.query_foundry_sql() - Simplified
Dataset.to_polars()to use the new return type instead of converting from Arrow
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| libs/foundry-dev-tools/src/foundry_dev_tools/utils/api_types.py | Added "polars" to SQLReturnType literal and updated documentation |
| libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py | Implemented polars return type handling in query_foundry_sql with proper import checking |
| libs/foundry-dev-tools/src/foundry_dev_tools/resources/dataset.py | Updated type hints and simplified to_polars() to use new return type |
| libs/foundry-dev-tools/src/foundry_dev_tools/foundry_api_client.py | Updated method signatures and documentation for polars support |
| tests/integration/resources/test_dataset.py | Added test coverage for polars return type in Dataset.query_foundry_sql |
| tests/integration/clients/test_foundry_sql_server.py | Added test coverage for polars return type in FoundrySqlServerClient |
| docs/examples/dataset.md | Updated documentation example to show direct polars return type usage |
| docs/dev/contribute.md | Updated development environment setup instructions for Python version |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py
Show resolved
Hide resolved
libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py
Outdated
Show resolved
Hide resolved
libs/foundry-dev-tools/src/foundry_dev_tools/clients/foundry_sql_server.py
Outdated
Show resolved
Hide resolved
| ).json() | ||
| assert_in_literal(return_type, SQLReturnType, "return_type") | ||
|
|
||
| if return_type == "raw": |
There was a problem hiding this comment.
Moving this up will break the fallback that is tested in test_legacy_fallback defined in foundry-dev-tools/tests/integration/clients/test_foundry_sql_server.py
I introduced this so data can be queried in constrained environments where arrow is not installed.
There was a problem hiding this comment.
Good catch. I also realized this is related to the try...catch as well. In the original code, if return_type != "raw" and it runs into either FoundrySqlSerializationFormatNotImplementedError or ImportError and the return_type is also != "arrow", then it will just swallow the exception and fall back to query_foundry_sql_legacy.
My suggestion: I will only add the if return_type == polars part and keep everything else as it was before but add some comments to document this behavior. What do you think?
There was a problem hiding this comment.
Changes:
- got rid of the dead code
- documented the fallback logic to legacy client (I now understand that ImportErrors are indirectly raised in FakeModule when trying to access an attribute of the fake module)
- implemented polars in queray_foundry_sql_legacy so that the final uncaught ImportError will be raised at this point
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Summary
Extend
query_foundry_sql()in the classesDataset,FoundryRestClient, andFoundrySqlServerClientwith the return_type polars.Checklist