// 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.

#include "kudu/common/partition_pruner.h"

#include <algorithm>
#include <boost/optional.hpp>
#include <cstring>
#include <memory>
#include <numeric>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>

#include "kudu/common/key_encoder.h"
#include "kudu/common/key_util.h"
#include "kudu/common/partition.h"
#include "kudu/common/scan_spec.h"
#include "kudu/common/schema.h"
#include "kudu/gutil/map-util.h"
#include "kudu/gutil/strings/join.h"
#include "kudu/gutil/strings/substitute.h"

using boost::optional;
using std::distance;
using std::find;
using std::get;
using std::iota;
using std::lower_bound;
using std::make_tuple;
using std::memcpy;
using std::min;
using std::move;
using std::string;
using std::tuple;
using std::unique_ptr;
using std::unordered_map;
using std::vector;

namespace kudu {
namespace {

// Returns true if the partition schema's range columns are a prefix of the
// primary key columns.
bool AreRangeColumnsPrefixOfPrimaryKey(const Schema& schema,
                                       const vector<ColumnId>& range_columns) {
  CHECK(range_columns.size() <= schema.num_key_columns());
  for (int32_t col_idx = 0; col_idx < range_columns.size(); col_idx++) {
    if (schema.column_id(col_idx) != range_columns[col_idx]) {
      return false;
    }
  }
  return true;
}

// Translates the scan primary key bounds into range keys. This should only be
// used when the range columns are a prefix of the primary key columns.
void EncodeRangeKeysFromPrimaryKeyBounds(const Schema& schema,
                                         const ScanSpec& scan_spec,
                                         size_t num_range_columns,
                                         string* range_key_start,
                                         string* range_key_end) {
  if (scan_spec.lower_bound_key() == nullptr && scan_spec.exclusive_upper_bound_key() == nullptr) {
    // Don't bother if there are no lower and upper PK bounds
    return;
  }

  if (num_range_columns == schema.num_key_columns()) {
    // The range columns are the primary key columns, so the range key is the
    // primary key.
    if (scan_spec.lower_bound_key() != nullptr) {
      *range_key_start = scan_spec.lower_bound_key()->encoded_key().ToString();
    }
    if (scan_spec.exclusive_upper_bound_key() != nullptr) {
      *range_key_end = scan_spec.exclusive_upper_bound_key()->encoded_key().ToString();
    }
  } else {
    // The range columns are a prefix of the primary key columns. Copy
    // the column values over to a row, and then encode the row as a range key.

    vector<int32_t> col_idxs(num_range_columns);
    iota(col_idxs.begin(), col_idxs.end(), 0);

    unique_ptr<uint8_t[]> buf(new uint8_t[schema.key_byte_size()]);
    ContiguousRow row(&schema, buf.get());

    if (scan_spec.lower_bound_key() != nullptr) {
      for (int32_t idx : col_idxs) {
        memcpy(row.mutable_cell_ptr(idx),
               scan_spec.lower_bound_key()->raw_keys()[idx],
               schema.column(idx).type_info()->size());
      }
      key_util::EncodeKey(col_idxs, row, range_key_start);
    }

    if (scan_spec.exclusive_upper_bound_key() != nullptr) {
      for (int32_t idx : col_idxs) {
        memcpy(row.mutable_cell_ptr(idx),
               scan_spec.exclusive_upper_bound_key()->raw_keys()[idx],
               schema.column(idx).type_info()->size());
      }
      key_util::EncodeKey(col_idxs, row, range_key_end);
    }
  }
}

// Push the scan predicates into the range keys.
void EncodeRangeKeysFromPredicates(const Schema& schema,
                                   const unordered_map<string, ColumnPredicate>& predicates,
                                   const vector<ColumnId>& range_columns,
                                   string* range_key_start,
                                   string* range_key_end) {
  // Find the column indexes of the range columns.
  vector<int32_t> col_idxs;
  col_idxs.reserve(range_columns.size());
  for (ColumnId column : range_columns) {
    int32_t col_idx = schema.find_column_by_id(column);
    CHECK(col_idx != Schema::kColumnNotFound);
    CHECK(col_idx < schema.num_key_columns());
    col_idxs.push_back(col_idx);
  }

  // Arenas must be at least the minimum chunk size, and we require at least
  // enough space for the range key columns.
  Arena arena(max<size_t>(Arena::kMinimumChunkSize, schema.key_byte_size()), 4096);
  uint8_t* buf = static_cast<uint8_t*>(CHECK_NOTNULL(arena.AllocateBytes(schema.key_byte_size())));
  ContiguousRow row(&schema, buf);

  if (key_util::PushLowerBoundKeyPredicates(col_idxs, predicates, &row, &arena) > 0) {
    key_util::EncodeKey(col_idxs, row, range_key_start);
  }

  if (key_util::PushUpperBoundKeyPredicates(col_idxs, predicates, &row, &arena) > 0) {
    key_util::EncodeKey(col_idxs, row, range_key_end);
  }
}
} // anonymous namespace

void PartitionPruner::Init(const Schema& schema,
                           const PartitionSchema& partition_schema,
                           const ScanSpec& scan_spec) {
  // If we can already short circuit the scan we don't need to bother with
  // partition pruning. This also allows us to assume some invariants of the
  // scan spec, such as no None predicates and that the lower bound PK < upper
  // bound PK.
  if (scan_spec.CanShortCircuit()) { return; }

  // Build a set of partition key ranges which cover the tablets necessary for
  // the scan.
  //
  // Example predicate sets and resulting partition key ranges, based on the
  // following tablet schema:
  //
  // CREATE TABLE t (a INT32, b INT32, c INT32) PRIMARY KEY (a, b, c)
  // DISTRIBUTE BY RANGE (c)
  //               HASH (a) INTO 2 BUCKETS
  //               HASH (b) INTO 3 BUCKETS;
  //
  // Assume that hash(0) = 0 and hash(2) = 2.
  //
  // | Predicates | Partition Key Ranges                                   |
  // +------------+--------------------------------------------------------+
  // | a = 0      | [(bucket=0, bucket=2, c=0), (bucket=0, bucket=2, c=1)) |
  // | b = 2      |                                                        |
  // | c = 0      |                                                        |
  // +------------+--------------------------------------------------------+
  // | a = 0      | [(bucket=0, bucket=2), (bucket=0, bucket=3))           |
  // | b = 2      |                                                        |
  // +------------+--------------------------------------------------------+
  // | a = 0      | [(bucket=0, bucket=0, c=0), (bucket=0, bucket=0, c=1)) |
  // | c = 0      | [(bucket=0, bucket=1, c=0), (bucket=0, bucket=1, c=1)) |
  // |            | [(bucket=0, bucket=2, c=0), (bucket=0, bucket=2, c=1)) |
  // +------------+--------------------------------------------------------+
  // | b = 2      | [(bucket=0, bucket=2, c=0), (bucket=0, bucket=2, c=1)) |
  // | c = 0      | [(bucket=1, bucket=2, c=0), (bucket=1, bucket=2, c=1)) |
  // +------------+--------------------------------------------------------+
  // | a = 0      | [(bucket=0), (bucket=1))                               |
  // +------------+--------------------------------------------------------+
  // | b = 2      | [(bucket=0, bucket=2), (bucket=0, bucket=3))           |
  // |            | [(bucket=1, bucket=2), (bucket=1, bucket=3))           |
  // +------------+--------------------------------------------------------+
  // | c = 0      | [(bucket=0, bucket=0, c=0), (bucket=0, bucket=0, c=1)) |
  // |            | [(bucket=0, bucket=1, c=0), (bucket=0, bucket=1, c=1)) |
  // |            | [(bucket=0, bucket=2, c=0), (bucket=0, bucket=2, c=1)) |
  // |            | [(bucket=1, bucket=0, c=0), (bucket=1, bucket=0, c=1)) |
  // |            | [(bucket=1, bucket=1, c=0), (bucket=1, bucket=1, c=1)) |
  // |            | [(bucket=1, bucket=2, c=0), (bucket=1, bucket=2, c=1)) |
  // +------------+--------------------------------------------------------+
  // | None       | [(), ())                                               |
  //
  // If the partition key is considered as a sequence of the hash bucket
  // components and a range component, then a few patterns emerge from the
  // examples above:
  //
  // 1) The partition keys are truncated after the final constrained component
  //    (hash bucket components are constrained when the scan is limited to a
  //    single bucket via equality predicates on that component, while range
  //    components are constrained if they have an upper or lower bound via
  //    range or equality predicates on that component).
  //
  // 2) If the final constrained component is a hash bucket, then the
  //    corresponding bucket in the upper bound is incremented in order to make
  //    it an exclusive key.
  //
  // 3) The number of partition key ranges in the result is equal to the product
  //    of the number of buckets of each unconstrained hash component which come
  //    before a final constrained component. If there are no unconstrained hash
  //    components, then the number of partition key ranges is one.

  // Step 1: Build the range portion of the partition key.
  string range_lower_bound;
  string range_upper_bound;
  const vector<ColumnId>& range_columns = partition_schema.range_schema_.column_ids;
  if (!range_columns.empty()) {
    if (AreRangeColumnsPrefixOfPrimaryKey(schema, range_columns)) {
      EncodeRangeKeysFromPrimaryKeyBounds(schema,
                                          scan_spec,
                                          range_columns.size(),
                                          &range_lower_bound,
                                          &range_upper_bound);
    } else {
      EncodeRangeKeysFromPredicates(schema,
                                    scan_spec.predicates(),
                                    range_columns,
                                    &range_lower_bound,
                                    &range_upper_bound);
    }
  }

  // Step 2: Create the hash bucket portion of the partition key.

  // The list of hash buckets per hash component, or none if the component is
  // not constrained.
  vector<optional<uint32_t>> hash_buckets;
  hash_buckets.reserve(partition_schema.hash_bucket_schemas_.size());
  for (int hash_idx = 0; hash_idx < partition_schema.hash_bucket_schemas_.size(); hash_idx++) {
    const auto& hash_bucket_schema = partition_schema.hash_bucket_schemas_[hash_idx];
    string encoded_columns;
    bool can_prune = true;
    for (int col_offset = 0; col_offset < hash_bucket_schema.column_ids.size(); col_offset++) {
      const ColumnSchema& column = schema.column_by_id(hash_bucket_schema.column_ids[col_offset]);
      const ColumnPredicate* predicate = FindOrNull(scan_spec.predicates(), column.name());
      if (predicate == nullptr || predicate->predicate_type() != PredicateType::Equality) {
        can_prune = false;
        break;
      }

      const KeyEncoder<string>& encoder = GetKeyEncoder<string>(column.type_info());
      encoder.Encode(predicate->raw_lower(),
                     col_offset + 1 == hash_bucket_schema.column_ids.size(),
                     &encoded_columns);
    }
    if (can_prune) {
      hash_buckets.push_back(partition_schema.BucketForEncodedColumns(encoded_columns,
                                                                      hash_bucket_schema));
    } else {
      hash_buckets.push_back(boost::none);
    }
  }

  // The index of the final constrained component in the partition key.
  int constrained_index;
  if (!range_lower_bound.empty() || !range_upper_bound.empty()) {
    // The range component is constrained.
    constrained_index = partition_schema.hash_bucket_schemas_.size();
  } else {
    // Search the hash bucket constraints from right to left, looking for the
    // first constrained component.
    constrained_index = partition_schema.hash_bucket_schemas_.size() -
                        distance(hash_buckets.rbegin(),
                                 find_if(hash_buckets.rbegin(),
                                         hash_buckets.rend(),
                                         [] (const optional<uint32_t>& x) { return x; }));
  }

  // Build up a set of partition key ranges out of the hash components.
  //
  // Each constrained hash component simply appends its bucket number to the
  // partition key ranges (possibly incrementing the upper bound by one bucket
  // number if this is the final constraint, see note 2 in the example above).
  //
  // Each unconstrained hash component results in creating a new partition key
  // range for each bucket of the hash component.
  vector<tuple<string, string>> partition_key_ranges(1);
  const KeyEncoder<string>& hash_encoder = GetKeyEncoder<string>(GetTypeInfo(UINT32));
  for (int hash_idx = 0; hash_idx < constrained_index; hash_idx++) {
    // This is the final partition key component if this is the final constrained
    // bucket, and the range upper bound is empty. In this case we need to
    // increment the bucket on the upper bound to convert from inclusive to
    // exclusive.
    bool is_last = hash_idx + 1 == constrained_index && range_upper_bound.empty();

    if (hash_buckets[hash_idx]) {
      // This hash component is constrained by equality predicates to a single
      // hash bucket.
      uint32_t bucket = *hash_buckets[hash_idx];
      uint32_t bucket_upper = is_last ? bucket + 1 : bucket;
      for (auto& partition_key_range : partition_key_ranges) {
        hash_encoder.Encode(&bucket, &get<0>(partition_key_range));
        hash_encoder.Encode(&bucket_upper, &get<1>(partition_key_range));
      }
    } else {
      const auto& hash_bucket_schema = partition_schema.hash_bucket_schemas_[hash_idx];
      // Add a partition key range for each possible hash bucket.
      vector<tuple<string, string>> new_partition_key_ranges;
      for (const auto& partition_key_range : partition_key_ranges) {
        for (uint32_t bucket = 0; bucket < hash_bucket_schema.num_buckets; bucket++) {
          uint32_t bucket_upper = is_last ? bucket + 1 : bucket;
          string lower = get<0>(partition_key_range);
          string upper = get<1>(partition_key_range);
          hash_encoder.Encode(&bucket, &lower);
          hash_encoder.Encode(&bucket_upper, &upper);
          new_partition_key_ranges.push_back(make_tuple(move(lower), move(upper)));
        }
      }
      partition_key_ranges.swap(new_partition_key_ranges);
    }
  }

  // Step 3: append the (possibly empty) range bounds to the partition key ranges.
  for (auto& range : partition_key_ranges) {
    get<0>(range).append(range_lower_bound);
    get<1>(range).append(range_upper_bound);
  }

  // Step 4: remove all partition key ranges past the scan spec's upper bound partition key.
  if (!scan_spec.exclusive_upper_bound_partition_key().empty()) {
    for (auto range = partition_key_ranges.rbegin();
         range != partition_key_ranges.rend();
         range++) {
      if (!get<1>(*range).empty() &&
          scan_spec.exclusive_upper_bound_partition_key() >= get<1>(*range)) {
        break;
      }
      if (scan_spec.exclusive_upper_bound_partition_key() <= get<0>(*range)) {
        partition_key_ranges.pop_back();
      } else {
        get<1>(*range) = scan_spec.exclusive_upper_bound_partition_key();
      }
    }
  }

  // Step 5: Reverse the order of the partition key ranges, so that it is
  // efficient to remove the partition key ranges from the vector in ascending order.
  partition_key_ranges_.resize(partition_key_ranges.size());
  move(partition_key_ranges.rbegin(), partition_key_ranges.rend(), partition_key_ranges_.begin());

  // Step 6: Remove all partition key ranges before the scan spec's lower bound partition key.
  if (!scan_spec.lower_bound_partition_key().empty()) {
    RemovePartitionKeyRange(scan_spec.lower_bound_partition_key());
  }
}

bool PartitionPruner::HasMorePartitionKeyRanges() const {
  return !partition_key_ranges_.empty();
}

const string& PartitionPruner::NextPartitionKey() const {
  CHECK(HasMorePartitionKeyRanges());
  return get<0>(partition_key_ranges_.back());
}

void PartitionPruner::RemovePartitionKeyRange(const string& upper_bound) {
  if (upper_bound.empty()) {
    partition_key_ranges_.clear();
    return;
  }

  for (auto range = partition_key_ranges_.rbegin();
       range != partition_key_ranges_.rend();
       range++) {
    if (upper_bound <= get<0>(*range)) { break; }
    if (get<1>(*range).empty() || upper_bound < get<1>(*range)) {
      get<0>(*range) = upper_bound;
    } else {
      partition_key_ranges_.pop_back();
    }
  }
}

bool PartitionPruner::ShouldPruneForTests(const Partition& partition) const {
  // range is an iterator that points to the first partition key range which
  // overlaps or is greater than the partition.
  auto range = lower_bound(partition_key_ranges_.rbegin(), partition_key_ranges_.rend(), partition,
    [] (const tuple<string, string>& scan_range, const Partition& partition) {
      // return true if scan_range < partition
      const string& scan_upper = get<1>(scan_range);
      return !scan_upper.empty() && scan_upper <= partition.partition_key_start();
    });

  return range == partition_key_ranges_.rend() ||
         (!partition.partition_key_end().empty() &&
          partition.partition_key_end() <= get<0>(*range));
}

string PartitionPruner::ToString(const Schema& schema,
                                 const PartitionSchema& partition_schema) const {
  vector<string> strings;
  for (auto range = partition_key_ranges_.rbegin();
       range != partition_key_ranges_.rend();
       range++) {
    strings.push_back(strings::Substitute(
          "[($0), ($1))",
          get<0>(*range).empty() ? "<start>" :
              partition_schema.PartitionKeyDebugString(get<0>(*range), schema),
          get<1>(*range).empty() ? "<end>" :
              partition_schema.PartitionKeyDebugString(get<1>(*range), schema)));
  }

  return JoinStrings(strings, ", ");
}

} // namespace kudu