testdata/datasets/tpch/preload

#!/bin/bash
# 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.

set -euo pipefail
. $IMPALA_HOME/bin/report_build_error.sh
setup_report_build_error

IMPALA_DATA=${IMPALA_HOME}/testdata/impala-data
TPC_H_DATA=${IMPALA_DATA}/tpch

SCALE_FACTOR=1
if [[ $# == 1 && $1 -gt 1 ]]
then
  SCALE_FACTOR=$1
  TPC_H_DATA=${TPC_H_DATA}${SCALE_FACTOR}
fi

if [ -f ${TPC_H_DATA}/SUCCESS ]; then
  echo "Reuse existing TPC-H data in ${TPC_H_DATA}"
  exit 0
fi

TPC_H_HOME=${IMPALA_TOOLCHAIN_PACKAGES_HOME}/tpc-h-${IMPALA_TPC_H_VERSION}
TPC_H_DBGEN=${TPC_H_HOME}/bin/dbgen

if [ ! -x ${TPC_H_DBGEN} ]; then
  echo "Could not find TPC-H data generator executable: ${TPC_H_DBGEN}"
  exit 1
fi

# Delete any preexisting data or symlinks
# Need to change permissions to work around an old TPC-H data tarball that had a
# non-writable top-level directory when extracted.
chmod +w ${TPC_H_DATA} || true
rm -rf ${TPC_H_DATA}
mkdir -p ${TPC_H_DATA}
# Create symlink if scale factor is 1
if [ ${SCALE_FACTOR} -eq 1 ]
then
  rm -rf ${TPC_H_DATA}${SCALE_FACTOR}
  ln -s ${TPC_H_DATA} ${TPC_H_DATA}${SCALE_FACTOR}
fi
cd ${TPC_H_DATA}

if [ -t 1 ]
then
  # Output is terminal, show progress verbosely
  VERBOSITY='-v'
else
  VERBOSITY=''
fi

function load-table {
  local chunk_number=$1
  local parallelism=$2
  if (( ${chunk_number} > ${parallelism} )); then
    echo "Invalid chunk number: ${chunk_number} parallelism: ${parallelism}"
  fi
  local tmp_log_file=$(mktemp)
  set +e
  if [[ ${parallelism} == 1 ]]; then
    # TPC-H doesn't allow specifying -C 1 to use a single chunk, so run without those
    # options
    ${TPC_H_DBGEN} ${VERBOSITY} -f -s ${SCALE_FACTOR}  2> ${tmp_log_file}
  else
    ${TPC_H_DBGEN} ${VERBOSITY} -f -s ${SCALE_FACTOR} \
        -C ${parallelism} -S ${chunk_number} 2> ${tmp_log_file}
  fi
  RET_VAL=$?
  set -e
  if [[ $RET_VAL != 0 ]]; then
    echo "Failed to generate chunk ${chunk_number}/${parallelism}"
    cat ${tmp_log_file}
    return 1
  fi
  rm -f ${tmp_log_file}
  echo "Generated chunk ${chunk_number}/${parallelism}"
}

# Use the regular generation process for scale factor 1 to avoid changes
# in the numbers of files for tests. If this is a higher scale factor,
# generate the tables in parallel to save time. The TPC-H dbgen utility
# has an ability to generate dataset in chunks in separate threads.
# Unfortunately, a couple tables (nation and region) are not generated
# in chunks. Instead, they are duplicated across all the chunks. To avoid
# races, produce the different chunks in different directories, then move
# them into place. Use one chunk per CPU.
if [[ $SCALE_FACTOR == 1 ]]; then
  parallelism=1
else
  parallelism=$(nproc)
fi

echo "Generating TPC-H data with parallelism ${parallelism} into ${TPC_H_DATA}"
pids=()
i=0
for chunk in $(seq ${parallelism}); do
  mkdir chunk_${chunk}
  pushd chunk_${chunk} > /dev/null
  load-table ${chunk} ${parallelism} &
  popd > /dev/null
  pids[${i}]=$!
  let i=i+1
done

# wait for all pids
for pid in ${pids[*]}; do
  wait $pid
done

echo "Done generating data"

tables=(
  "customer"
  "lineitem"
  "nation"
  "orders"
  "part"
  "partsupp"
  "region"
  "supplier"
)

# Verify that all the expected files exit. Impala expects each table to be in its own
# subdirectory, so this moves them.
for table in ${tables[*]}; do
  mkdir -p ${table}
  # If this didn't go parallel or are a non-parallel table (region or nation), the data
  # file is in the first chunk and the file does not have a suffix.
  if [[ ${parallelism} == 1 || ${table} == "region" || ${table} == "nation" ]]; then
    data_file="chunk_1/${table}.tbl"
    if ! [[ -f ${data_file} ]]; then
      echo "Missing expected data file ${data_file}"
      exit 1
    fi
    mv ${data_file} ${table}
  else
    # Other tables went parallel and have a file in each chunk
    for i in $(seq 1 ${parallelism}); do
      data_file="chunk_${i}/${table}.tbl.${i}"
      if ! [[ -f ${data_file} ]]; then
        echo "Missing expected data file ${data_file}"
        exit 1
      fi
      mv ${data_file} ${table}
    done
  fi
done

# Cleanup directories
rm -rf chunk_*

touch SUCCESS