Skip to content

Instantly share code, notes, and snippets.

@aaronsteers
Forked from lmatthieu/spark_read_csv.py
Created October 5, 2016 15:09
Show Gist options
  • Save aaronsteers/d5cbb912875f108ea624258b69f26f94 to your computer and use it in GitHub Desktop.
Save aaronsteers/d5cbb912875f108ea624258b69f26f94 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