feat(python): Support pre_execution_query parameter from connectorx #21288
Conversation
jsjasonseba
requested review from
ritchie46,
c-peters,
alexander-beedie,
MarcoGorelli and
reswqa
as code owners
github-actions
bot
added
enhancement
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #21288 +/- ##
==========================================
+ Coverage 79.82% 80.03% +0.20%
==========================================
Files 1596 1604 +8
Lines 228565 231339 +2774
Branches 2608 2642 +34
==========================================
+ Hits 182462 185142 +2680
- Misses 45507 45588 +81
- Partials 596 609 +13 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
@alexander-beedie can you review this one? |
|
@alexander-beedie got time to review this? I gave a few suggestions but haven't heard back from the PR creator, so if they've disappeared I can just take this one over--this would help my team quite a bit, as our postgres implementation (Greenplum) keeps the query optimizer disabled by default and this is the only way to enable it. |
|
Hi @mcrumiller , which suggestions are you talking about? In the meantime, I used the pre_execution_query param in connectorx read_sql and set the destination to arrow, and then use the from_arrow in polars. This might be a viable workaround for you too |
There was a problem hiding this comment.
Sorry @jsjasonseba, I used the review mechanism instead of normal comments and forgot to click the submit button. You should see my recommendations now.
| ) | ||
| elif engine == "adbc": | ||
| if not isinstance(query, str): | ||
| msg = "only a single SQL query string is accepted for adbc" | ||
| raise ValueError(msg) | ||
| if pre_execution_query: | ||
| msg = ( | ||
| "the 'connectorx' engine does not support use of `pre_execution_query`" |
There was a problem hiding this comment.
This should say adbc, not connectorx.
| query="SELECT * FROM test_data", | ||
| protocol="sqlite", | ||
| errclass=ValueError, | ||
| errmsg="the 'connectorx' engine does not support use of `pre_execution_query`", |
| partition_num=partition_num, | ||
| protocol=protocol, | ||
| pre_execution_query=pre_execution_query, | ||
| ) |
There was a problem hiding this comment.
Minor refactor, but this section would be a bit cleaner as:
try:
if pre_execution_query:
if parse_version(cx.__version__) < (0, 4, 2):
msg = "pre_execution_query is only supported in connectorx version 0.4.2 or later."
raise ValueError(msg)
return_type = "arrow2"
pre_execution_args = {"pre_execution_query": pre_execution_query}
else:
return_type = "arrow"
pre_execution_args = {}
tbl = cx.read_sql(
conn=connection_uri,
query=query,
return_type=return_type,
partition_on=partition_on,
partition_range=partition_range,
partition_num=partition_num,
protocol=protocol,
**pre_execution_args,
)There was a problem hiding this comment.
Agreed. Appreciate your feedback!
| pre_execution_args = {} | ||
| else: | ||
| return_type = "arrow" | ||
| pre_execution_args = {"pre_execution_query": pre_execution_query} |
There was a problem hiding this comment.
Thanks, this is more correct than my suggestion.
There was a problem hiding this comment.
yes, made some slight correction. But I got your point clearly. Thanks!
Implements #21287