blob: b461d7afdd823c573ec6fa65755b6fcf79bbbed9 [file] [log] [blame]
#!/usr/bin/env impala-python
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# This script is used to load the proper datasets for the specified workloads. It loads
# all data via Hive except for parquet data which needs to be loaded via Impala.
# Most ddl commands are executed by Impala.
import collections
import getpass
import logging
import multiprocessing
import os
import re
import shutil
import sqlparse
import subprocess
import sys
import tempfile
import time
import traceback
from optparse import OptionParser
from tests.beeswax.impala_beeswax import *
from multiprocessing.pool import ThreadPool
LOG = logging.getLogger('')
parser = OptionParser()
parser.add_option("-e", "--exploration_strategy", dest="exploration_strategy",
help="The exploration strategy for schema gen: 'core', "\
"'pairwise', or 'exhaustive'")
parser.add_option("--hive_warehouse_dir", dest="hive_warehouse_dir",
help="The HDFS path to the base Hive test warehouse directory")
parser.add_option("-w", "--workloads", dest="workloads",
help="Comma-separated list of workloads to load data for. If 'all' is "\
"specified then data for all workloads is loaded.")
parser.add_option("-s", "--scale_factor", dest="scale_factor", default="",
help="An optional scale factor to generate the schema for")
parser.add_option("-f", "--force_reload", dest="force_reload", action="store_true",
default=False, help='Skips HDFS exists check and reloads all tables')
parser.add_option("--impalad", dest="impalad", default="localhost",
help="Impala daemon to connect to")
parser.add_option("--hive_hs2_hostport", dest="hive_hs2_hostport",
help="HS2 host:Port to issue Hive queries against using beeline")
parser.add_option("--table_names", dest="table_names", default=None,
help="Only load the specified tables - specified as a comma-seperated "\
"list of base table names")
parser.add_option("--table_formats", dest="table_formats", default=None,
help="Override the test vectors and load using the specified table "\
"formats. Ex. --table_formats=seq/snap/block,text/none")
parser.add_option("--hdfs_namenode", dest="hdfs_namenode", default="localhost:20500",
help="HDFS name node for Avro schema URLs, default localhost:20500")
parser.add_option("--workload_dir", dest="workload_dir",
help="Directory that contains Impala workloads")
parser.add_option("--dataset_dir", dest="dataset_dir",
help="Directory that contains Impala datasets")
parser.add_option("--use_kerberos", action="store_true", default=False,
help="Load data on a kerberized cluster.")
parser.add_option("--principal", default=None, dest="principal",
help="Kerberos service principal, required if --use_kerberos is set")
parser.add_option("--num_processes", type="int", default=multiprocessing.cpu_count(),
dest="num_processes", help="Number of parallel processes to use.")
options, args = parser.parse_args()
WORKLOAD_DIR = options.workload_dir
DATASET_DIR = options.dataset_dir
TESTDATA_BIN_DIR = os.path.join(os.environ['IMPALA_HOME'], 'testdata/bin')
AVRO_SCHEMA_DIR = "avro_schemas"
GENERATE_SCHEMA_CMD = " --exploration_strategy=%s "\
"--workload=%s --scale_factor=%s --verbose"
# Load data using Hive's beeline because the Hive shell has regressed (HIVE-5515).
# The Hive shell is stateful, meaning that certain series of actions lead to problems.
# Examples of problems due to the statefullness of the Hive shell:
# - Creating an HBase table changes the replication factor to 1 for subsequent LOADs.
# - INSERTs into an HBase table fail if they are the first stmt executed in a session.
# However, beeline itself also has bugs. For example, inserting a NULL literal into
# a string-typed column leads to an NPE. We work around these problems by using LOAD from
# a datafile instead of doing INSERTs.
HIVE_CMD = os.path.join(os.environ['HIVE_HOME'], 'bin/beeline')
hive_auth = "auth=none"
if options.use_kerberos:
if not options.principal:
print "--principal is required when --use_kerberos is specified"
hive_auth = "principal=" + options.principal
HIVE_ARGS = '-n %s -u "jdbc:hive2://%s/default;%s" --verbose=true'\
% (getpass.getuser(), options.hive_hs2_hostport, hive_auth)
HADOOP_CMD = os.path.join(os.environ['HADOOP_HOME'], 'bin/hadoop')
def available_workloads(workload_dir):
return [subdir for subdir in os.listdir(workload_dir)
if os.path.isdir(os.path.join(workload_dir, subdir))]
def validate_workloads(all_workloads, workloads):
for workload in workloads:
if workload not in all_workloads:
LOG.error('Workload \'%s\' not found in workload directory' % workload)
LOG.error('Available workloads: ' + ', '.join(all_workloads))
def exec_cmd(cmd, error_msg=None, exit_on_error=True, out_file=None):
"""Run the given command in the shell returning whether the command
succeeded. If 'error_msg' is set, log the error message on failure.
If 'exit_on_error' is True, exit the program on failure.
If 'out_file' is specified, log all output to that file."""
success = True
if out_file:
with open(out_file, 'w') as f:
ret_val =, shell=True, stderr=f, stdout=f)
ret_val =, shell=True)
if ret_val != 0:
if error_msg:
if exit_on_error: sys.exit(ret_val)
success = False
return success
def exec_hive_query_from_file_beeline(file_name):
if not os.path.exists(file_name):"Error: File {0} not found".format(file_name))
return False"Beginning execution of hive SQL: {0}".format(file_name))
# When HiveServer2 is configured to use "local" mode (i.e., MR jobs are run
# in-process rather than on YARN), Hadoop's LocalDistributedCacheManager has a
# race, wherein it tires to localize jars into
# /tmp/hadoop-$USER/mapred/local/<millis>. Two simultaneous Hive queries
# against HS2 can conflict here. Weirdly LocalJobRunner handles a similar issue
# (with the staging directory) by appending a random number. To over come this,
# in the case that HS2 is on the local machine (which we conflate with also
# running MR jobs locally), we move the temporary directory into a unique
# directory via configuration. This block can be removed when
# is resolved.
hive_args = HIVE_ARGS
unique_dir = None
if options.hive_hs2_hostport.startswith("localhost:"):
unique_dir = tempfile.mkdtemp(prefix="hive-data-load-")
hive_args += ' --hiveconf "mapreduce.cluster.local.dir=%s"' % unique_dir
output_file = file_name + ".log"
hive_cmd = "{0} {1} -f {2}".format(HIVE_CMD, hive_args, file_name)
is_success = exec_cmd(hive_cmd, exit_on_error=False, out_file=output_file)
if unique_dir:
if is_success:"Finished execution of hive SQL: {0}".format(file_name))
else:"Error executing hive SQL: {0} See: {1}".format(file_name, \
return is_success
def exec_hbase_query_from_file(file_name):
if not os.path.exists(file_name): return
hbase_cmd = "hbase shell %s" % file_name'Executing HBase Command: %s' % hbase_cmd)
exec_cmd(hbase_cmd, error_msg='Error executing hbase create commands')
# KERBEROS TODO: fails when kerberized and impalad principal isn't "impala"
def exec_impala_query_from_file(file_name):
"""Execute each query in an Impala query file individually"""
if not os.path.exists(file_name):"Error: File {0} not found".format(file_name))
return False"Beginning execution of impala SQL on {0}: {1}".format(
options.impalad, file_name))
is_success = True
impala_client = ImpalaBeeswaxClient(options.impalad, use_kerberos=options.use_kerberos)
output_file = file_name + ".log"
query = None
with open(output_file, 'w') as out_file:
with open(file_name, 'r+') as query_file:
queries = sqlparse.split(
for query in queries:
query = sqlparse.format(query.rstrip(';'), strip_comments=True)
if query.strip() != "":
result = impala_client.execute(query)
out_file.write("{0}\n{1}\n".format(query, result))
except Exception as e:
if query:
out_file.write("ERROR: {0}\n".format(query))
out_file.write("Encounter errors before parsing any queries.\n")
is_success = False
if is_success:"Finished execution of impala SQL: {0}".format(file_name))
else:"Error executing impala SQL: {0} See: {1}".format(file_name, \
return is_success
def run_dataset_preload(dataset):
"""Execute a preload script if present in dataset directory. E.g. to generate data
before loading"""
dataset_preload_script = os.path.join(DATASET_DIR, dataset, "preload")
if os.path.exists(dataset_preload_script):"Running preload script for " + dataset)
if options.scale_factor > 1:
dataset_preload_script += " " + str(options.scale_factor)
exec_cmd(dataset_preload_script, error_msg="Error executing preload script for " + dataset,
def generate_schema_statements(workload):
generate_cmd = GENERATE_SCHEMA_CMD % (options.exploration_strategy, workload,
if options.table_names:
generate_cmd += " --table_names=%s" % options.table_names
if options.force_reload:
generate_cmd += " --force_reload"
if options.table_formats:
generate_cmd += " --table_formats=%s" % options.table_formats
if options.hive_warehouse_dir is not None:
generate_cmd += " --hive_warehouse_dir=%s" % options.hive_warehouse_dir
if options.hdfs_namenode is not None:
generate_cmd += " --hdfs_namenode=%s" % options.hdfs_namenode
generate_cmd += " --backend=%s" % options.impalad'Executing Generate Schema Command: ' + generate_cmd)
schema_cmd = os.path.join(TESTDATA_BIN_DIR, generate_cmd)
error_msg = 'Error generating schema statements for workload: ' + workload
exec_cmd(schema_cmd, error_msg=error_msg)
def get_dataset_for_workload(workload):
dimension_file_name = os.path.join(WORKLOAD_DIR, workload,
'%s_dimensions.csv' % workload)
if not os.path.isfile(dimension_file_name):
LOG.error('Dimension file not found: ' + dimension_file_name)
with open(dimension_file_name, 'rb') as input_file:
match ='dataset:\s*([\w\-\.]+)',
if match:
LOG.error('Dimension file does not contain dataset for workload \'%s\'' % (workload))
def copy_avro_schemas_to_hdfs(schemas_dir):
"""Recursively copies all of schemas_dir to the test warehouse."""
if not os.path.exists(schemas_dir):'Avro schema dir (%s) does not exist. Skipping copy to HDFS.' % schemas_dir)
exec_hadoop_fs_cmd("-mkdir -p " + options.hive_warehouse_dir)
exec_hadoop_fs_cmd("-put -f %s %s/" % (schemas_dir, options.hive_warehouse_dir))
def exec_hadoop_fs_cmd(args, exit_on_error=True):
cmd = "%s fs %s" % (HADOOP_CMD, args)"Executing Hadoop command: " + cmd)
exec_cmd(cmd, error_msg="Error executing Hadoop command, exiting",
def exec_query_files_parallel(thread_pool, query_files, execution_type):
"""Executes the query files provided using the execution engine specified
in parallel using the given thread pool. Aborts immediately if any execution
encounters an error."""
assert(execution_type == 'impala' or execution_type == 'hive')
if len(query_files) == 0: return
if execution_type == 'impala':
execution_function = exec_impala_query_from_file
elif execution_type == 'hive':
execution_function = exec_hive_query_from_file_beeline
for result in thread_pool.imap_unordered(execution_function, query_files):
if not result:
def impala_exec_query_files_parallel(thread_pool, query_files):
exec_query_files_parallel(thread_pool, query_files, 'impala')
def hive_exec_query_files_parallel(thread_pool, query_files):
exec_query_files_parallel(thread_pool, query_files, 'hive')
def main():
logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%H:%M:%S')
# Having the actual command line at the top of each data-load-* log can help
# when debugging dataload issues.
LOG.debug(' '.join(sys.argv))
all_workloads = available_workloads(WORKLOAD_DIR)
workloads = []
if options.workloads is None:
LOG.error("At least one workload name must be specified.")
elif options.workloads == 'all':'Loading data for all workloads.')
workloads = all_workloads
workloads = options.workloads.split(",")
validate_workloads(all_workloads, workloads)'Starting data load for the following workloads: ' + ', '.join(workloads))'Running with {0} threads'.format(options.num_processes))
# Note: The processes are in whatever the caller's directory is, so all paths
# passed to the pool need to be absolute paths. This will allow the pool
# to be used for different workloads (and thus different directories)
# simultaneously.
thread_pool = ThreadPool(processes=options.num_processes)
loading_time_map = collections.defaultdict(float)
for workload in workloads:
start_time = time.time()
dataset = get_dataset_for_workload(workload)
# This script is tightly coupled with testdata/bin/
# Specifically, this script is expecting the following:
# 1. generates files and puts them in the
# directory ${IMPALA_DATA_LOADING_SQL_DIR}/${workload}
# (e.g. ${IMPALA_HOME}/logs/data_loading/sql/tpch)
# 2. populates the subdirectory
# avro_schemas/${workload} with JSON files specifying the Avro schema for the
# tables being loaded.
# 3. uses a particular naming scheme to distinguish
# between SQL files of different load phases.
# Using the following variables:
# workload_exploration = ${workload}-${exploration_strategy} and
# file_format_suffix = ${file_format}-${codec}-${compression_type}
# A. Impala table creation scripts run in Impala to create tables, partitions,
# and views. There is one for each file format. They take the form:
# create-${workload_exploration}-impala-generated-${file_format_suffix}.sql
# B. Hive creation/load scripts run in Hive to load data into tables and create
# tables or views that Impala does not support. There is one for each
# file format. They take the form:
# load-${workload_exploration}-hive-generated-${file_format_suffix}.sql
# C. HBase creation script runs through the hbase commandline to create
# HBase tables. (Only generated if loading HBase table.) It takes the form:
# load-${workload_exploration}-hbase-generated.create
# D. HBase postload script runs through the hbase commandline to flush
# HBase tables. (Only generated if loading HBase table.) It takes the form:
# post-load-${workload_exploration}-hbase-generated.sql
# E. Impala load scripts run in Impala to load data. Only Parquet and Kudu
# are loaded through Impala. There is one for each of those formats loaded.
# They take the form:
# load-${workload_exploration}-impala-generated-${file_format_suffix}.sql
# F. Invalidation script runs through Impala to invalidate/refresh metadata
# for tables. It takes the form:
# invalidate-${workload_exploration}-impala-generated.sql
# Determine the directory from #1
sql_dir = os.path.join(SQL_OUTPUT_DIR, dataset)
assert os.path.isdir(sql_dir),\
("Could not find the generated SQL files for loading dataset '%s'.\
\nExpected to find the SQL files in: %s" % (dataset, sql_dir))
# Copy the avro schemas (see #2) into HDFS
avro_schemas_path = os.path.join(sql_dir, AVRO_SCHEMA_DIR)
# List all of the files in the sql directory to sort out the various types of
# files (see #3).
dataset_dir_contents = [os.path.join(sql_dir, f) for f in os.listdir(sql_dir)]
workload_exploration = "%s-%s" % (workload, options.exploration_strategy)
# Remove the AVRO_SCHEMA_DIR from the list of files
if os.path.exists(avro_schemas_path):
# Match for Impala create files (3.A)
impala_create_match = 'create-%s-impala-generated' % workload_exploration
# Match for Hive create/load files (3.B)
hive_load_match = 'load-%s-hive-generated' % workload_exploration
# Match for HBase creation script (3.C)
hbase_create_match = 'load-%s-hbase-generated.create' % workload_exploration
# Match for HBase post-load script (3.D)
hbase_postload_match = 'post-load-%s-hbase-generated.sql' % workload_exploration
# Match for Impala load scripts (3.E)
impala_load_match = 'load-%s-impala-generated' % workload_exploration
# Match for Impala invalidate script (3.F)
invalidate_match = 'invalidate-%s-impala-generated' % workload_exploration
impala_create_files = []
hive_load_text_files = []
hive_load_nontext_files = []
hbase_create_files = []
hbase_postload_files = []
impala_load_files = []
invalidate_files = []
for filename in dataset_dir_contents:
if impala_create_match in filename:
elif hive_load_match in filename:
if 'text-none-none' in filename:
elif hbase_create_match in filename:
elif hbase_postload_match in filename:
elif impala_load_match in filename:
elif invalidate_match in filename:
assert False, "Unexpected input file {0}".format(filename)
# Simple helper function to dump a header followed by the filenames
def log_file_list(header, file_list):
if (len(file_list) == 0): return
map(LOG.debug, map(os.path.basename, file_list))
log_file_list("Impala Create Files:", impala_create_files)
log_file_list("Hive Load Text Files:", hive_load_text_files)
log_file_list("Hive Load Non-Text Files:", hive_load_nontext_files)
log_file_list("HBase Create Files:", hbase_create_files)
log_file_list("HBase Post-Load Files:", hbase_postload_files)
log_file_list("Impala Load Files:", impala_load_files)
log_file_list("Impala Invalidate Files:", invalidate_files)
# Execute the data loading scripts.
# Creating tables in Impala has no dependencies, so we execute them first.
# HBase table inserts are done via hive, so the hbase tables need to be created before
# running the hive scripts. Some of the Impala inserts depend on hive tables,
# so they're done at the end. Finally, the Hbase Tables that have been filled with data
# need to be flushed.
impala_exec_query_files_parallel(thread_pool, impala_create_files)
# There should be at most one hbase creation script
assert(len(hbase_create_files) <= 1)
for hbase_create in hbase_create_files:
# If this is loading text tables plus multiple other formats, the text tables
# need to be loaded first
assert(len(hive_load_text_files) <= 1)
hive_exec_query_files_parallel(thread_pool, hive_load_text_files)
hive_exec_query_files_parallel(thread_pool, hive_load_nontext_files)
assert(len(hbase_postload_files) <= 1)
for hbase_postload in hbase_postload_files:
# Invalidate so that Impala sees the loads done by Hive before loading Parquet/Kudu
# Note: This only invalidates tables for this workload.
assert(len(invalidate_files) <= 1)
if impala_load_files:
impala_exec_query_files_parallel(thread_pool, invalidate_files)
impala_exec_query_files_parallel(thread_pool, impala_load_files)
# Final invalidate for this workload
impala_exec_query_files_parallel(thread_pool, invalidate_files)
loading_time_map[workload] = time.time() - start_time
total_time = 0.0
for workload, load_time in loading_time_map.iteritems():
total_time += load_time'Data loading for workload \'%s\' completed in: %.2fs'\
% (workload, load_time))'Total load time: %.2fs\n' % total_time)
if __name__ == "__main__": main()