Load csv file, infer types and save the results in Spark SQL parquet file

spark_read_csv.py

from pyspark import SparkContext, SparkConf
from pyspark.sql import HiveContext, SQLContext
import pandas as pd

# sc: Spark context
# file_name: csv file_name
# table_name: output table name
# sep: csv file separator
# infer_limit: pandas type inference nb rows
def read_csv(sc, file_name, table_name, sep=",", infer_limit=10000):
    hc = HiveContext(sc)
    df = pd.read_csv(file_name, sep=sep, nrows=infer_limit)
    names = df.columns
    types = []
    for i in range(len(names)):
        tp = names[i] + " "
        if df.dtypes[i] == "O":
            tp += "STRING"
        elif df.dtypes[i] == "int64":
            tp += "INT"
        else:
            tp += "DOUBLE"
        types.append(tp)
    hc.sql('drop table if exists %s' %table_name)
    hc.sql("""CREATE TABLE IF NOT EXISTS %s (%s) row format delimited fields terminated by '%s'
LINES TERMINATED BY '\n' tblproperties ('skip.header.line.count'='1')""" %(table_name, ','.join(types), sep))
    hc.sql("LOAD DATA LOCAL INPATH '%s' OVERWRITE INTO TABLE %s" %(file_name, table_name))
    rdd = hc.sql("SELECT * FROM %s" %table_name)
    rdd.saveAsParquetFile("%s" %table_name)
    
# to read parquet file
# rdd = SQLContext(sc).parquetFile("parquet_dir")
# rdd.registerTempTable("mytable")
# rdd.sql("select count(*) from mytable")