Function upload_df

sdk/python/feast/infra/utils/snowflake/snowflake_utils.py:388–433

Args: df: Dataframe we'd like to write back. cursor: cursor to be used to communicate with Snowflake. stage_name: stage name in Snowflake connection. chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once

(
    df: pd.DataFrame,
    cursor: SnowflakeCursor,
    stage_name: str,
    chunk_size: Optional[int] = None,
    parallel: int = 4,
    compression: str = "gzip",
)

def upload_df(
    df: pd.DataFrame,
    cursor: SnowflakeCursor,
    stage_name: str,
    chunk_size: Optional[int] = None,
    parallel: int = 4,
    compression: str = "gzip",
):
    """
    Args:
        df: Dataframe we'd like to write back.
        cursor: cursor to be used to communicate with Snowflake.
        stage_name: stage name in Snowflake connection.
        chunk_size: Number of elements to be inserted once, if not provided all elements will be dumped once
            (Default value = None).
        parallel: Number of threads to be used when uploading chunks, default follows documentation at:
            https://docs.snowflake.com/en/sql-reference/sql/put.html#optional-parameters (Default value = 4).
        compression: The compression used on the Parquet files, can only be gzip, or snappy. Gzip gives supposedly a
            better compression, while snappy is faster. Use whichever is more appropriate (Default value = 'gzip').

    """
    if chunk_size is None:
        chunk_size = len(df)

    with TemporaryDirectory() as tmp_folder:
        for i, chunk in chunk_helper(df, chunk_size):
            chunk_path = os.path.join(tmp_folder, "file{}.txt".format(i))
            # Dump chunk into parquet file
            chunk.to_parquet(
                chunk_path,
                compression=compression,
                use_deprecated_int96_timestamps=True,
            )
            # Upload parquet file
            upload_sql = (
                "PUT /* Python:feast.infra.utils.snowflake_utils.upload_df() */ "
                "'file://{path}' @\"{stage_name}\" PARALLEL={parallel}"
            ).format(
                path=chunk_path.replace("\\", "\\\\").replace("'", "\\'"),
                stage_name=stage_name,
                parallel=parallel,
            )
            logger.debug(f"uploading files with '{upload_sql}'")
            cursor.execute(upload_sql, _is_internal=True)
            # Remove chunk file
            os.remove(chunk_path)


def upload_local_pq(