From 6c5dfdeb1e755eacd0c3d70921b95b6aa9c17fec Mon Sep 17 00:00:00 2001 From: MoooCat Date: Fri, 21 Jun 2024 23:39:57 +0800 Subject: [PATCH] Improving the robustness of the 'remove_columns' method --- sdgx/data_processors/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/sdgx/data_processors/base.py b/sdgx/data_processors/base.py index f69773de..21f97ecc 100644 --- a/sdgx/data_processors/base.py +++ b/sdgx/data_processors/base.py @@ -1,9 +1,9 @@ from __future__ import annotations from typing import Any, Dict - import pandas as pd +from sdgx.log import logger from sdgx.data_models.metadata import Metadata from sdgx.exceptions import SynthesizerProcessorError @@ -72,7 +72,10 @@ def remove_columns(tabular_data: pd.DataFrame, column_name_to_remove: list) -> p result_data = tabular_data.copy() # Remove specified columns - result_data = result_data.drop(columns=column_name_to_remove) + try: + result_data = result_data.drop(columns=column_name_to_remove) + except KeyError: + logger.warning('Duplicate column removal occurred, which might lead to unintended consequences.') return result_data