7
\$\begingroup\$

My Goal: To efficiently move data from Amazon S3 to Amazon Redshift.

Basically, I am moving all CSV files on my S3 to Redshift using the below code. I parse through part of the file, build a table structure and then use the copy command to load data into redshift.

'''
Created on Feb 25, 2015
@author: Siddartha.Reddy
'''

import sys
from boto.s3 import connect_to_region
from boto.s3.connection import Location
import csv
import itertools
import psycopg2

''' ARGUMENTS TO PASS '''
AWS_KEY = sys.argv[1]
AWS_SECRET_KEY = sys.argv[2]
S3_DOWNLOAD_PATH = sys.argv[3]
REDSHIFT_SCHEMA = sys.argv[4]
TABLE_NAME = sys.argv[5]

UTILS = S3_DOWNLOAD_PATH.split('/')

class UTIL():

    global UTILS

    def bucket_name(self):
        self.BUCKET_NAME = UTILS[0]
        return self.BUCKET_NAME

    def path(self):
        self.PATH = ''
        offset = 0
        for value in UTILS:
            if offset == 0:
                offset += 1
            else:
                self.PATH = self.PATH + value + '/'
        return self.PATH[:-1]

def GETDATAINMEMORY():
    conn = connect_to_region(Location.USWest2,aws_access_key_id = AWS_KEY,
        aws_secret_access_key = AWS_SECRET_KEY,
        is_secure=False,host='s3-us-west-2.amazonaws.com'
        )
    ut = util()
    BUCKET_NAME = ut.bucket_name()
    PATH = ut.path()
    filelist = conn.lookup(BUCKET_NAME)

    ''' Fecth part of the data from S3 '''
    for path in filelist:
        if PATH in path.name:
            DATA = path.get_contents_as_string(headers={'Range': 'bytes=%s-%s' % (0,100000000)}) 

    return DATA

def TRAVERSEDATA():
    DATA = getdatainmemory()
    CREATE_TABLE_QUERY = 'CREATE TABLE ' + REDSHIFT_SCHEMA + '.' + TABLE_NAME + '( '
    JUNKED_OUT = DATA[3:]
    PROCESSED_DATA = JUNKED_OUT.split('\n')
    CSV_DATA = csv.reader(PROCESSED_DATA,delimiter=',')
    COUNTER,STRING,NUMBER = 0,0,0
    COLUMN_TYPE = []

    ''' GET COLUMN NAMES AND COUNT '''
    for line in CSV_DATA:
        NUMBER_OF_COLUMNS = len(line)
        COLUMN_NAMES = line
        break;

    ''' PROCESS COLUMN NAMES '''
    a = 0
    for REMOVESPACE in COLUMN_NAMES:
        TEMPHOLDER = REMOVESPACE.split(' ')
        temp1 = ''
        for x in TEMPHOLDER:
            temp1 = temp1 + x 
        COLUMN_NAMES[a] = temp1
        a = a + 1

    ''' GET COLUMN DATA TYPES '''
    # print(NUMBER_OF_COLUMNS,COLUMN_NAMES,COUNTER)
    # print(NUMBER_OF_COLUMNS)
    i,j,a= 0,500,0 
    while COUNTER < NUMBER_OF_COLUMNS:
        for COLUMN in itertools.islice(CSV_DATA,i,j+1):
            if COLUMN[COUNTER].isdigit():
                NUMBER = NUMBER + 1
            else:
                STRING = STRING + 1
        if NUMBER == 501:
            COLUMN_TYPE.append('INTEGER')
            # print('I CAME IN')
            NUMBER = 0
        else:
            COLUMN_TYPE.append('VARCHAR(2500)')
            STRING = 0
        COUNTER = COUNTER + 1
        # print(COUNTER)

    COUNTER = 0
    ''' BUILD SCHEMA '''
    while COUNTER < NUMBER_OF_COLUMNS:
        if COUNTER == 0:
            CREATE_TABLE_QUERY = CREATE_TABLE_QUERY + COLUMN_NAMES[COUNTER] + ' ' + COLUMN_TYPE[COUNTER] + ' NOT NULL,'
        else:
            CREATE_TABLE_QUERY = CREATE_TABLE_QUERY + COLUMN_NAMES[COUNTER] + ' ' + COLUMN_TYPE[COUNTER] + ' ,'
        COUNTER += 1
    CREATE_TABLE_QUERY = CREATE_TABLE_QUERY[:-2]+ ')'

    return CREATE_TABLE_QUERY

def COPY_COMMAND():
    S3_PATH = 's3://' + S3_DOWNLOAD_PATH
    COPY_COMMAND = "COPY "+REDSHIFT_SCHEMA+"."+TABLE_NAME+" from '"+S3_PATH+"' credentials 'aws_access_key_id="+AWS_KEY+";aws_secret_access_key="+AWS_SECRET_KEY+"' REGION 'us-west-2' csv delimiter ',' ignoreheader as 1 TRIMBLANKS maxerror as 500"
    return COPY_COMMAND

def S3TOREDSHIFT():
    conn = psycopg2.connect("dbname='xxx' port='5439' user='xxx' host='xxxxxx' password='xxxxx'")
    cursor = conn.cursor()
    cursor.execute('DROP TABLE IF EXISTS '+ REDSHIFT_SCHEMA + "." + TABLE_NAME)
    SCHEMA = TRAVERSEDATA()
    print(SCHEMA)
    cursor.execute(SCHEMA)
    COPY = COPY_COMMAND()
    print(COPY)
    cursor.execute(COPY)
    conn.commit()

S3TOREDSHIFT()

Current Challenges:

Challenges with creating the table structure:

  1. Field lengths: Right now I am just hardcoding the VARCHAR fields to 2500. All my files are > 30gb and parsing through the whole file to calculate length of a field takes lot of processing time.
  2. Determining if a column is null: I am simply hard coding the first column to NOT NULL using the COUNTER variable. (All my files have ID as first column). I would like to know if there is a better way of doing it.

Is there any data structure I can use? I am always interested in learning new ways to improve the performance.

\$\endgroup\$
  • 1
    \$\begingroup\$ I'd suggest browsing the major points of PEP-8 guidelines, otherwise every Python dev that sees your code will call you out on your code - not to mention it makes your code far easier to read. \$\endgroup\$ – jsanc623 Mar 5 '15 at 15:14
  • \$\begingroup\$ If you care more about speed than about using Python you could check out a specialised program for loading data into PostgreSQL, pgloader; you'd still have to script the downloading from S3 though. \$\endgroup\$ – ferada Mar 6 '15 at 12:41
2
\$\begingroup\$

Your capitalization convention strikes me as odd, not just for Python, but for any language I've ever used. Not only is it hard to read, but things cased LIKE_THIS are usually constants in Python, so this usage will confuse Python programmers. Many other things about your code formatting are also confusing or unconventional. Please do check out PEP-8.

Your function traverse_data is quite long. I'd recommend trying to split everything under ''' One of these ''' into its own function. Also note that the comment marker in Python is #. The syntax ''' like this ''' is for strings. Python allows you to put one of those as the first thing in a function or class, as a documentation string. You can then read documentation interactively on the command line. If the triple quoted string isn't the first thing in a function or class, it's just a string literal that never gets assigned to anything and disappears. As a personal request from me, please don't document your code like this. The performance hit is probably negligible, but it's just weird and wrong. Stick to a doc string at the start of a function or class, and regular comments for things that aren't going to be part of the code.

It's more common when you have a script you want to run from the command line to write it like this:

if __name__ == "__main__":
    # code currently inside s3_to_redshift
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.