[WIP]: Datapipe ml opts: добавлена динамическая фильтрация индексов

.github/workflows/test_examples.yaml

@@ -25,6 +25,8 @@ jobs:
           - many_to_zero
           - model_inference
           - one_to_many_pipeline
+          - batch_transform_with_filters/simple-example
+          - batch_transform_with_filters/filters-as-function
         executor:
           - SingleThreadExecutor
           - RayExecutor

CHANGELOG.md

@@ -13,6 +13,8 @@ Major changes:
 
 See "Migration from v0.13 to v0.14" for more details
 
+* `BatchTransform` has new argument `filters`. It's using to filter processing transform indexes using only that indexes that as indicated in `filters`. See `docs/concepts.md` for more details.
+
 # 0.13.14
 
 * Fix [#334](https://github.com/epoch8/datapipe/issues/334)

Makefile

@@ -0,0 +1,6 @@
+black:
+	autoflake -r --in-place --remove-all-unused-imports steps/ *.py brandlink_utils/
+	black --verbose --config black.toml steps/ alembic *.py brandlink_utils/
+
+mypy:
+	mypy -p datapipe --ignore-missing-imports --follow-imports=silent --namespace-packages

black.toml

@@ -0,0 +1,2 @@
+[tool.black]
+line-length = 120

datapipe/cli.py

@@ -5,6 +5,7 @@
 from typing import Dict, List, Optional, cast
 
 import click
+from datapipe.run_config import RunConfig
 import pandas as pd
 import rich
 from opentelemetry import trace
@@ -318,6 +319,7 @@ def step(
     steps = filter_steps_by_labels_and_name(app, labels=labels_list, name_prefix=name)
 
     ctx.obj["steps"] = steps
+    ctx.obj["labels"] = labels_list
 
 
 def to_human_repr(step: ComputeStep, extra_args: Optional[Dict] = None) -> str:
@@ -379,6 +381,7 @@ def step_list(ctx: click.Context, status: bool) -> None:  # noqa
 def step_run(ctx: click.Context, loop: bool, loop_delay: int) -> None:
     app: DatapipeApp = ctx.obj["pipeline"]
     steps_to_run: List[ComputeStep] = ctx.obj["steps"]
+    run_config = RunConfig(labels={k: v for k, v in ctx.obj["labels"]})
 
     executor: Executor = ctx.obj["executor"]
 
@@ -387,7 +390,7 @@ def step_run(ctx: click.Context, loop: bool, loop_delay: int) -> None:
 
     while True:
         if len(steps_to_run) > 0:
-            run_steps(app.ds, steps_to_run, executor=executor)
+            run_steps(app.ds, steps_to_run, run_config=run_config, executor=executor)
 
         if not loop:
             break

datapipe/datatable.py

@@ -7,6 +7,11 @@
 from datapipe.event_logger import EventLogger
 from datapipe.meta.sql_meta import MetaTable
 from datapipe.run_config import RunConfig
+from datapipe.sql_util import (
+    sql_apply_idx_filter_to_table,
+    sql_apply_runconfig_filters,
+)
+from datapipe.store.database import DBConn, MetaKey
 from datapipe.store.database import DBConn
 from datapipe.store.table_store import TableStore
 from datapipe.types import DataDF, IndexDF, MetadataDF, data_to_index, index_difference

datapipe/meta/sql_meta.py

@@ -20,7 +20,7 @@
 import sqlalchemy as sa
 
 from datapipe.run_config import RunConfig
-from datapipe.sql_util import sql_apply_idx_filter_to_table, sql_apply_runconfig_filter
+from datapipe.sql_util import sql_apply_idx_filter_to_table, sql_apply_runconfig_filters
 from datapipe.store.database import DBConn, MetaKey
 from datapipe.types import (
     DataDF,
@@ -389,7 +389,7 @@ def get_stale_idx(
             )
         )
 
-        sql = sql_apply_runconfig_filter(
+        sql = sql_apply_runconfig_filters(
             sql, self.sql_table, self.primary_keys, run_config
         )
 
@@ -447,7 +447,7 @@ def get_agg_cte(
             sql = sql.group_by(*key_cols)
 
         sql = sql_apply_filters_idx_to_subquery(sql, keys, filters_idx)
-        sql = sql_apply_runconfig_filter(sql, tbl, self.primary_keys, run_config)
+        sql = sql_apply_runconfig_filters_to_subquery(sql, self.primary_keys, run_config)
 
         return (keys, sql.cte(name=f"{tbl.name}__update"))
 
@@ -649,7 +649,7 @@ def mark_all_rows_unprocessed(
             .where(self.sql_table.c.is_success == True)
         )
 
-        sql = sql_apply_runconfig_filter(
+        sql = sql_apply_runconfig_filters(
             update_sql, self.sql_table, self.primary_keys, run_config
         )
 
@@ -680,6 +680,18 @@ def sql_apply_filters_idx_to_subquery(
     return sql
 
 
+def sql_apply_runconfig_filters_to_subquery(
+    sql: Any,
+    keys: List[str],
+    run_config: Optional[RunConfig] = None,
+) -> Any:
+    if run_config is not None:
+        filters_idx = pd.DataFrame(run_config.filters)
+        sql = sql_apply_filters_idx_to_subquery(sql, keys, filters_idx)
+
+    return sql
+
+
 @dataclass
 class ComputeInputCTE:
     cte: Any
@@ -795,6 +807,7 @@ def build_changed_idx_sql(
     )
 
     out = sql_apply_filters_idx_to_subquery(out, transform_keys, filters_idx)
+    out = sql_apply_runconfig_filters_to_subquery(out, transform_keys, run_config)
 
     out = out.cte(name="transform")
 

datapipe/run_config.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from typing import Any, List, Dict, Optional
 
-LabelDict = Dict[str, Any]
+import pandas as pd
+from datapipe.types import LabelDict
 
 
 @dataclass
@@ -10,7 +11,7 @@ class RunConfig:
     # если не пуст, то во время запуска обрабатываются только те строки,
     # которые строго соответствуют фильтру
     # (в случае, если у таблицы есть идентификатор с совпадающим именем).
-    filters: LabelDict = field(default_factory=dict)
+    filters: List[LabelDict] = field(default_factory=list)
     labels: LabelDict = field(default_factory=dict)
 
     @classmethod
@@ -22,3 +23,16 @@ def add_labels(cls, rc: Optional["RunConfig"], labels: LabelDict) -> "RunConfig"
             )
         else:
             return RunConfig(labels=labels)
+
+    @classmethod
+    def add_filters(cls, rc: Optional["RunConfig"], filters: List[LabelDict]) -> "RunConfig":
+        if rc is not None:
+            return RunConfig(
+                filters=list(
+                    pd.concat([pd.DataFrame(rc.filters), pd.DataFrame(filters)], ignore_index=True)
+                    .drop_duplicates()
+                    .apply(lambda row : row.dropna().to_dict(), axis=1)
+                ),
+            )
+        else:
+            return RunConfig(filters=filters)

datapipe/sql_util.py

@@ -1,10 +1,12 @@
-from typing import Any, Dict, List, Optional
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, cast
 
-from sqlalchemy import Column, Integer, String, Table, tuple_
+import pandas as pd
+from sqlalchemy import Column, Integer, String, Table, column, tuple_
+from sqlalchemy.sql.expression import and_, or_
 
 from datapipe.run_config import RunConfig
-from datapipe.types import IndexDF
-
+from datapipe.types import IndexDF, LabelDict
 
 def sql_apply_idx_filter_to_table(
     sql: Any,
@@ -22,22 +24,28 @@ def sql_apply_idx_filter_to_table(
         keys = tuple_(*[table.c[key] for key in primary_keys])  # type: ignore
 
         sql = sql.where(
-            keys.in_([tuple([r[key] for key in primary_keys]) for r in idx.to_dict(orient="records")])  # type: ignore
+            keys.in_(
+                [
+                    tuple([r[key] for key in primary_keys])  # type: ignore
+                    for r in idx.to_dict(orient="records")
+                ]
+            )
         )
 
     return sql
 
 
-def sql_apply_runconfig_filter(
+def sql_apply_runconfig_filters(
     sql: Any,
     table: Table,
-    primary_keys: List[str],
+    keys: List[str],
     run_config: Optional[RunConfig] = None,
 ) -> Any:
     if run_config is not None:
-        for k, v in run_config.filters.items():
-            if k in primary_keys:
-                sql = sql.where(table.c[k] == v)
+        filters_idx = pd.DataFrame(run_config.filters)
+        primary_keys = [key for key in keys if key in table.c and key in filters_idx.columns]
+        if len(filters_idx) > 0 and len(primary_keys) > 0:
+            sql = sql_apply_idx_filter_to_table(sql, table, primary_keys, cast(IndexDF, filters_idx))
 
     return sql
 
@@ -49,4 +57,4 @@ def sql_apply_runconfig_filter(
 
 
 def sql_schema_to_dtype(schema: List[Column]) -> Dict[str, Any]:
-    return {i.name: SCHEMA_TO_DTYPE_LOOKUP[i.type.__class__] for i in schema}  # type: ignore
+    return {i.name: SCHEMA_TO_DTYPE_LOOKUP[i.type.__class__] for i in schema}