Skip to content

Fix NaN in read_rows #378

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
elephantum wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix NaN in read_rows #378

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 31 additions & 4 deletions datapipe/meta/sql_meta.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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]

Expand Down Expand Up @@ -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]

Expand Down
13 changes: 11 additions & 2 deletions datapipe/store/database.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,15 +272,23 @@ 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)

return pd.concat(res, ignore_index=True)

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,
Expand All @@ -296,4 +304,5 @@ def read_rows_meta_pseudo_df(
sql,
con=con,
chunksize=chunksize,
dtype_backend="pyarrow",
)
47 changes: 46 additions & 1 deletion datapipe/store/tests/abstract.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions datapipe/tests/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
1 change: 1 addition & 0 deletions examples/datatable_batch_transform/app.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ def count_tbl(
return pd.read_sql_query(
sql,
con=con,
dtype_backend="pyarrow",
)


Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -23,6 +23,7 @@ dependencies = [
"opentelemetry-instrumentation-sqlalchemy",
"click>=7.1.2",
"rich>=13.3.2",
"pyarrow[pyarrow]>=21.0.0",
]

[project.optional-dependencies]
Expand All @@ -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"
Expand Down
Loading