Create a Custom pandas compatible Dataframe class #1319

kaxil · 2022-11-28T22:29:47Z

This is an attempt to create a pandas-compatible Dataframe class that has serialize and deserialize methods needed by Airflow 2.5's BaseXCom class

codecov · 2022-11-28T22:56:55Z

Codecov Report

Base: 95.05% // Head: 95.71% // Increases project coverage by +0.65% 🎉

Coverage data is based on head (2a01f45) compared to base (48b8cb0).
Patch has no changes to coverable lines.

❗ Current head 2a01f45 differs from pull request most recent head 4e25c7e. Consider uploading reports for the commit 4e25c7e to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1319      +/-   ##
==========================================
+ Coverage   95.05%   95.71%   +0.65%     
==========================================
  Files          72       19      -53     
  Lines        3400      677    -2723     
  Branches      391       68     -323     
==========================================
- Hits         3232      648    -2584     
+ Misses        102       18      -84     
+ Partials       66       11      -55

Impacted Files	Coverage Δ
python-sdk/src/astro/databases/base.py
python-sdk/src/astro/files/types/csv.py
python-sdk/src/astro/files/types/json.py
python-sdk/src/astro/files/types/ndjson.py
python-sdk/src/astro/files/types/parquet.py
python-sdk/src/astro/settings.py
python-sdk/src/astro/sql/operators/dataframe.py
python-sdk/src/astro/utils/dataframe.py
python-sdk/src/astro/sql/operators/cleanup.py
python-sdk/src/astro/utils/table.py
... and 43 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

phanikumv

Great implementation ! You have covered a lot of ground in this PR 👍

python-sdk/dev/Dockerfile

python-sdk/docs/configurations.rst

Create a pandas-compatible Dataframe class that has `serialize` and `deserialize` methods needed by Airflow 2.5's BaseXCom class

There was a bug in #1319 where I was converting Dataframe to str, but Airflow's XCom expects it to be a dict. This meant that when I tried to check Airflow's Xcom page, I saw the following error in Webserver logs: ``` dev-airflow-webserver-1 | [2022-11-30 15:22:34,709] {app.py:1741} ERROR - Exception on /xcom/list/ [GET] dev-airflow-webserver-1 | Traceback (most recent call last): dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 2525, in wsgi_app dev-airflow-webserver-1 | response = self.full_dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1822, in full_dispatch_request dev-airflow-webserver-1 | rv = self.handle_user_exception(e) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1820, in full_dispatch_request dev-airflow-webserver-1 | rv = self.dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1796, in dispatch_request dev-airflow-webserver-1 | return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/security/decorators.py", line 133, in wraps dev-airflow-webserver-1 | return f(self, *args, **kwargs) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/views.py", line 554, in list dev-airflow-webserver-1 | widgets = self._list() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1164, in _list dev-airflow-webserver-1 | widgets = self._get_list_widget( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1063, in _get_list_widget dev-airflow-webserver-1 | count, lst = self.datamodel.query( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/models/sqla/interface.py", line 471, in query dev-airflow-webserver-1 | query_results = query.all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/query.py", line 2759, in all dev-airflow-webserver-1 | return self._iter().all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1361, in all dev-airflow-webserver-1 | return self._allrows() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 400, in _allrows dev-airflow-webserver-1 | rows = self._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1274, in _fetchall_impl dev-airflow-webserver-1 | return self._real_result._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1686, in _fetchall_impl dev-airflow-webserver-1 | return list(self.iterator) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in chunks dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in <listcomp> dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 984, in _instance dev-airflow-webserver-1 | state.manager.dispatch.load(state, context) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/event/attr.py", line 343, in __call__ dev-airflow-webserver-1 | fn(*args, **kw) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/mapper.py", line 3600, in _event_on_load dev-airflow-webserver-1 | instrumenting_mapper._reconstructor(state.obj()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 125, in init_on_load dev-airflow-webserver-1 | self.value = self.orm_deserialize_value() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 672, in orm_deserialize_value dev-airflow-webserver-1 | return BaseXCom._deserialize_value(self, True) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 654, in _deserialize_value dev-airflow-webserver-1 | return json.loads(result.value.decode("UTF-8"), cls=XComDecoder, object_hook=object_hook) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/__init__.py", line 359, in loads dev-airflow-webserver-1 | return cls(**kw).decode(s) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 337, in decode dev-airflow-webserver-1 | obj, end = self.raw_decode(s, idx=_w(s, 0).end()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 353, in raw_decode dev-airflow-webserver-1 | obj, end = self.scan_once(s, idx) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/utils/json.py", line 255, in orm_object_hook dev-airflow-webserver-1 | for k, v in data.items(): dev-airflow-webserver-1 | AttributeError: 'str' object has no attribute 'items' ```

There was a bug in #1319 where I was converting Dataframe to str, but Airflow's XCom expects it to be a dict as described in http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/latest/concepts/taskflow.html#passing-arbitrary-objects-as-arguments This meant that when I tried to check Airflow's Xcom page, I saw the following error in Webserver logs: ``` dev-airflow-webserver-1 | [2022-11-30 15:22:34,709] {app.py:1741} ERROR - Exception on /xcom/list/ [GET] dev-airflow-webserver-1 | Traceback (most recent call last): dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 2525, in wsgi_app dev-airflow-webserver-1 | response = self.full_dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1822, in full_dispatch_request dev-airflow-webserver-1 | rv = self.handle_user_exception(e) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1820, in full_dispatch_request dev-airflow-webserver-1 | rv = self.dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1796, in dispatch_request dev-airflow-webserver-1 | return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/security/decorators.py", line 133, in wraps dev-airflow-webserver-1 | return f(self, *args, **kwargs) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/views.py", line 554, in list dev-airflow-webserver-1 | widgets = self._list() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1164, in _list dev-airflow-webserver-1 | widgets = self._get_list_widget( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1063, in _get_list_widget dev-airflow-webserver-1 | count, lst = self.datamodel.query( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/models/sqla/interface.py", line 471, in query dev-airflow-webserver-1 | query_results = query.all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/query.py", line 2759, in all dev-airflow-webserver-1 | return self._iter().all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1361, in all dev-airflow-webserver-1 | return self._allrows() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 400, in _allrows dev-airflow-webserver-1 | rows = self._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1274, in _fetchall_impl dev-airflow-webserver-1 | return self._real_result._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1686, in _fetchall_impl dev-airflow-webserver-1 | return list(self.iterator) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in chunks dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in <listcomp> dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 984, in _instance dev-airflow-webserver-1 | state.manager.dispatch.load(state, context) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/event/attr.py", line 343, in __call__ dev-airflow-webserver-1 | fn(*args, **kw) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/mapper.py", line 3600, in _event_on_load dev-airflow-webserver-1 | instrumenting_mapper._reconstructor(state.obj()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 125, in init_on_load dev-airflow-webserver-1 | self.value = self.orm_deserialize_value() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 672, in orm_deserialize_value dev-airflow-webserver-1 | return BaseXCom._deserialize_value(self, True) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 654, in _deserialize_value dev-airflow-webserver-1 | return json.loads(result.value.decode("UTF-8"), cls=XComDecoder, object_hook=object_hook) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/__init__.py", line 359, in loads dev-airflow-webserver-1 | return cls(**kw).decode(s) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 337, in decode dev-airflow-webserver-1 | obj, end = self.raw_decode(s, idx=_w(s, 0).end()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 353, in raw_decode dev-airflow-webserver-1 | obj, end = self.scan_once(s, idx) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/utils/json.py", line 255, in orm_object_hook dev-airflow-webserver-1 | for k, v in data.items(): dev-airflow-webserver-1 | AttributeError: 'str' object has no attribute 'items' ```

This is an attempt to create a pandas-compatible Dataframe class that has `serialize` and `deserialize` methods needed by Airflow 2.5's BaseXCom class (cherry picked from commit a7308e1)

There was a bug in #1319 where I was converting Dataframe to str, but Airflow's XCom expects it to be a dict as described in http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/latest/concepts/taskflow.html#passing-arbitrary-objects-as-arguments This meant that when I tried to check Airflow's Xcom page, I saw the following error in Webserver logs: ``` dev-airflow-webserver-1 | [2022-11-30 15:22:34,709] {app.py:1741} ERROR - Exception on /xcom/list/ [GET] dev-airflow-webserver-1 | Traceback (most recent call last): dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 2525, in wsgi_app dev-airflow-webserver-1 | response = self.full_dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1822, in full_dispatch_request dev-airflow-webserver-1 | rv = self.handle_user_exception(e) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1820, in full_dispatch_request dev-airflow-webserver-1 | rv = self.dispatch_request() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask/app.py", line 1796, in dispatch_request dev-airflow-webserver-1 | return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/security/decorators.py", line 133, in wraps dev-airflow-webserver-1 | return f(self, *args, **kwargs) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/views.py", line 554, in list dev-airflow-webserver-1 | widgets = self._list() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1164, in _list dev-airflow-webserver-1 | widgets = self._get_list_widget( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/baseviews.py", line 1063, in _get_list_widget dev-airflow-webserver-1 | count, lst = self.datamodel.query( dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/flask_appbuilder/models/sqla/interface.py", line 471, in query dev-airflow-webserver-1 | query_results = query.all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/query.py", line 2759, in all dev-airflow-webserver-1 | return self._iter().all() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1361, in all dev-airflow-webserver-1 | return self._allrows() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 400, in _allrows dev-airflow-webserver-1 | rows = self._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1274, in _fetchall_impl dev-airflow-webserver-1 | return self._real_result._fetchall_impl() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/engine/result.py", line 1686, in _fetchall_impl dev-airflow-webserver-1 | return list(self.iterator) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in chunks dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 151, in <listcomp> dev-airflow-webserver-1 | rows = [proc(row) for row in fetch] dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/loading.py", line 984, in _instance dev-airflow-webserver-1 | state.manager.dispatch.load(state, context) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/event/attr.py", line 343, in __call__ dev-airflow-webserver-1 | fn(*args, **kw) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/sqlalchemy/orm/mapper.py", line 3600, in _event_on_load dev-airflow-webserver-1 | instrumenting_mapper._reconstructor(state.obj()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 125, in init_on_load dev-airflow-webserver-1 | self.value = self.orm_deserialize_value() dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 672, in orm_deserialize_value dev-airflow-webserver-1 | return BaseXCom._deserialize_value(self, True) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/models/xcom.py", line 654, in _deserialize_value dev-airflow-webserver-1 | return json.loads(result.value.decode("UTF-8"), cls=XComDecoder, object_hook=object_hook) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/__init__.py", line 359, in loads dev-airflow-webserver-1 | return cls(**kw).decode(s) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 337, in decode dev-airflow-webserver-1 | obj, end = self.raw_decode(s, idx=_w(s, 0).end()) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/json/decoder.py", line 353, in raw_decode dev-airflow-webserver-1 | obj, end = self.scan_once(s, idx) dev-airflow-webserver-1 | File "/usr/local/lib/python3.9/site-packages/airflow/utils/json.py", line 255, in orm_object_hook dev-airflow-webserver-1 | for k, v in data.items(): dev-airflow-webserver-1 | AttributeError: 'str' object has no attribute 'items' ``` (cherry picked from commit 53d6bb7)

kaxil force-pushed the compat-dataframe branch 2 times, most recently from 16d7142 to 731c51d Compare November 28, 2022 22:37

sunank200 added this to the 1.3.0 milestone Nov 29, 2022

kaxil force-pushed the compat-dataframe branch 2 times, most recently from 2a01f45 to ce28b8d Compare November 30, 2022 14:13

kaxil marked this pull request as ready for review November 30, 2022 14:13

kaxil requested review from dimberman, tatiana, utkarsharma2, sunank200, pankajastro, pankajkoti and feluelle as code owners November 30, 2022 14:13

kaxil force-pushed the compat-dataframe branch from ce28b8d to 45f973e Compare November 30, 2022 14:13

phanikumv reviewed Nov 30, 2022

View reviewed changes

python-sdk/dev/Dockerfile Outdated Show resolved Hide resolved

python-sdk/docs/configurations.rst Outdated Show resolved Hide resolved

kaxil force-pushed the compat-dataframe branch from 45f973e to db441a4 Compare November 30, 2022 14:32

Create a Custom pandas compatible Dataframe class

4e25c7e

Create a pandas-compatible Dataframe class that has `serialize` and `deserialize` methods needed by Airflow 2.5's BaseXCom class

kaxil force-pushed the compat-dataframe branch from db441a4 to 4e25c7e Compare November 30, 2022 14:34

phanikumv approved these changes Nov 30, 2022

View reviewed changes

kaxil merged commit a7308e1 into main Nov 30, 2022

kaxil deleted the compat-dataframe branch November 30, 2022 15:10

kaxil mentioned this pull request Nov 30, 2022

Serialize dataframe to dict instead of JSON for XCom #1331

Merged

2 tasks

This was referenced Dec 1, 2022

Test the XCOM Pickling changes as part of Airflow 2.5.0 Rc #1232

Closed

Fix CI failures found while testing SDK with airflow 2.5.0 #1337

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Create a Custom pandas compatible Dataframe class #1319

Create a Custom pandas compatible Dataframe class #1319

kaxil commented Nov 28, 2022 •

edited

Loading

codecov bot commented Nov 28, 2022 •

edited

Loading

phanikumv left a comment

Create a Custom pandas compatible Dataframe class #1319

Create a Custom pandas compatible Dataframe class #1319

Conversation

kaxil commented Nov 28, 2022 • edited Loading

codecov bot commented Nov 28, 2022 • edited Loading

Codecov Report

phanikumv left a comment

Choose a reason for hiding this comment

kaxil commented Nov 28, 2022 •

edited

Loading

codecov bot commented Nov 28, 2022 •

edited

Loading