datapipe-core 0.14.2.dev2__tar.gz → 0.15.0.dev1__tar.gz

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/.gitignore

@@ -1,5 +1,4 @@
 .DS_Store
-.idea
 .mypy_cache/
 .pytest_cache/
 .venv/

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datapipe-core
-Version: 0.14.2.dev2
+Version: 0.15.0.dev1
 Summary: `datapipe` is a realtime incremental ETL library for Python application
 Project-URL: Repository, https://github.com/epoch8/datapipe
 Author-email: Andrey Tatarinov <a@tatarinov.co>
@@ -24,8 +24,6 @@ Requires-Dist: sqlalchemy<3.0.0,>=2.0.0
 Requires-Dist: tqdm-loggable<0.3,>=0.2
 Requires-Dist: traceback-with-variables<3,>=2.0.4
 Requires-Dist: types-tqdm>=4.67.0.20241221
-Provides-Extra: elastic
-Requires-Dist: elasticsearch>=8.17.1; extra == 'elastic'
 Provides-Extra: excel
 Requires-Dist: openpyxl>=3.0.7; extra == 'excel'
 Requires-Dist: xlrd>=2.0.1; extra == 'excel'

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/datapipe/store/database.py

@@ -14,7 +14,7 @@ from sqlalchemy.sql.expression import delete, select
 
 from datapipe.run_config import RunConfig
 from datapipe.sql_util import sql_apply_idx_filter_to_table, sql_apply_runconfig_filter
-from datapipe.store.table_store import TableStore, TableStoreCaps
+from datapipe.store.table_store import TableStore
 from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema, OrmTable, TAnyDF
 
 logger = logging.getLogger("datapipe.store.database")
@@ -121,14 +121,6 @@ class MetaKey(SchemaItem):
 
 
 class TableStoreDB(TableStore):
-    caps = TableStoreCaps(
-        supports_delete=True,
-        supports_get_schema=True,
-        supports_read_all_rows=True,
-        supports_read_nonexistent_rows=True,
-        supports_read_meta_pseudo_df=True,
-    )
-
     def __init__(
         self,
         dbconn: Union["DBConn", str],
@@ -144,9 +136,9 @@ class TableStoreDB(TableStore):
 
         if orm_table is not None:
             assert name is None, "name should be None if orm_table is provided"
-            assert data_sql_schema is None, (
-                "data_sql_schema should be None if orm_table is provided"
-            )
+            assert (
+                data_sql_schema is None
+            ), "data_sql_schema should be None if orm_table is provided"
 
             orm_table__table = orm_table.__table__  # type: ignore
             self.data_table = cast(Table, orm_table__table)
@@ -169,12 +161,12 @@ class TableStoreDB(TableStore):
             ]
 
         else:
-            assert name is not None, (
-                "name should be provided if data_table is not provided"
-            )
-            assert data_sql_schema is not None, (
-                "data_sql_schema should be provided if data_table is not provided"
-            )
+            assert (
+                name is not None
+            ), "name should be provided if data_table is not provided"
+            assert (
+                data_sql_schema is not None
+            ), "data_sql_schema should be provided if data_table is not provided"
 
             self.name = name
 

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/datapipe/store/filedir.py

@@ -15,7 +15,7 @@ from PIL import Image
 from sqlalchemy import Column, Integer, String
 
 from datapipe.run_config import RunConfig
-from datapipe.store.table_store import TableStore, TableStoreCaps
+from datapipe.store.table_store import TableStore
 from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema
 
 
@@ -103,9 +103,9 @@ def _pattern_to_attrnames(pat: str) -> List[str]:
     assert len(attrnames) > 0, "The scheme is not valid."
     if len(attrnames) >= 2:
         duplicates_attrnames = list(duplicates(attrnames))
-        assert len(duplicates_attrnames) == 0, (
-            f"Some keys are repeated: {duplicates_attrnames}. Rename them."
-        )
+        assert (
+            len(duplicates_attrnames) == 0
+        ), f"Some keys are repeated: {duplicates_attrnames}. Rename them."
 
     return attrnames
 
@@ -161,14 +161,6 @@ class Replacer:
 
 
 class TableStoreFiledir(TableStore):
-    caps = TableStoreCaps(
-        supports_delete=True,
-        supports_get_schema=False,
-        supports_read_all_rows=True,
-        supports_read_nonexistent_rows=False,
-        supports_read_meta_pseudo_df=True,
-    )
-
     def __init__(
         self,
         filename_pattern: Union[str, Path],
@@ -286,8 +278,7 @@ class TableStoreFiledir(TableStore):
                 for attrname in self.attrnames
             ]
         self.attrname_to_cls = {
-            column.name: type_to_cls[type(column.type)]
-            for column in self.primary_schema  # type: ignore
+            column.name: type_to_cls[type(column.type)] for column in self.primary_schema  # type: ignore
         }
 
     def get_primary_schema(self) -> DataSchema:
@@ -331,9 +322,9 @@ class TableStoreFiledir(TableStore):
         """
         _, filepath = fsspec.core.split_protocol(filepath)
         m = re.match(self.filename_match, filepath)
-        assert m is not None, (
-            f"Filepath {filepath} does not match the pattern {self.filename_match}"
-        )
+        assert (
+            m is not None
+        ), f"Filepath {filepath} does not match the pattern {self.filename_match}"
 
         data = {}
         for attrname in self.attrnames:

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/datapipe/store/redis.py

@@ -7,7 +7,7 @@ from redis.cluster import RedisCluster
 from sqlalchemy import Column
 
 from datapipe.store.database import MetaKey
-from datapipe.store.table_store import TableStore, TableStoreCaps
+from datapipe.store.table_store import TableStore
 from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema, data_to_index
 
 
@@ -24,30 +24,14 @@ def _to_itertuples(df: DataDF, colnames):
 
 
 class RedisStore(TableStore):
-    caps = TableStoreCaps(
-        supports_delete=True,
-        supports_get_schema=False,
-        supports_read_all_rows=False,
-        supports_read_nonexistent_rows=False,  # TODO check
-        supports_read_meta_pseudo_df=False,
-    )
-
     def __init__(
-        self,
-        connection: str,
-        name: str,
-        data_sql_schema: List[Column],
-        cluster_mode: bool = False,
+        self, connection: str, name: str, data_sql_schema: List[Column], cluster_mode: bool = False
     ) -> None:
         self.connection = connection
         if not cluster_mode:
-            self.redis_connection: Union[Redis, RedisCluster] = Redis.from_url(
-                connection, decode_responses=True
-            )
+            self.redis_connection: Union[Redis, RedisCluster] = Redis.from_url(connection, decode_responses=True)
         else:
-            self.redis_connection = RedisCluster.from_url(
-                connection, decode_responses=True
-            )
+            self.redis_connection = RedisCluster.from_url(connection, decode_responses=True)
 
         self.name = name
         self.data_sql_schema = data_sql_schema

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/datapipe/store/table_store.py

@@ -1,5 +1,4 @@
 from abc import ABC
-from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator, List, Optional, Union
 
@@ -10,18 +9,7 @@ from datapipe.run_config import RunConfig
 from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema, data_to_index
 
 
-@dataclass
-class TableStoreCaps:
-    supports_delete: bool
-    supports_get_schema: bool
-    supports_read_all_rows: bool
-    supports_read_nonexistent_rows: bool
-    supports_read_meta_pseudo_df: bool
-
-
 class TableStore(ABC):
-    caps: TableStoreCaps
-
     def get_primary_schema(self) -> DataSchema:
         raise NotImplementedError
 
@@ -58,14 +46,6 @@ class TableStore(ABC):
 
 
 class TableDataSingleFileStore(TableStore):
-    caps = TableStoreCaps(
-        supports_delete=True,
-        supports_get_schema=False,
-        supports_read_all_rows=True,
-        supports_read_nonexistent_rows=True,
-        supports_read_meta_pseudo_df=True,
-    )
-
     def __init__(
         self,
         filename: Union[Path, str, None] = None,

{datapipe_core-0.14.2.dev2 → datapipe_core-0.15.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "datapipe-core"
-version = "0.14.2-dev.2"
+version = "0.15.0-dev.1"
 description = "`datapipe` is a realtime incremental ETL library for Python application"
 authors = [{ name = "Andrey Tatarinov", email = "a@tatarinov.co" }]
 readme = "README.md"
@@ -40,7 +40,6 @@ redis = ["redis>=4.3.4,<5"]
 qdrant = ["qdrant-client>=1.1.7,<2"]
 ray = ["ray[default]>=2.5.0,<3"]
 gcp = ["opentelemetry-exporter-gcp-trace"]
-elastic = ["elasticsearch>=8.17.1"]
 
 [project.urls]
 Repository = "https://github.com/epoch8/datapipe"
@@ -63,9 +62,6 @@ dev = [
     "types-PyYAML",
 ]
 
-[tool.ruff]
-line-length = 120
-
 [tool.hatch.build.targets.sdist]
 include = ["datapipe", "datapipe/py.typed"]
 

datapipe_core-0.14.2.dev2/datapipe/store/elastic.py

@@ -1,211 +0,0 @@
-import base64
-import hashlib
-from typing import Any, Dict, Iterable, Iterator, List, Optional, TypedDict
-
-import pandas as pd
-from elastic_transport import ObjectApiResponse
-from elasticsearch import Elasticsearch, helpers
-from sqlalchemy import Column
-
-from datapipe.run_config import RunConfig
-from datapipe.store.database import MetaKey
-from datapipe.store.table_store import TableStore, TableStoreCaps
-from datapipe.types import DataDF, DataSchema, IndexDF, MetaSchema
-
-
-def get_elastic_id(keys: Iterable[Any], length: int = 20) -> str:
-    concatenated_keys = "".join([str(key) for key in keys])
-    needed_bytes = length * 3 // 4
-    hash_object = hashlib.sha256(concatenated_keys.encode("utf-8"))
-    hash_bytes = hash_object.digest()[:needed_bytes]
-    base64_encoded_id = base64.urlsafe_b64encode(hash_bytes).decode("utf-8")
-    return base64_encoded_id[:length]
-
-
-def _to_itertuples(df: DataDF, colnames):
-    return list(df[colnames].itertuples(index=False, name=None))
-
-
-def remap_dict_keys(data: Dict[str, Any], key_name_remapping: Dict[str, str]) -> Dict[str, Any]:
-    return {key_name_remapping.get(key, key): value for key, value in data.items()}
-
-
-class ElasticStoreState(TypedDict):
-    index: str
-    data_sql_schema: List[Column]
-    es_kwargs: Dict[str, Any]
-    key_name_remapping: Optional[Dict[str, str]]
-    mapping: Optional[dict]
-
-
-class ElasticStore(TableStore):
-    caps = TableStoreCaps(
-        supports_delete=True,
-        supports_read_all_rows=True,
-        supports_get_schema=True,
-        supports_read_meta_pseudo_df=True,
-        supports_read_nonexistent_rows=False,
-    )
-
-    def __init__(
-        self,
-        index: str,
-        data_sql_schema: List[Column],
-        es_kwargs: Dict[str, Any],
-        key_name_remapping: Optional[Dict[str, str]] = None,
-        mapping: Optional[dict] = None,
-    ) -> None:
-        self.index = index
-        self.data_sql_schema = data_sql_schema
-        self.key_name_remapping = key_name_remapping or {}
-        self.primary_key_columns = [column.name for column in self.data_sql_schema if column.primary_key]
-        self.value_key_columns = [column.name for column in self.data_sql_schema if not column.primary_key]
-        self.primary_key_column_rename = "_dtp_orig_{pk}"
-        self.mapping = mapping
-
-        self.es_kwargs = es_kwargs
-        self.es_client = Elasticsearch(**es_kwargs)
-
-    def __getstate__(self) -> ElasticStoreState:
-        return {
-            "index": self.index,
-            "data_sql_schema": self.data_sql_schema,
-            "es_kwargs": self.es_kwargs,
-            "mapping": self.mapping,
-            "key_name_remapping": self.key_name_remapping,
-        }
-
-    def __setstate__(self, state: ElasticStoreState) -> None:
-        ElasticStore.__init__(
-            self,
-            index=state["index"],
-            data_sql_schema=state["data_sql_schema"],
-            es_kwargs=state["es_kwargs"],
-            key_name_remapping=state["key_name_remapping"],
-            mapping=state["mapping"],
-        )
-
-    def insert_rows(self, df: DataDF) -> None:
-        if df.empty:
-            return
-
-        # previously index was implicitly created by the bulk api call, now explicit with mapping
-        index_exists = self.es_client.indices.exists(index=self.index)
-        if not index_exists:
-            self.es_client.indices.create(index=self.index, body=self.mapping)
-
-        actions = []
-        for row in df.to_dict(orient="records"):  # type: ignore
-            # I need to retrieve data in chunks and restore the ids
-            # here ids are hashed, so I need to store the original ide values in _source
-            # since I cannot store the _id in source (ES will not validate request), I rename these fields
-            row_data: Dict[str, Any] = {key: row[key] for key in self.value_key_columns}
-            row_id = get_elastic_id([row[key] for key in self.primary_key_columns])
-            row_data = remap_dict_keys(row_data, self.key_name_remapping)
-            row_data.update(
-                {self.primary_key_column_rename.format(pk=key): row[key] for key in self.primary_key_columns}
-            )
-            actions.append({"_index": self.index, "_source": row_data, "_id": row_id})
-
-        helpers.bulk(client=self.es_client, actions=actions, refresh=True)
-
-    def read_rows(self, idx: Optional[IndexDF] = None) -> DataDF:
-        if idx is not None:
-            if idx.empty:
-                return pd.DataFrame(columns=[column.name for column in self.data_sql_schema])
-
-            key_rows = _to_itertuples(idx, self.primary_key_columns)
-            rows_ids = [get_elastic_id(row) for row in key_rows]
-            data = self.es_client.mget(index=self.index, body={"ids": rows_ids}, source=True)
-            data = data["docs"]
-        else:
-            # elasticsearch has default limit of 10000 per query
-            # I assume you will use the read_rows_meta_pseudo_df for larger result sets
-            data = self.es_client.search(index=self.index, query={"match_all": {}}, size=10000)
-            data = data["hits"]["hits"]
-
-        remapping_with_primary_keys = {
-            **self.key_name_remapping,
-            **{
-                self.primary_key_column_rename.format(pk=primary_key): f"{primary_key}"
-                for primary_key in self.primary_key_columns
-            },
-        }
-        result = [
-            remap_dict_keys(item["_source"], remapping_with_primary_keys)  # type: ignore
-            for item in data
-        ]
-        if result:
-            return pd.DataFrame(result)
-        else:
-            return pd.DataFrame(columns=self.primary_key_columns)
-
-    def read_rows_meta_pseudo_df(
-        self, chunksize: int = 1000, run_config: Optional[RunConfig] = None
-    ) -> Iterator[DataDF]:
-        pit_timeout = "5m"
-
-        pit_resp = self.es_client.open_point_in_time(index=self.index, keep_alive=pit_timeout)
-        pit_id = pit_resp["id"]
-
-        query: dict
-        if run_config:
-            # run_config is not taken into account now
-            query = {"match_all": {}}
-        else:
-            query = {"match_all": {}}
-
-        data_resp: ObjectApiResponse[Any] | None
-        data_resp = self.es_client.search(
-            query=query,
-            sort=["_doc"],
-            pit={"id": pit_id, "keep_alive": pit_timeout},
-            size=chunksize,
-        )
-        if data_resp and len(data_resp["hits"]["hits"]) == 0:
-            data_resp = None
-            yield pd.DataFrame(columns=self.primary_key_columns)
-
-        while data_resp:
-            data = data_resp["hits"]["hits"]
-            last_search_result = data[-1]["sort"]
-
-            remapping_with_primary_keys = {
-                **self.key_name_remapping,
-                **{
-                    self.primary_key_column_rename.format(pk=primary_key): f"{primary_key}"
-                    for primary_key in self.primary_key_columns
-                },
-            }
-            result = [remap_dict_keys(item["_source"], remapping_with_primary_keys) for item in data]
-            yield pd.DataFrame(result)
-
-            data_resp = self.es_client.search(
-                query=query,
-                search_after=last_search_result,
-                sort=["_doc"],
-                pit={"id": pit_id, "keep_alive": pit_timeout},
-                size=chunksize,
-            )
-            if len(data_resp["hits"]["hits"]) == 0:
-                data_resp = None
-
-        self.es_client.close_point_in_time(id=pit_id)
-
-    def delete_rows(self, idx: IndexDF) -> None:
-        if idx.empty:
-            return
-        key_rows = _to_itertuples(idx, self.primary_key_columns)
-        rows_ids = [get_elastic_id(row) for row in key_rows]
-        actions = [{"_op_type": "delete", "_index": self.index, "_id": row_id} for row_id in rows_ids]
-        helpers.bulk(client=self.es_client, actions=actions, refresh=True)
-
-    def get_schema(self) -> DataSchema:
-        return self.data_sql_schema
-
-    def get_primary_schema(self) -> DataSchema:
-        return [column for column in self.data_sql_schema if column.primary_key]
-
-    def get_meta_schema(self) -> MetaSchema:
-        meta_key_prop = MetaKey.get_property_name()
-        return [column for column in self.data_sql_schema if hasattr(column, meta_key_prop)]

datapipe_core-0.14.2.dev2/datapipe/store/tests/abstract.py

@@ -1,307 +0,0 @@
-# This is copy of concept of reusable test classes from `fsspec`
-# https://github.com/fsspec/filesystem_spec/tree/master/fsspec/tests/abstract
-
-from typing import Callable, Iterable, cast
-
-import cloudpickle
-import pandas as pd
-import pytest
-from sqlalchemy import Column, String
-
-from datapipe.run_config import RunConfig
-from datapipe.store.table_store import TableStore
-from datapipe.store.tests.stubs import DATA_PARAMS
-from datapipe.tests.util import assert_df_equal, assert_ts_contains
-from datapipe.types import DataDF, DataSchema, IndexDF, data_to_index
-
-TableStoreMaker = Callable[[DataSchema], TableStore]
-
-
-class AbstractBaseStoreFixtures:
-    @pytest.fixture
-    def store_maker(self) -> TableStoreMaker:
-        raise NotImplementedError("This function must be overridden in derived classes")
-
-
-class AbstractBaseStoreTests:
-    def test_cloudpickle(self, store_maker: TableStoreMaker) -> None:
-        store = store_maker(
-            [
-                Column("id", String(), primary_key=True),
-            ]
-        )
-        ser = cloudpickle.dumps(store)
-        cloudpickle.loads(ser)
-
-        # TODO assert store is the same
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_get_schema(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_get_schema:
-            raise pytest.skip("Store does not support get_schema")
-
-        assert store.get_schema() == schema
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_write_read_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-        store.insert_rows(data_df)
-
-        assert_ts_contains(store, data_df)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_write_read_full_rows(
-        self, store_maker: TableStoreMaker, data_df: pd.DataFrame, schema: DataSchema
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_all_rows:
-            raise pytest.skip("Store does not support read_all_rows")
-
-        store.insert_rows(data_df)
-
-        assert_df_equal(store.read_rows(), data_df, index_cols=store.primary_keys)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_insert_identical_rows_twice_and_read_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        store.insert_rows(data_df)
-
-        test_df_mod = data_df.copy()
-        test_df_mod.loc[50:, "price"] = test_df_mod.loc[50:, "price"] + 1
-
-        store.insert_rows(test_df_mod.loc[50:])
-
-        assert_ts_contains(store, test_df_mod)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_read_non_existent_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_nonexistent_rows:
-            raise pytest.skip("Store does not support read_nonexistent_rows")
-
-        test_df_to_store = data_df.drop(range(1, 5))
-
-        store.insert_rows(test_df_to_store)
-
-        assert_df_equal(
-            store.read_rows(data_to_index(data_df, store.primary_keys)),
-            test_df_to_store,
-            index_cols=store.primary_keys,
-        )
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_read_empty_df(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-        store.insert_rows(data_df)
-
-        df_empty = pd.DataFrame()
-
-        df_result = store.read_rows(cast(IndexDF, df_empty))
-        assert df_result.empty
-        assert all(col in df_result.columns for col in store.primary_keys)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_insert_empty_df(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_all_rows:
-            raise pytest.skip("Store does not support read_all_rows")
-
-        df_empty = pd.DataFrame()
-        store.insert_rows(df_empty)
-
-        df_result = store.read_rows()
-        assert df_result.empty
-        assert all(col in df_result.columns for col in store.primary_keys)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_update_empty_df(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_all_rows:
-            raise pytest.skip("Store does not support read_all_rows")
-
-        df_empty = pd.DataFrame()
-        store.update_rows(df_empty)
-
-        df_result = store.read_rows()
-        assert df_result.empty
-        assert all(col in df_result.columns for col in store.primary_keys)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_partial_update_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-        store.insert_rows(data_df)
-
-        assert_ts_contains(store, data_df)
-
-        test_df_mod = data_df.copy()
-        test_df_mod.loc[50:, "price"] = test_df_mod.loc[50:, "price"] + 1
-
-        store.update_rows(test_df_mod.loc[50:])
-
-        assert_ts_contains(store, test_df_mod)
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_full_update_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-        store.insert_rows(data_df)
-
-        assert_ts_contains(store, data_df)
-
-        data_df_mod = data_df.copy()
-        data_df_mod.loc[:, "price"] = data_df_mod.loc[:, "price"] + 1
-
-        store.update_rows(data_df_mod)
-
-        assert_ts_contains(store, data_df_mod)
-
-    # TODO add test which does not require read_all_rows support
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_delete_rows(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_delete:
-            raise pytest.skip("Store does not support delete")
-        if not store.caps.supports_read_all_rows:
-            raise pytest.skip("Store does not support read_all_rows")
-
-        store.insert_rows(data_df)
-
-        assert_df_equal(
-            store.read_rows(data_to_index(data_df, store.primary_keys)),
-            data_df,
-            index_cols=store.primary_keys,
-        )
-
-        store.delete_rows(cast(IndexDF, data_df.loc[20:50, store.primary_keys]))
-
-        assert_df_equal(
-            store.read_rows(),
-            pd.concat([data_df.loc[0:19], data_df.loc[51:]]),
-            index_cols=store.primary_keys,
-        )
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_read_rows_meta_pseudo_df(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_meta_pseudo_df:
-            raise pytest.skip("Store does not support read_meta_pseudo_df")
-
-        store.insert_rows(data_df)
-
-        assert_ts_contains(store, data_df)
-
-        pseudo_df_iter = store.read_rows_meta_pseudo_df()
-
-        assert isinstance(pseudo_df_iter, Iterable)
-
-        pseudo_df = pd.concat(pseudo_df_iter, ignore_index=True)
-
-        for pk in store.primary_keys:
-            assert pk in pseudo_df.columns
-
-        # TODO check that ids of pseudo_df equal to ids of data_df
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_read_empty_rows_meta_pseudo_df(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_meta_pseudo_df:
-            raise pytest.skip("Store does not support read_meta_pseudo_df")
-
-        pseudo_df_iter = store.read_rows_meta_pseudo_df()
-        assert isinstance(pseudo_df_iter, Iterable)
-        for pseudo_df in pseudo_df_iter:
-            assert isinstance(pseudo_df, DataDF)
-            pseudo_df[store.primary_keys]  # Empty df must have primary keys columns
-
-    @pytest.mark.parametrize("data_df,schema", DATA_PARAMS)
-    def test_read_rows_meta_pseudo_df_with_runconfig(
-        self,
-        store_maker: TableStoreMaker,
-        data_df: pd.DataFrame,
-        schema: DataSchema,
-    ) -> None:
-        store = store_maker(schema)
-
-        if not store.caps.supports_read_meta_pseudo_df:
-            raise pytest.skip("Store does not support read_meta_pseudo_df")
-
-        store.insert_rows(data_df)
-
-        assert_ts_contains(store, data_df)
-
-        # TODO проверять, что runconfig реально влияет на результирующие данные
-        pseudo_df_iter = store.read_rows_meta_pseudo_df(
-            run_config=RunConfig(filters={"a": 1})
-        )
-        assert isinstance(pseudo_df_iter, Iterable)
-        for pseudo_df in pseudo_df_iter:
-            assert isinstance(pseudo_df, DataDF)

datapipe_core-0.14.2.dev2/datapipe/store/tests/stubs.py

@@ -1,89 +0,0 @@
-import pandas as pd
-import pytest
-from sqlalchemy import Column, Integer, String
-
-DATA_PARAMS = [
-    pytest.param(
-        pd.DataFrame(
-            {
-                "id": range(100),
-                "name": [f"Product {i}" for i in range(100)],
-                "price": [1000 + i for i in range(100)],
-            }
-        ),
-        [
-            Column("id", Integer, primary_key=True),
-            Column("name", String(100)),
-            Column("price", Integer),
-        ],
-        id="int_id",
-    ),
-    pytest.param(
-        pd.DataFrame(
-            {
-                "id": [f"id_{i}" for i in range(100)],
-                "name": [f"Product {i}" for i in range(100)],
-                "price": [1000 + i for i in range(100)],
-            }
-        ),
-        [
-            Column("id", String(100), primary_key=True),
-            Column("name", String(100)),
-            Column("price", Integer),
-        ],
-        id="str_id",
-    ),
-    pytest.param(
-        pd.DataFrame(
-            {
-                "id_int": range(100),
-                "id_str": [f"id_{i}" for i in range(100)],
-                "name": [f"Product {i}" for i in range(100)],
-                "price": [1000 + i for i in range(100)],
-            }
-        ),
-        [
-            Column("id_int", Integer, primary_key=True),
-            Column("id_str", String(100), primary_key=True),
-            Column("name", String(100)),
-            Column("price", Integer),
-        ],
-        id="multi_id",
-    ),
-    pytest.param(
-        pd.DataFrame(
-            {
-                "id1": [f"id_{i}" for i in range(1000)],
-                "id2": [f"id_{i}" for i in range(1000)],
-                "name": [f"Product {i}" for i in range(1000)],
-                "price": [1000 + i for i in range(1000)],
-            }
-        ),
-        [
-            Column("id1", String(100), primary_key=True),
-            Column("id2", String(100), primary_key=True),
-            Column("name", String(100)),
-            Column("price", Integer),
-        ],
-        id="double_id_1000_records",
-    ),
-    pytest.param(
-        pd.DataFrame(
-            {
-                "id1": [f"id_{i}" for i in range(1000)],
-                "id2": [f"id_{i}" for i in range(1000)],
-                "id3": [f"id_{i}" for i in range(1000)],
-                "name": [f"Product {i}" for i in range(1000)],
-                "price": [1000 + i for i in range(1000)],
-            }
-        ),
-        [
-            Column("id1", String(100), primary_key=True),
-            Column("id2", String(100), primary_key=True),
-            Column("id3", String(100), primary_key=True),
-            Column("name", String(100)),
-            Column("price", Integer),
-        ],
-        id="triple_id_1000_records",
-    ),
-]

datapipe_core-0.14.2.dev2/datapipe/tests/util.py

@@ -1,59 +0,0 @@
-from typing import List, cast
-
-import pandas as pd
-
-from datapipe.datatable import DataTable
-from datapipe.store.table_store import TableStore
-from datapipe.types import DataDF, IndexDF, data_to_index
-
-
-def assert_idx_equal(a, b):
-    a = sorted(list(a))
-    b = sorted(list(b))
-
-    assert a == b
-
-
-def assert_df_equal(a: pd.DataFrame, b: pd.DataFrame, index_cols=["id"]) -> bool:
-    a = a.set_index(index_cols)
-    b = b.set_index(index_cols)
-
-    assert_idx_equal(a.index, b.index)
-
-    eq_rows = (a.sort_index() == b.sort_index()).all(axis="columns")
-
-    if eq_rows.all():
-        return True
-
-    else:
-        print("Difference")
-        print("A:")
-        print(a.loc[-eq_rows])
-        print("B:")
-        print(b.loc[-eq_rows])
-
-        raise AssertionError
-
-
-def assert_datatable_equal(a: DataTable, b: DataDF) -> bool:
-    return assert_df_equal(a.get_data(), b, index_cols=a.primary_keys)
-
-
-def assert_ts_contains(ts: TableStore, df: DataDF):
-    assert_df_equal(
-        ts.read_rows(data_to_index(df, ts.primary_keys)),
-        df,
-        index_cols=ts.primary_keys,
-    )
-
-
-def assert_idx_no_duplicates(idx: IndexDF, index_cols: List[str]) -> bool:
-    duplicates = cast(IndexDF, idx[idx[index_cols].duplicated()])
-    if len(duplicates) == 0:
-        return True
-    else:
-        idx = cast(IndexDF, idx.loc[idx.index].sort_values(index_cols))
-        print("Duplicated found:")
-        print(idx)
-
-        raise AssertionError