src/trace_processor/db/storage/numeric_storage.cc

/*
 * Copyright (C) 2023 The Android Open Source Project
 *
 * Licensed 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 "src/trace_processor/db/storage/numeric_storage.h"

#include <variant>

namespace perfetto {
namespace trace_processor {
namespace storage {

namespace {

// Templated part of FastPathComparison.
template <typename T>
inline void TypedFastPathComparison(std::optional<NumericValue> val,
                                    FilterOp op,
                                    const T* start,
                                    uint32_t num_elements,
                                    BitVector::Builder& builder) {
  if (!val) {
    builder.Skip(num_elements);
    return;
  }
  std::visit(
      [val, start, num_elements, &builder](auto comparator) {
        T typed_val = std::get<T>(*val);
        for (uint32_t i = 0; i < num_elements; i += BitVector::kBitsInWord) {
          uint64_t word = 0;
          // This part should be optimised by SIMD and is expected to be fast.
          for (uint32_t k = 0; k < BitVector::kBitsInWord; ++k) {
            bool comp_result = comparator(start[i + k], typed_val);
            word |= static_cast<uint64_t>(comp_result) << k;
          }
          builder.AppendWord(word);
        }
      },
      GetFilterOpVariant<T>(op));
}

// Templated part of SlowPathComparison.
template <typename T>
inline void TypedSlowPathComparison(std::optional<NumericValue> val,
                                    FilterOp op,
                                    const T* start,
                                    uint32_t num_elements,
                                    BitVector::Builder& builder) {
  if (!val) {
    builder.Skip(num_elements);
    return;
  }
  std::visit(
      [val, start, num_elements, &builder](auto comparator) {
        T typed_val = std::get<T>(*val);
        for (uint32_t i = 0; i < num_elements; ++i) {
          builder.Append(comparator(start[i], typed_val));
        }
      },
      GetFilterOpVariant<T>(op));
}

}  // namespace

void NumericStorage::StableSort(uint32_t* rows, uint32_t rows_size) const {
  NumericValue val = *GetNumericTypeVariant(type_, SqlValue::Long(0));
  std::visit(
      [this, &rows, rows_size](auto val_data) {
        using T = decltype(val_data);
        const T* typed_start = static_cast<const T*>(data_);
        std::stable_sort(rows, rows + rows_size,
                         [typed_start](uint32_t a_idx, uint32_t b_idx) {
                           T first_val = typed_start[a_idx];
                           T second_val = typed_start[b_idx];
                           return first_val < second_val;
                         });
      },
      val);
}

void NumericStorage::Sort(uint32_t*, uint32_t) const {}

// Responsible for invoking templated version of FastPathComparison.
void NumericStorage::LinearSearchAligned(FilterOp op,
                                         SqlValue sql_val,
                                         uint32_t offset,
                                         uint32_t num_elements,
                                         BitVector::Builder& builder) const {
  PERFETTO_DCHECK(num_elements % BitVector::kBitsInWord == 0);
  std::optional<NumericValue> val = GetNumericTypeVariant(type_, sql_val);

  // If the value is invalid we should just ignore those elements.
  if (!val.has_value() || op == FilterOp::kIsNotNull ||
      op == FilterOp::kIsNull || op == FilterOp::kGlob) {
    builder.Skip(num_elements);
    return;
  }
  std::visit(
      [this, op, offset, num_elements, &builder](auto num_val) {
        using T = decltype(num_val);
        auto* typed_start = static_cast<const T*>(data_) + offset;
        TypedFastPathComparison(num_val, op, typed_start, num_elements,
                                builder);
      },
      *val);
}

// Responsible for invoking templated version of SlowPathComparison.
void NumericStorage::LinearSearchUnaligned(FilterOp op,
                                           SqlValue sql_val,
                                           uint32_t offset,
                                           uint32_t num_elements,
                                           BitVector::Builder& builder) const {
  std::optional<NumericValue> val = GetNumericTypeVariant(type_, sql_val);

  // If the value is invalid we should just ignore those elements.
  if (!val.has_value() || op == FilterOp::kIsNotNull ||
      op == FilterOp::kIsNull || op == FilterOp::kGlob) {
    builder.Skip(num_elements);
    return;
  }

  std::visit(
      [this, op, offset, num_elements, &builder](auto val) {
        using T = decltype(val);
        auto* typed_start = static_cast<const T*>(data_) + offset;
        TypedSlowPathComparison(val, op, typed_start, num_elements, builder);
      },
      *val);
}

uint32_t NumericStorage::UpperBoundIndex(NumericValue val,
                                         Range search_range) const {
  return std::visit(
      [this, search_range](auto val_data) {
        using T = decltype(val_data);
        const T* typed_start = static_cast<const T*>(data_);
        auto upper = std::upper_bound(typed_start + search_range.start,
                                      typed_start + search_range.end, val_data);
        return static_cast<uint32_t>(std::distance(typed_start, upper));
      },
      val);
}

// As we don't template those functions, we need to use std::visitor to type
// `start`, hence this wrapping.
uint32_t NumericStorage::LowerBoundIndex(NumericValue val,
                                         Range search_range) const {
  return std::visit(
      [this, search_range](auto val_data) {
        using T = decltype(val_data);
        const T* typed_start = static_cast<const T*>(data_);
        auto lower = std::lower_bound(typed_start + search_range.start,
                                      typed_start + search_range.end, val_data);
        return static_cast<uint32_t>(std::distance(typed_start, lower));
      },
      val);
}

std::optional<Range> NumericStorage::BinarySearch(FilterOp op,
                                                  SqlValue sql_val,
                                                  Range search_range) const {
  std::optional<NumericValue> val = GetNumericTypeVariant(type_, sql_val);
  if (op == FilterOp::kIsNotNull)
    return Range(0, size());

  if (!val)
    return std::nullopt;

  switch (op) {
    case FilterOp::kEq:
      return Range(LowerBoundIndex(*val, search_range),
                   UpperBoundIndex(*val, search_range));
    case FilterOp::kLe:
      return Range(0, UpperBoundIndex(*val, search_range));
    case FilterOp::kLt:
      return Range(0, LowerBoundIndex(*val, search_range));
    case FilterOp::kGe:
      return Range(LowerBoundIndex(*val, search_range), size_);
    case FilterOp::kGt:
      return Range(UpperBoundIndex(*val, search_range), size_);
    case FilterOp::kNe:
    case FilterOp::kIsNull:
    case FilterOp::kIsNotNull:
    case FilterOp::kGlob:
      return std::nullopt;
  }
  return std::nullopt;
}

uint32_t NumericStorage::UpperBoundIndex(NumericValue val,
                                         uint32_t* order,
                                         Range search_range) const {
  return std::visit(
      [this, order, search_range](auto val_data) {
        using T = decltype(val_data);
        const T* typed_start = static_cast<const T*>(data_);
        auto upper = std::upper_bound(order + search_range.start,
                                      order + search_range.end, val_data,
                                      [typed_start](T val, uint32_t index) {
                                        return val < *(typed_start + index);
                                      });
        return static_cast<uint32_t>(std::distance(order, upper));
      },
      val);
}

// As we don't template those functions, we need to use std::visitor to type
// `start`, hence this wrapping.
uint32_t NumericStorage::LowerBoundIndex(NumericValue val,
                                         uint32_t* order,
                                         Range search_range) const {
  return std::visit(
      [this, order, search_range](auto val_data) {
        using T = decltype(val_data);
        const T* typed_start = static_cast<const T*>(data_);
        auto lower = std::lower_bound(order + search_range.start,
                                      order + search_range.end, val_data,
                                      [typed_start](uint32_t index, T val) {
                                        return *(typed_start + index) < val;
                                      });
        return static_cast<uint32_t>(std::distance(order, lower));
      },
      val);
}

std::optional<Range> NumericStorage::BinarySearchWithIndex(
    FilterOp op,
    SqlValue sql_val,
    uint32_t* order,
    Range search_range) const {
  std::optional<NumericValue> val = GetNumericTypeVariant(type_, sql_val);

  if (op == FilterOp::kIsNotNull)
    return Range(0, size());

  if (!val.has_value())
    return std::nullopt;

  switch (op) {
    case FilterOp::kEq:
      return Range(LowerBoundIndex(*val, order, search_range),
                   UpperBoundIndex(*val, order, search_range));
    case FilterOp::kLe:
      return Range(0, UpperBoundIndex(*val, order, search_range));
    case FilterOp::kLt:
      return Range(0, LowerBoundIndex(*val, order, search_range));
    case FilterOp::kGe:
      return Range(LowerBoundIndex(*val, order, search_range), size_);
    case FilterOp::kGt:
      return Range(UpperBoundIndex(*val, order, search_range), size_);
    case FilterOp::kNe:
    case FilterOp::kIsNull:
    case FilterOp::kIsNotNull:
    case FilterOp::kGlob:
      return std::nullopt;
  }
  return std::nullopt;
}

}  // namespace storage
}  // namespace trace_processor
}  // namespace perfetto