bin/load-data.py

#!/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
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# 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.
from __future__ import absolute_import, division, print_function
import collections
import getpass
import logging
import multiprocessing
import os
import re
import sqlparse
import subprocess
import sys
import time
import traceback

from optparse import OptionParser
from tests.beeswax.impala_beeswax import ImpalaBeeswaxClient
from multiprocessing.pool import ThreadPool

LOG = logging.getLogger('load-data.py')

parser = OptionParser()
parser.add_option("-e", "--exploration_strategy", dest="exploration_strategy",
                  default="core",
                  help="The exploration strategy for schema gen: 'core', "\
                  "'pairwise', or 'exhaustive'")
parser.add_option("--hive_warehouse_dir", dest="hive_warehouse_dir",
                  default="/test-warehouse",
                  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",
                  default="localhost:11050",
                  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",
                  default=os.environ['IMPALA_WORKLOAD_DIR'],
                  help="Directory that contains Impala workloads")
parser.add_option("--dataset_dir", dest="dataset_dir",
                  default=os.environ['IMPALA_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()

SQL_OUTPUT_DIR = os.environ['IMPALA_DATA_LOADING_SQL_DIR']
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 = "generate-schema-statements.py --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")
    exit(1)
  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))
      sys.exit(1)

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 = subprocess.call(cmd, shell=True, stderr=f, stdout=f)
  else:
    ret_val = subprocess.call(cmd, shell=True)
  if ret_val != 0:
    if error_msg: LOG.info(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):
    LOG.info("Error: File {0} not found".format(file_name))
    return False

  LOG.info("Beginning execution of hive SQL: {0}".format(file_name))

  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 is_success:
    LOG.info("Finished execution of hive SQL: {0}".format(file_name))
  else:
    LOG.info("Error executing hive SQL: {0} See: {1}".format(file_name, \
             output_file))

  return is_success


def exec_hbase_query_from_file(file_name, step_name):
  if not os.path.exists(file_name): return
  LOG.info('Begin step "%s".' % step_name)
  start_time = time.time()
  hbase_cmd = "hbase shell %s" % file_name
  LOG.info('Executing HBase Command: %s' % hbase_cmd)
  exec_cmd(hbase_cmd, error_msg='Error executing hbase create commands')
  total_time = time.time() - start_time
  LOG.info('End step "%s". Total time: %.2fs\n' % (step_name, total_time))


# 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):
    LOG.info("Error: File {0} not found".format(file_name))
    return False

  LOG.info("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:
    try:
      impala_client.connect()
      with open(file_name, 'r+') as query_file:
        queries = sqlparse.split(query_file.read())
        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))
      else:
        out_file.write("Encounter errors before parsing any queries.\n")
      traceback.print_exc(file=out_file)
      is_success = False

  if is_success:
    LOG.info("Finished execution of impala SQL: {0}".format(file_name))
  else:
    LOG.info("Error executing impala SQL: {0} See: {1}".format(file_name, \
             output_file))

  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):
    LOG.info("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,
        exit_on_error=True)

def generate_schema_statements(workload):
  generate_cmd = GENERATE_SCHEMA_CMD % (options.exploration_strategy, workload,
                                        options.scale_factor)
  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
  LOG.info('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)
    sys.exit(1)
  with open(dimension_file_name, 'rb') as input_file:
    match = re.search('dataset:\s*([\w\-\.]+)', input_file.read())
    if match:
      return match.group(1)
    else:
      LOG.error('Dimension file does not contain dataset for workload \'%s\'' % (workload))
      sys.exit(1)

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):
    LOG.info('Avro schema dir (%s) does not exist. Skipping copy to HDFS.' % schemas_dir)
    return

  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)
  LOG.info("Executing Hadoop command: " + cmd)
  exec_cmd(cmd, error_msg="Error executing Hadoop command, exiting",
      exit_on_error=exit_on_error)


def exec_query_files_parallel(thread_pool, query_files, execution_type, step_name):
  """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

  LOG.info('Begin step "%s".' % step_name)
  start_time = time.time()
  for result in thread_pool.imap_unordered(execution_function, query_files):
    if not result:
      thread_pool.terminate()
      sys.exit(1)
  total_time = time.time() - start_time
  LOG.info('End step "%s". Total time: %.2fs\n' % (step_name, total_time))


def impala_exec_query_files_parallel(thread_pool, query_files, step_name):
  exec_query_files_parallel(thread_pool, query_files, 'impala', step_name)


def hive_exec_query_files_parallel(thread_pool, query_files, step_name):
  exec_query_files_parallel(thread_pool, query_files, 'hive', step_name)


def main():
  logging.basicConfig(format='%(asctime)s %(message)s', datefmt='%H:%M:%S')
  LOG.setLevel(logging.DEBUG)

  # 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.")
    parser.print_help()
    sys.exit(1)
  elif options.workloads == 'all':
    LOG.info('Loading data for all workloads.')
    workloads = all_workloads
  else:
    workloads = options.workloads.split(",")
    validate_workloads(all_workloads, workloads)

  LOG.info('Starting data load for the following workloads: ' + ', '.join(workloads))
  LOG.info('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)
    run_dataset_preload(dataset)
    # This script is tightly coupled with testdata/bin/generate-schema-statements.py
    # Specifically, this script is expecting the following:
    # 1. generate-schema-statements.py generates files and puts them in the
    #    directory ${IMPALA_DATA_LOADING_SQL_DIR}/${workload}
    #    (e.g. ${IMPALA_HOME}/logs/data_loading/sql/tpch)
    # 2. generate-schema-statements.py populates the subdirectory
    #    avro_schemas/${workload} with JSON files specifying the Avro schema for the
    #    tables being loaded.
    # 3. generate-schema-statements.py 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
    generate_schema_statements(workload)

    # 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)
    copy_avro_schemas_to_hdfs(avro_schemas_path)

    # 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):
      dataset_dir_contents.remove(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_orc_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:
        impala_create_files.append(filename)
      elif hive_load_match in filename:
        if 'text-none-none' in filename:
          hive_load_text_files.append(filename)
        elif 'orc-def-block' in filename:
          hive_load_orc_files.append(filename)
        else:
          hive_load_nontext_files.append(filename)
      elif hbase_create_match in filename:
        hbase_create_files.append(filename)
      elif hbase_postload_match in filename:
        hbase_postload_files.append(filename)
      elif impala_load_match in filename:
        impala_load_files.append(filename)
      elif invalidate_match in filename:
        invalidate_files.append(filename)
      else:
        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
      LOG.debug(header)
      list(map(LOG.debug, list(map(os.path.basename, file_list))))
      LOG.debug("\n")

    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 Orc Files:", hive_load_orc_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, "Impala Create")

    # There should be at most one hbase creation script
    assert(len(hbase_create_files) <= 1)
    for hbase_create in hbase_create_files:
      exec_hbase_query_from_file(hbase_create, "HBase Create")

    # 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 Load Text")
    # IMPALA-9923: Run ORC serially separately from other non-text formats. This hacks
    # around flakiness seen when loading this in parallel (see IMPALA-12630 comments for
    # broken tests). This should be removed as soon as possible.
    assert(len(hive_load_orc_files) <= 1)
    hive_exec_query_files_parallel(thread_pool, hive_load_orc_files, "Hive Load ORC")

    # Load all non-text formats (goes parallel)
    hive_exec_query_files_parallel(thread_pool, hive_load_nontext_files,
        "Hive Load Non-Text")

    assert(len(hbase_postload_files) <= 1)
    for hbase_postload in hbase_postload_files:
      exec_hbase_query_from_file(hbase_postload, "HBase Post-Load")

    # 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 Invalidate 1")
      impala_exec_query_files_parallel(thread_pool, impala_load_files, "Impala Load")
    # Final invalidate for this workload
    impala_exec_query_files_parallel(thread_pool, invalidate_files, "Impala Invalidate 2")
    loading_time_map[workload] = time.time() - start_time

  total_time = 0.0
  thread_pool.close()
  thread_pool.join()
  for workload, load_time in loading_time_map.items():
    total_time += load_time
    LOG.info('Data loading for workload \'%s\' completed in: %.2fs'\
        % (workload, load_time))
  LOG.info('Total load time: %.2fs\n' % total_time)

if __name__ == "__main__": main()