How to write pandas with variant columns in Snowflake?

Viewed 410

I've created a df_to_snowflake function that uses pd_writer to write the dataframe to a table in Snowflake which works great with tables that contain standard values (string, bool, int, float). However, I have a table that contains a few columns that are either lists or dictionaries. I need to have them converted as variants but when I try to use this function I get this error:

ArrowNotImplementedError: Cannot write struct type 'CUSTOM_VARIABLES' with no child field to Parquet. Consider adding a dummy child field.

I tried converting the columns to strings and it works but then they are strings in snowflake which again is not what I need. How to change this function to automatically detect dict and list types and convert properly?

def df_to_snowflake(df, db, sc, tb, if_exists='append'):
    engine=create_engine(URL(
        account=acc,
        user=usr,
        warehouse=whs,
        role=rol,
        database=db,
        schema=sc,
        authenticator='externalbrowser'
    ))
    
    df.columns = df.columns.str.upper()
    df.to_sql(name=tb, con=engine, index=False, method=pd_writer, if_exists=if_exists)

UPDATE: A bit overkill but I got it working with this function instead. Any improvements would greatly be appreciated.

conn = sf.connect(
    account=acc,
    user=usr,
    warehouse=whs,
    role=rol,
    authenticator='externalbrowser'
)

def df_to_snowflake(df, db, sc, tb, if_exists='append'):
    engine=create_engine(URL(
        account=acc,
        user=usr,
        warehouse=whs,
        role=rol,
        database=db,
        schema=sc,
        authenticator='externalbrowser'
    ))
    
    df.columns = df.columns.str.upper()
    
    variant_type = {}
    has_variants = False
    row = df.iloc[0]
    for col in df.columns:
        if isinstance(row[col], dict) or isinstance(row[col], list):
            df[col] = pd.Series([json.dumps(x) for x in df[col]])
            has_variants = True
            variant_type[col] = True
        else:
            variant_type[col] = False
    
    df.to_sql(name=tb, con=engine, index=False, method=pd_writer, if_exists=if_exists)
    
    if has_variants:
        sql = []
        for k,v in variant_type.items():
            if v:
                sql.append(f"parse_json({k}) {k}")
            else:
                sql.append(f"{k}")
        sql = ",\n".join(sql)
        sql = f"create or replace table {db}.{sc}.{tb} as select {sql} from {db}.{sc}.{tb}"
        cursor = conn.cursor()
        cursor.execute(sql)
        cursor.close()
0 Answers
Related