[DOP-22350] Add transformations for Transfers with dataframe column filtering (#186)

Author: IlyasDevelopment

docs/changelog/next_release/186.feature.rst

@@ -0,0 +1 @@
+Add transformations for **Transfers** with dataframe column filtering

syncmaster/schemas/v1/transfers/__init__.py

@@ -29,6 +29,9 @@
     S3ReadTransferTarget,
 )
 from syncmaster.schemas.v1.transfers.strategy import FullStrategy, IncrementalStrategy
+from syncmaster.schemas.v1.transfers.transformations.dataframe_columns_filter import (
+    DataframeColumnsFilter,
+)
 from syncmaster.schemas.v1.transfers.transformations.dataframe_rows_filter import (
     DataframeRowsFilter,
 )
@@ -102,7 +105,7 @@
     | None
 )
 
-TransformationSchema = DataframeRowsFilter
+TransformationSchema = DataframeRowsFilter | DataframeColumnsFilter
 
 
 class CopyTransferSchema(BaseModel):

syncmaster/schemas/v1/transfers/transformations/dataframe_columns_filter.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: 2023-2024 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from syncmaster.schemas.v1.transformation_types import DATAFRAME_COLUMNS_FILTER
+
+
+class BaseColumnsFilter(BaseModel):
+    field: str
+
+
+class IncludeFilter(BaseColumnsFilter):
+    type: Literal["include"]
+
+
+class RenameFilter(BaseColumnsFilter):
+    type: Literal["rename"]
+    to: str
+
+
+class CastFilter(BaseColumnsFilter):
+    type: Literal["cast"]
+    as_type: str
+
+
+ColumnsFilter = IncludeFilter | RenameFilter | CastFilter
+
+
+class DataframeColumnsFilter(BaseModel):
+    type: DATAFRAME_COLUMNS_FILTER
+    filters: list[Annotated[ColumnsFilter, Field(..., discriminator="type")]] = Field(default_factory=list)

syncmaster/schemas/v1/transfers/transformations/dataframe_rows_filter.py

@@ -7,74 +7,74 @@
 from syncmaster.schemas.v1.transformation_types import DATAFRAME_ROWS_FILTER
 
 
-class BaseRowFilter(BaseModel):
+class BaseRowsFilter(BaseModel):
     field: str
 
 
-class IsNullFilter(BaseRowFilter):
+class IsNullFilter(BaseRowsFilter):
     type: Literal["is_null"]
 
 
-class IsNotNullFilter(BaseRowFilter):
+class IsNotNullFilter(BaseRowsFilter):
     type: Literal["is_not_null"]
 
 
-class EqualFilter(BaseRowFilter):
+class EqualFilter(BaseRowsFilter):
     type: Literal["equal"]
     value: str
 
 
-class NotEqualFilter(BaseRowFilter):
+class NotEqualFilter(BaseRowsFilter):
     type: Literal["not_equal"]
     value: str
 
 
-class GreaterThanFilter(BaseRowFilter):
+class GreaterThanFilter(BaseRowsFilter):
     type: Literal["greater_than"]
     value: str
 
 
-class GreaterOrEqualFilter(BaseRowFilter):
+class GreaterOrEqualFilter(BaseRowsFilter):
     type: Literal["greater_or_equal"]
     value: str
 
 
-class LessThanFilter(BaseRowFilter):
+class LessThanFilter(BaseRowsFilter):
     type: Literal["less_than"]
     value: str
 
 
-class LessOrEqualFilter(BaseRowFilter):
+class LessOrEqualFilter(BaseRowsFilter):
     type: Literal["less_or_equal"]
     value: str
 
 
-class LikeFilter(BaseRowFilter):
+class LikeFilter(BaseRowsFilter):
     type: Literal["like"]
     value: str
 
 
-class ILikeFilter(BaseRowFilter):
+class ILikeFilter(BaseRowsFilter):
     type: Literal["ilike"]
     value: str
 
 
-class NotLikeFilter(BaseRowFilter):
+class NotLikeFilter(BaseRowsFilter):
     type: Literal["not_like"]
     value: str
 
 
-class NotILikeFilter(BaseRowFilter):
+class NotILikeFilter(BaseRowsFilter):
     type: Literal["not_ilike"]
     value: str
 
 
-class RegexpFilter(BaseRowFilter):
+class RegexpFilter(BaseRowsFilter):
     type: Literal["regexp"]
     value: str
 
 
-RowFilter = (
+RowsFilter = (
     IsNullFilter
     | IsNotNullFilter
     | EqualFilter
@@ -93,4 +93,4 @@ class RegexpFilter(BaseRowFilter):
 
 class DataframeRowsFilter(BaseModel):
     type: DATAFRAME_ROWS_FILTER
-    filters: list[Annotated[RowFilter, Field(..., discriminator="type")]] = Field(default_factory=list)
+    filters: list[Annotated[RowsFilter, Field(..., discriminator="type")]] = Field(default_factory=list)

syncmaster/schemas/v1/transformation_types.py

@@ -3,3 +3,4 @@
 from typing import Literal
 
 DATAFRAME_ROWS_FILTER = Literal["dataframe_rows_filter"]
+DATAFRAME_COLUMNS_FILTER = Literal["dataframe_columns_filter"]

syncmaster/worker/handlers/db/base.py

@@ -38,7 +38,8 @@ def read(self) -> DataFrame:
         reader = DBReader(
             connection=self.connection,
             table=self.transfer_dto.table_name,
-            where=self._get_filter_expression(),
+            where=self._get_rows_filter_expression(),
+            columns=self._get_columns_filter_expressions(),
         )
         return reader.run()
 
@@ -53,13 +54,47 @@ def write(self, df: DataFrame) -> None:
     def _normalize_column_names(self, df: DataFrame) -> DataFrame: ...
 
     @abstractmethod
-    def _make_filter_expression(self, filters: list[dict]) -> str | None: ...
+    def _make_rows_filter_expression(self, filters: list[dict]) -> str | None: ...
 
-    def _get_filter_expression(self) -> str | None:
-        filters = []
+    def _make_columns_filter_expressions(self, filters: list[dict]) -> list[str] | None:
+        expressions = []
+        for filter in filters:
+            filter_type = filter["type"]
+            field = self._quote_field(filter["field"])
+
+            if filter_type == "include":
+                expressions.append(field)
+            elif filter_type == "rename":
+                new_name = self._quote_field(filter["to"])
+                expressions.append(f"{field} AS {new_name}")
+            elif filter_type == "cast":
+                cast_type = filter["as_type"]
+                expressions.append(f"CAST({field} AS {cast_type}) AS {field}")
+
+        return expressions or None
+
+    def _get_rows_filter_expression(self) -> str | None:
+        expressions = []
         for transformation in self.transfer_dto.transformations:
             if transformation["type"] == "dataframe_rows_filter":
-                filters.extend(transformation["filters"])
-        if filters:
-            return self._make_filter_expression(filters)
+                expressions.extend(transformation["filters"])
+
+        if expressions:
+            return self._make_rows_filter_expression(expressions)
+
         return None
+
+    def _get_columns_filter_expressions(self) -> list[str] | None:
+        expressions = []
+        for transformation in self.transfer_dto.transformations:
+            if transformation["type"] == "dataframe_columns_filter":
+                expressions.extend(transformation["filters"])
+
+        if expressions:
+            return self._make_columns_filter_expressions(expressions)
+
+        return None
+
+    @staticmethod
+    def _quote_field(field: str) -> str:
+        return f'"{field}"'

syncmaster/worker/handlers/db/clickhouse.py

@@ -64,10 +64,10 @@ def _normalize_column_names(self, df: DataFrame) -> DataFrame:
             df = df.withColumnRenamed(column_name, column_name.lower())
         return df
 
-    def _make_filter_expression(self, filters: list[dict]) -> str | None:
+    def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []
         for filter in filters:
-            field = f'"{filter["field"]}"'
+            field = self._quote_field(filter["field"])
             op = self._operators[filter["type"]]
             value = filter.get("value")
 

syncmaster/worker/handlers/db/hive.py

@@ -40,11 +40,11 @@ def _normalize_column_names(self, df: DataFrame) -> DataFrame:
             df = df.withColumnRenamed(column_name, column_name.lower())
         return df
 
-    def _make_filter_expression(self, filters: list[dict]) -> str | None:
+    def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []
         for filter in filters:
             op = self._operators[filter["type"]]
-            field = f"`{filter["field"]}`"
+            field = self._quote_field(filter["field"])
             value = filter.get("value")
 
             if value is None:
@@ -59,3 +59,7 @@ def _make_filter_expression(self, filters: list[dict]) -> str | None:
                 expressions.append(f"{field} {op} '{value}'")
 
         return " AND ".join(expressions) or None
+
+    @staticmethod
+    def _quote_field(field):
+        return f"`{field}`"

syncmaster/worker/handlers/db/mssql.py

@@ -43,11 +43,11 @@ def _normalize_column_names(self, df: DataFrame) -> DataFrame:
             df = df.withColumnRenamed(column_name, column_name.lower())
         return df
 
-    def _make_filter_expression(self, filters: list[dict]) -> str | None:
+    def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []
         for filter in filters:
             op = self._operators[filter["type"]]
-            field = f'"{filter["field"]}"'
+            field = self._quote_field(filter["field"])
             value = filter.get("value")
 
             if value is None:

syncmaster/worker/handlers/db/mysql.py

@@ -40,11 +40,11 @@ def _normalize_column_names(self, df: DataFrame) -> DataFrame:
             df = df.withColumnRenamed(column_name, column_name.lower())
         return df
 
-    def _make_filter_expression(self, filters: list[dict]) -> str | None:
+    def _make_rows_filter_expression(self, filters: list[dict]) -> str | None:
         expressions = []
         for filter in filters:
             op = self._operators[filter["type"]]
-            field = f"`{filter["field"]}`"
+            field = self._quote_field(filter["field"])
             value = filter.get("value")
 
             if value is None:
@@ -59,3 +59,7 @@ def _make_filter_expression(self, filters: list[dict]) -> str | None:
                 expressions.append(f"{field} {op} '{value}'")
 
         return " AND ".join(expressions) or None
+
+    @staticmethod
+    def _quote_field(field: str) -> str:
+        return f"`{field}`"