def METHOD_NAME(
    self, context: InputContext, table_slice: TableSlice, connection
) -> pyspark.sql.DataFrame:
    """Loads the return of the query as the correct type."""
    spark = SparkSession.builder.getOrCreate()  # type: ignore
    if table_slice.partition_dimensions and len(context.asset_partition_keys) == 0:
        return spark.createDataFrame([], StructType([]))
    pd_df = connection.execute(DuckDbClient.get_select_statement(table_slice)).fetchdf()
    return spark.createDataFrame(pd_df)