Add a column mapping task

wildlife-dynamics · Jun 20, 2024 · 38df604 · 38df604
1 parent 57e47cf
commit 38df604
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/ecoscope_workflows/tasks/transformation/__init__.py b/ecoscope_workflows/tasks/transformation/__init__.py
diff --git a/ecoscope_workflows/tasks/transformation/column_mapping.py b/ecoscope_workflows/tasks/transformation/column_mapping.py
@@ -0,0 +1,23 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from ecoscope_workflows.annotations import DataFrame, JsonSerializableDataFrameModel
+from ecoscope_workflows.decorators import distributed
+
+
+@distributed
+def map_columns(
+    df: DataFrame[JsonSerializableDataFrameModel],
+    drop_columns: Annotated[list[str], Field()],
+    retain_columns: Annotated[list[str], Field()],
+    rename_columns: Annotated[dict[str, str], Field()],
+) -> DataFrame[JsonSerializableDataFrameModel]:
+
+    df.drop(columns=drop_columns, inplace=True)
+    if retain_columns:
+        df.reindex(columns=[retain_columns], inplace=True)
+    if rename_columns:
+        df.rename(columns=rename_columns, inplace=True)
+
+    return df