카테고리 없음

Polars type 변경

유키공 2025. 5. 26. 13:49
import polars as pl

def process_dataframe_optimized_pl(dict_df_types: dict, df: pl.DataFrame) -> pl.DataFrame:
    def handle_column(col: str, dtype: str) -> pl.Expr:
        try:
            expr = pl.col(col)

            if dtype == 'int':
                return expr.cast(pl.Int32).fill_null(0).alias(col)
            elif dtype == 'float':
                return expr.cast(pl.Float32).fill_null(0).alias(col)
            elif dtype == 'bool':
                return expr.cast(pl.Utf8).str.to_lowercase().is_in(['true', 't', '1']).alias(col)
            elif dtype == 'datetime':
                return expr.cast(pl.Utf8).str.strptime(pl.Datetime, strict=False).alias(col)
            elif dtype == 'string':
                return expr.cast(pl.String).alias(col)
            elif dtype == 'category':
                return expr.cast(pl.Categorical).alias(col)
            else:
                return expr  # return original if dtype not recognized
        except Exception as e:
            print(f"컬럼 '{col}' 처리 중 오류 발생: {e}")
            return expr  # return original on error

    # Get intersection of DataFrame columns and dictionary keys
    valid_cols = set(df.columns) & set(dict_df_types.keys())
    
    # Filter for only valid types we want to process
    valid_types = {'int', 'float', 'bool', 'datetime', 'string', 'category'}
    
    # Create expressions for columns that need processing
    exprs = [
        handle_column(col, dict_df_types[col])
        for col in valid_cols
        if dict_df_types.get(col) in valid_types
    ]
    
    return df.with_columns(exprs)