Skip to content

Instantly share code, notes, and snippets.

@lmatthieu
Last active October 5, 2016 15:09
Show Gist options
  • Save lmatthieu/20c5fe13f491af8a6a60 to your computer and use it in GitHub Desktop.
Save lmatthieu/20c5fe13f491af8a6a60 to your computer and use it in GitHub Desktop.
Load csv file, infer types and save the results in Spark SQL parquet file
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")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment