diff --git a/fieldExtraction/src/constants/investment_columns.py b/fieldExtraction/src/constants/investment_columns.py index f51968b..b30d8d8 100644 --- a/fieldExtraction/src/constants/investment_columns.py +++ b/fieldExtraction/src/constants/investment_columns.py @@ -84,8 +84,8 @@ COLUMN_ORDER = [ "FEE_SCHEDULE_VERSION", "AARETE_DERIVED_FEE_SCHEDULE_VERSION", "SERVICE_TERM", - "PROCEDURE_CD", - "PROCEDURE_CD_DESC", + "CPT4_PROC_CD", + "CPT4_PROC_CD_DESC", "CPT4_PROC_MOD", "CPT4_PROC_MOD_DESC", "REVENUE_CD", diff --git a/fieldExtraction/src/investment/investment_postprocessing_funcs.py b/fieldExtraction/src/investment/investment_postprocessing_funcs.py index e8db1e7..0ec42da 100644 --- a/fieldExtraction/src/investment/investment_postprocessing_funcs.py +++ b/fieldExtraction/src/investment/investment_postprocessing_funcs.py @@ -20,6 +20,26 @@ BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__fil MAPPINGS_DIR = os.path.join(BASE_DIR, 'crosswalk', 'mappings') +def rename_columns(df): + """ + Use the provided mapping to rename columns in the data + + Args: + df (pd.DataFrame): The input DataFrame with columns to be renamed. + Returns: + pd.DataFrame: The DataFrame with renamed columns. + """ + # Check if the DataFrame is empty, return it as is if it is + if df.empty: + return df + + # If not empty, proceed with column renaming + rename_map = {"PROCEDURE_CD" : "CPT4_PROC_CD", + "PROCEDURE_CD_DESC" : "CPT4_PROC_CD_DESC"} + # Rename columns using the mapping + df = df.rename(columns=rename_map) + return df + def normalize_indicator_field(value: str) -> str: """ diff --git a/fieldExtraction/src/investment/postprocess.py b/fieldExtraction/src/investment/postprocess.py index 274c616..1131c93 100644 --- a/fieldExtraction/src/investment/postprocess.py +++ b/fieldExtraction/src/investment/postprocess.py @@ -8,16 +8,18 @@ from src.prompts.investment_prompts import FieldSet, invoke_derived_term_date def postprocess(df): if df.shape[0] > 0: - # Add postprocessing here + df['CLIENT_NAME'] = config.CLIENT_NAME + # Rename columns + df = investment_postprocessing_funcs.rename_columns(df) + # Sort to ensure consistent ordering sort_columns = ['FILE_NAME'] if 'EXHIBIT_PAGE' in df.columns: sort_columns.append('EXHIBIT_PAGE') if 'REIMB_LESSER_OF_ID' in df.columns: sort_columns.append('REIMB_LESSER_OF_ID') - df = df.sort_values(by=sort_columns).reset_index(drop=True) # Postprocess date fields diff --git a/fieldExtraction/tests/investment_postprocess_test.py b/fieldExtraction/tests/investment_postprocess_test.py index 9e26cdb..7f8b255 100644 --- a/fieldExtraction/tests/investment_postprocess_test.py +++ b/fieldExtraction/tests/investment_postprocess_test.py @@ -15,6 +15,7 @@ from src.investment.investment_postprocessing_funcs import ( process_patient_age_range, remove_redundant_reimb_info, deduplicate_provider_columns, + rename_columns ) @@ -47,6 +48,59 @@ class TestPostprocessFunctions(unittest.TestCase): self.assertEqual(flatten_singleton_string_list("invalid"), "invalid") self.assertEqual(flatten_singleton_string_list(None), "") + def test_rename_columns(self): + """Tests the rename_columns function to ensure it correctly renames specified columns. + + Tests: + 1. Basic column renaming from PROCEDURE_CD to CPT4_PROC_CD + 2. Multiple columns being renamed + 3. Handling of columns not in the rename map + 4. Empty DataFrame + """ + # Test case 1: Basic column renaming + input_df1 = pd.DataFrame({ + "PROCEDURE_CD": ["12345", "67890"], + "PROCEDURE_CD_DESC": ["Test Procedure", "Another Procedure"], + "OTHER_COLUMN": ["value1", "value2"] + }) + + expected_df1 = pd.DataFrame({ + "CPT4_PROC_CD": ["12345", "67890"], + "CPT4_PROC_CD_DESC": ["Test Procedure", "Another Procedure"], + "OTHER_COLUMN": ["value1", "value2"] + }) + + result_df1 = rename_columns(input_df1) + pd.testing.assert_frame_equal(result_df1, expected_df1) + + # Test case 2: Only some columns need renaming + input_df2 = pd.DataFrame({ + "PROCEDURE_CD": ["12345", "67890"], + "OTHER_COLUMN": ["value1", "value2"] + }) + + expected_df2 = pd.DataFrame({ + "CPT4_PROC_CD": ["12345", "67890"], + "OTHER_COLUMN": ["value1", "value2"] + }) + + result_df2 = rename_columns(input_df2) + pd.testing.assert_frame_equal(result_df2, expected_df2) + + # Test case 3: None of the columns need renaming + input_df3 = pd.DataFrame({ + "COLUMN_A": ["a", "b"], + "COLUMN_B": ["c", "d"] + }) + + result_df3 = rename_columns(input_df3) + pd.testing.assert_frame_equal(result_df3, input_df3) # Should be unchanged + + # Test case 4: Empty DataFrame + empty_df = pd.DataFrame() + result_empty_df = rename_columns(empty_df) + pd.testing.assert_frame_equal(result_empty_df, empty_df) # Should be unchanged + @patch("src.utils.llm_utils.invoke_claude") @patch("src.investment.investment_postprocessing_funcs.invoke_derived_term_date")