Separate and cast bso state columns

data.py

@@ -6,6 +6,7 @@ from huggingface_hub import snapshot_download
 import numpy as np
 
 from string import ascii_letters
+from itertools import product
 
 from convert import (
   aux_global_id_to_code, presult, 
@@ -47,10 +48,10 @@ for season in tqdm(SEASONS):
   _aux_sched_df = pl.read_parquet(os.path.join(DATA_PATH, str(season), 'aux_schedule.parquet'))
   aux_sched_df = pl.concat((aux_sched_df, _aux_sched_df))
 
-backup_data_df = data_df
-backup_aux_df = aux_df
-backup_sched_df = sched_df
-backup_aux_sched_df = aux_sched_df
+# backup_data_df = data_df
+# backup_aux_df = aux_df
+# backup_sched_df = sched_df
+# backup_aux_sched_df = aux_sched_df
 
 aux_df = (
     aux_df
@@ -155,6 +156,9 @@ data_df = (
         (
             aux_df.filter(~pl.col('ibb'))[['universal_code', 'battingResult', 'inning_pas', 'pa_pitches', 'beforeBso', 'bso']]
             .rename({'battingResult': 'aux_bresult', 'inning_pas': 'aux_inning_pas', 'pa_pitches': 'aux_pa_pitches'})
+            .with_columns(
+                *[pl.col(f'{bso}').struct.field(field).cast(pl.Int32).alias(f'{"after" if bso == "bso" else "before"}_{field}') for bso, field in product(['beforeBso', 'bso'], ['b', 's', 'o'])]
+            )
         ),
         on='universal_code',
         how='left'