web-dev-qa-db-fra.com

Comment créer une trame de données à partir d'AWS Athena à l'aide de la méthode Boto3 get_query_results

J'utilise AWS Athena pour interroger les données brutes de S3. Depuis qu'Athena écrit la sortie de la requête dans le bucket de sortie S3, je faisais:

df = pd.read_csv(OutputLocation)

Mais cela semble être un moyen coûteux. Récemment, j'ai remarqué le get_query_results méthode de boto3 qui renvoie un dictionnaire complexe des résultats.

client = boto3.client('athena')
response = client.get_query_results(
        QueryExecutionId=res['QueryExecutionId']
        )

Je suis confronté à deux problèmes principaux:

  1. Comment puis-je formater les résultats de get_query_results dans pandas trame de données?
  2. get_query_results ne renvoie que 1000 lignes. Comment puis-je l'utiliser pour obtenir deux millions de lignes?
pythonpandasamazon-web-servicesdataframeamazon-athena
7
Niv Cohen

get_query_results ne renvoie que 1000 lignes. Comment puis-je l'utiliser pour obtenir deux millions de lignes dans un Pandas dataframe?

Si vous essayez d'ajouter:

client.get_query_results(QueryExecutionId=res['QueryExecutionId'], MaxResults=2000)

Vous obtiendrez l'erreur suivante:

Une erreur s'est produite (InvalidRequestException) lors de l'appel de l'opération GetQueryResults: MaxResults est supérieur à la longueur maximale autorisée 1000.

Vous pouvez obtenir des millions de lignes si vous obtenez le fichier directement à partir de votre compartiment s3 (dans l'exemple suivant dans un Pandas Dataframe):

def obtain_data_from_s3(self):
    self.resource = boto3.resource('s3', 
                          region_name = self.region_name, 
                          aws_access_key_id = self.aws_access_key_id,
                          aws_secret_access_key= self.aws_secret_access_key)

    response = self.resource \
    .Bucket(self.bucket) \
    .Object(key= self.folder + self.filename + '.csv') \
    .get()

    return pd.read_csv(io.BytesIO(response['Body'].read()), encoding='utf8')

Le self.filename peut être:

self.filename = response['QueryExecutionId'] + ".csv"

Parce qu'Athena nomme les fichiers comme QueryExecutionId. Je vais vous écrire tout mon code qui prend une requête et retourner un cadre de données avec toutes les lignes et colonnes.

import time
import boto3
import pandas as pd
import io

class QueryAthena:

    def __init__(self, query, database):
        self.database = database
        self.folder = 'my_folder/'
        self.bucket = 'my_bucket'
        self.s3_input = 's3://' + self.bucket + '/my_folder_input'
        self.s3_output =  's3://' + self.bucket + '/' + self.folder
        self.region_name = 'us-east-1'
        self.aws_access_key_id = "my_aws_access_key_id"
        self.aws_secret_access_key = "my_aws_secret_access_key"
        self.query = query

    def load_conf(self, q):
        try:
            self.client = boto3.client('athena', 
                              region_name = self.region_name, 
                              aws_access_key_id = self.aws_access_key_id,
                              aws_secret_access_key= self.aws_secret_access_key)
            response = self.client.start_query_execution(
                QueryString = q,
                    QueryExecutionContext={
                    'Database': self.database
                    },
                    ResultConfiguration={
                    'OutputLocation': self.s3_output,
                    }
            )
            self.filename = response['QueryExecutionId']
            print('Execution ID: ' + response['QueryExecutionId'])

        except Exception as e:
            print(e)
        return response                

    def run_query(self):
        queries = [self.query]
        for q in queries:
            res = self.load_conf(q)
        try:              
            query_status = None
            while query_status == 'QUEUED' or query_status == 'RUNNING' or query_status is None:
                query_status = self.client.get_query_execution(QueryExecutionId=res["QueryExecutionId"])['QueryExecution']['Status']['State']
                print(query_status)
                if query_status == 'FAILED' or query_status == 'CANCELLED':
                    raise Exception('Athena query with the string "{}" failed or was cancelled'.format(self.query))
                time.sleep(10)
            print('Query "{}" finished.'.format(self.query))

            df = self.obtain_data()
            return df

        except Exception as e:
            print(e)      

    def obtain_data(self):
        try:
            self.resource = boto3.resource('s3', 
                                  region_name = self.region_name, 
                                  aws_access_key_id = self.aws_access_key_id,
                                  aws_secret_access_key= self.aws_secret_access_key)

            response = self.resource \
            .Bucket(self.bucket) \
            .Object(key= self.folder + self.filename + '.csv') \
            .get()

            return pd.read_csv(io.BytesIO(response['Body'].read()), encoding='utf8')   
        except Exception as e:
            print(e)  


if __name__ == "__main__":       
    query = "SELECT * FROM bucket.folder"
    qa = QueryAthena(query=query, database='myAthenaDb')
    dataframe = qa.run_query()
9
Eric Bellet

J'ai une solution pour ma première question, en utilisant la fonction suivante

def results_to_df(results):

    columns = [
        col['Label']
        for col in results['ResultSet']['ResultSetMetadata']['ColumnInfo']
    ]

    listed_results = []
    for res in results['ResultSet']['Rows'][1:]:
         values = []
         for field in res['Data']:
            try:
                values.append(list(field.values())[0]) 
            except:
                values.append(list(' '))

        listed_results.append(
            dict(Zip(columns, values))
        )

    return listed_results

puis:

t = results_to_df(response)
pd.DataFrame(t)

Quant à ma 2e question et à la demande de @EricBellet, j'ajoute également mon approche de la pagination que je trouve inefficace et plus longue par rapport au chargement des résultats de la sortie Athena en S3:

def run_query(query, database, s3_output):
    ''' 
    Function for executing Athena queries and return the query ID 
    '''
    client = boto3.client('athena')
    response = client.start_query_execution(
        QueryString=query,
        QueryExecutionContext={
            'Database': database
            },
        ResultConfiguration={
            'OutputLocation': s3_output,
            }
        )
    print('Execution ID: ' + response['QueryExecutionId'])
    return response



def format_result(results):
    '''
    This function format the results toward append in the needed format.
    '''
    columns = [
        col['Label']
        for col in results['ResultSet']['ResultSetMetadata']['ColumnInfo']
    ]

    formatted_results = []

    for result in results['ResultSet']['Rows'][0:]:
        values = []
        for field in result['Data']:
            try:
                values.append(list(field.values())[0]) 
            except:
                values.append(list(' '))

        formatted_results.append(
            dict(Zip(columns, values))
        )
    return formatted_results



res = run_query(query_2, database, s3_ouput) #query Athena



import sys
import boto3

marker = None
formatted_results = []
query_id = res['QueryExecutionId']
i = 0
start_time = time.time()

while True:
    paginator = client.get_paginator('get_query_results')
    response_iterator = paginator.paginate( 
        QueryExecutionId=query_id,
        PaginationConfig={
            'MaxItems': 1000,
            'PageSize': 1000,
            'StartingToken': marker})

    for page in response_iterator:
        i = i + 1
        format_page = format_result(page)
        if i == 1:
            formatted_results = pd.DataFrame(format_page)
        Elif i > 1:
            formatted_results = formatted_results.append(pd.DataFrame(format_page))

    try:
        marker = page['NextToken']
    except KeyError:
        break

print ("My program took", time.time() - start_time, "to run")

Ce n'est pas si bien formaté mais je pense que ça fait l'affaire ...

5
Niv Cohen