sql_table() query_adapter_callback function with custom extraction query fails with KeyError during pipeline.extract() (pyarrow backend)
#2374
Comments
acaruso7
changed the title
sql_table() query_adapter_callback function with custom extraction query fails with KeyError (pyarrow backend)sql_table() query_adapter_callback function with custom extraction query fails with KeyError during pipeline.extract() (pyarrow backend)
|
@acaruso7 are you able to check the newest version of yield row_tuples_to_arrow(
partition,
columns=_add_missing_columns(self.columns, columns),
tz=backend_kwargs.get("tz", "UTC"),
)
the correct way of implementing this is to use def _add_change_tracking_columns(table: Table) -> None:
required_columns = [
("_dlt_sys_change_version", sa.BigInteger, {"nullable": True}),
("_dlt_deleted", sa.Text, {"default": None, "nullable": True}),
]
for col_name, col_type, col_kwargs in required_columns:
if col_name not in table.c:
table.append_column(
sa.Column(col_name, col_type, **col_kwargs) # type: ignore[arg-type]
)defines two additional columns that are computed columns added by the query adapter |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
dlt version
1.5.0
Describe the problem
https://dlthub-community.slack.com/archives/C04DQA7JJN6/p1740596179916109
I am using
sql_table()source withpyarrowbackend to extract data from a mysql db with a custom extraction query. The custom extraction query is passed toquery_adapter_callbackas a string, and selects a subset of columns from source table, as well as some additional derived columns (example below)Upon calling
pipeline.extract(), adlt.extract.exceptions.ResourceExtractionErroris thrown. The stack trace indicates aKeyErrorfor one of the table columns which is not defined in custom extraction query result setMore details with full reproducable example and stack trace below
Note: I am following along with the example in the documentation here
Expected behavior
query_adapter_callbackwith custom extraction query should extract only the columns defined in query into pyarrow table.KeyErrorshould not occurSteps to reproduce
Full reproducible example using public mysql db and duckdb destination:
Full stack trace:
I dug into the internals of the
arrow_helpers.row_tuples_to_arrowfunction a bit and noticed some odd behaviordlt/dlt/common/libs/pyarrow.py
Line 575 in e8c5e9b
The value of the
columnsvariable iswhich is not what I would expect. This output contains all columns from the source table, as well as my additional derived column,
my_custom_columnThe value of
len(rows[0])is3, which is what I would expect - 3 columns in results, which aligns to my custom extraction queryThe value of
columnar.keys()isdict_keys(['author_id', 'name', 'last_name'])- not what I would expect.last_nameis not included in the result set of my custom queryIt seems something may be going wrong with the operation here, where
columnsare zipped withrowsdlt/dlt/common/libs/pyarrow.py
Lines 595 to 597 in e8c5e9b
This may lead to a
KeyErrorlaterOperating system
Linux
Runtime environment
Docker, Docker Compose
Python version
3.10
dlt data source
sql_table()withpyarrowbackenddlt destination
No response
Other deployment details
No response
Additional information
Separately from the issue described here, I wonder whether I am going about this the right way. Is this the preferred approach for extracting custom column sets (including derived columns) from source tables?
An additional requirement I should mention for my own use case is that the custom extraction query must be passed the the
query_adapter_callbackas a raw SQL string - it cannot be built using native sqlalchemy syntax as I am re-using the same query elsewhereThe text was updated successfully, but these errors were encountered: