From 6ab13022a7909c15cd4ed4e7e1e80f53c36f9754 Mon Sep 17 00:00:00 2001 From: Andrey Tatarinov Date: Sat, 14 Mar 2026 15:48:00 +0400 Subject: [PATCH] Fix NaN in read_rows; migrate to pyarrow dtypes --- datapipe/meta/sql_meta.py | 35 +++++++++++++++-- datapipe/store/database.py | 13 ++++++- datapipe/store/tests/abstract.py | 47 ++++++++++++++++++++++- datapipe/tests/util.py | 3 ++ examples/datatable_batch_transform/app.py | 1 + pyproject.toml | 4 +- 6 files changed, 94 insertions(+), 9 deletions(-) diff --git a/datapipe/meta/sql_meta.py b/datapipe/meta/sql_meta.py index d4cdbde1..2c7e79cf 100644 --- a/datapipe/meta/sql_meta.py +++ b/datapipe/meta/sql_meta.py @@ -156,12 +156,25 @@ def get_metadata(self, idx: IndexDF | None = None, include_deleted: bool = False with self.dbconn.con.begin() as con: if idx is None: sql = self._build_metadata_query(sql, idx, include_deleted) - return cast(MetadataDF, pd.read_sql_query(sql, con=con)) + return cast( + MetadataDF, + pd.read_sql_query( + sql, + con=con, + dtype_backend="pyarrow", + ), + ) for chunk_idx in self._chunk_idx_df(idx): chunk_sql = self._build_metadata_query(sql, chunk_idx, include_deleted) - res.append(pd.read_sql_query(chunk_sql, con=con)) + res.append( + pd.read_sql_query( + chunk_sql, + con=con, + dtype_backend="pyarrow", + ) + ) if len(res) > 0: return cast(MetadataDF, pd.concat(res)) @@ -222,6 +235,7 @@ def get_existing_idx(self, idx: IndexDF | None = None) -> IndexDF: res_df: DataDF = pd.read_sql_query( sql, con=con, + dtype_backend="pyarrow", ) return data_to_index(res_df, self.primary_keys) @@ -347,7 +361,14 @@ def get_stale_idx( with self.dbconn.con.begin() as con: return cast( Iterator[IndexDF], - list(pd.read_sql_query(sql, con=con, chunksize=1000)), + list( + pd.read_sql_query( + sql, + con=con, + chunksize=1000, + dtype_backend="pyarrow", + ) + ), ) def get_agg_cte( @@ -503,7 +524,12 @@ def get_full_process_ids( def alter_res_df(): with ds.meta_dbconn.con.begin() as con: - for df in pd.read_sql_query(u1, con=con, chunksize=chunk_size): + for df in pd.read_sql_query( + u1, + con=con, + chunksize=chunk_size, + dtype_backend="pyarrow", + ): assert isinstance(df, pd.DataFrame) df = df[self.transform_keys] @@ -540,6 +566,7 @@ def get_change_list_process_ids( table_changes_df = pd.read_sql_query( sql, con=con, + dtype_backend="pyarrow", ) table_changes_df = table_changes_df[self.transform_keys] diff --git a/datapipe/store/database.py b/datapipe/store/database.py index 15f461f3..2899e32c 100644 --- a/datapipe/store/database.py +++ b/datapipe/store/database.py @@ -272,7 +272,11 @@ def read_rows(self, idx: IndexDF | None = None) -> pd.DataFrame: with self.dbconn.con.begin() as con: for chunk_idx in self._chunk_idx_df(idx): chunk_sql = sql_apply_idx_filter_to_table(sql, self.data_table, self.primary_keys, chunk_idx) - chunk_df = pd.read_sql_query(chunk_sql, con=con) + chunk_df = pd.read_sql_query( + chunk_sql, + con=con, + dtype_backend="pyarrow", + ) res.append(chunk_df) @@ -280,7 +284,11 @@ def read_rows(self, idx: IndexDF | None = None) -> pd.DataFrame: else: with self.dbconn.con.begin() as con: - return pd.read_sql_query(sql, con=con) + return pd.read_sql_query( + sql, + con=con, + dtype_backend="pyarrow", + ) def read_rows_meta_pseudo_df( self, @@ -296,4 +304,5 @@ def read_rows_meta_pseudo_df( sql, con=con, chunksize=chunksize, + dtype_backend="pyarrow", ) diff --git a/datapipe/store/tests/abstract.py b/datapipe/store/tests/abstract.py index d5a64d76..03ca9b33 100644 --- a/datapipe/store/tests/abstract.py +++ b/datapipe/store/tests/abstract.py @@ -6,7 +6,7 @@ import cloudpickle import pandas as pd import pytest -from sqlalchemy import Column, String +from sqlalchemy import Column, Integer, String from datapipe.run_config import RunConfig from datapipe.store.table_store import TableStore @@ -49,6 +49,51 @@ def test_get_schema( assert store.get_schema() == schema + def test_multiple_keys_with_none(self, store_maker: TableStoreMaker) -> None: + data_df = pd.DataFrame( + { + "id1": [1, 2, 3], + "id2": ["a", "b", "c"], + "value": [10, 20, 30], + } + ) + schema: list[Column] = [ + Column("id1", Integer, primary_key=True), + Column("id2", String(100), primary_key=True), + Column("value", Integer), + ] + + store = store_maker(schema) + store.insert_rows(data_df) + + assert_df_equal( + store.read_rows( + data_to_index( + pd.DataFrame.from_dict({"id1": [2, -1], "id2": ["b", None]}).convert_dtypes( + dtype_backend="pyarrow" + ), + ["id1", "id2"], + ) + ), + pd.DataFrame.from_dict({"id1": [2], "id2": ["b"], "value": [20]}), + index_cols=["id1", "id2"], + ) + + assert_df_equal( + store.read_rows( + data_to_index( + pd.DataFrame.from_dict({"id1": [2, None], "id2": ["b", "z"]}).convert_dtypes( + dtype_backend="pyarrow" + ), + ["id1", "id2"], + ) + ), + pd.DataFrame.from_dict({"id1": [2], "id2": ["b"], "value": [20]}), + index_cols=["id1", "id2"], + ) + + assert_ts_contains(store, data_df) + @pytest.mark.parametrize("data_df,schema", DATA_PARAMS) def test_write_read_rows( self, diff --git a/datapipe/tests/util.py b/datapipe/tests/util.py index f39f2342..ae5086e1 100644 --- a/datapipe/tests/util.py +++ b/datapipe/tests/util.py @@ -15,6 +15,9 @@ def assert_idx_equal(a, b): def assert_df_equal(a: pd.DataFrame, b: pd.DataFrame, index_cols=["id"]) -> bool: + a = a.convert_dtypes(dtype_backend="pyarrow") + b = b.convert_dtypes(dtype_backend="pyarrow") + a = a.set_index(index_cols) b = b.set_index(index_cols) diff --git a/examples/datatable_batch_transform/app.py b/examples/datatable_batch_transform/app.py index fbe4f457..b0ec25f9 100644 --- a/examples/datatable_batch_transform/app.py +++ b/examples/datatable_batch_transform/app.py @@ -76,6 +76,7 @@ def count_tbl( return pd.read_sql_query( sql, con=con, + dtype_backend="pyarrow", ) diff --git a/pyproject.toml b/pyproject.toml index 461f8e08..ae386678 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ readme = "README.md" requires-python = ">=3.10,<3.15" dependencies = [ "fsspec>=2021.11.1", - "pandas>=1.2.0", + "pandas>=1.2.0, <3.0", "numpy>=1.21.0, <3.0", "SQLAlchemy>=1.4.0, <3.0.0", "psycopg2_binary>=2.8.4", @@ -23,6 +23,7 @@ dependencies = [ "opentelemetry-instrumentation-sqlalchemy", "click>=7.1.2", "rich>=13.3.2", + "pyarrow[pyarrow]>=21.0.0", ] [project.optional-dependencies] @@ -40,7 +41,6 @@ ray = ["ray[default]>=2.5.0,<3"] gcp = ["opentelemetry-exporter-gcp-trace"] elastic = ["elasticsearch>=8.17.1"] neo4j = ["neo4j"] -pyarrow = ["pyarrow[pyarrow]>=21.0.0"] [project.urls] Repository = "https://github.com/epoch8/datapipe"