tp: let INTERVAL INTERSECTION's PER keys be of any one type

Android.bp

@@ -17823,6 +17823,7 @@ filegroup {
         "src/trace_processor/core/exec/column_view.cc",
         "src/trace_processor/core/exec/dataframe_scan.cc",
         "src/trace_processor/core/exec/interval_intersect.cc",
+        "src/trace_processor/core/exec/key_encoder.cc",
         "src/trace_processor/core/exec/operator.cc",
         "src/trace_processor/core/exec/pipeline.cc",
         "src/trace_processor/core/exec/row_batch.cc",
@@ -17848,6 +17849,7 @@ filegroup {
         "src/trace_processor/core/exec/dataframe_scan_unittest.cc",
         "src/trace_processor/core/exec/executor_contract_unittest.cc",
         "src/trace_processor/core/exec/interval_intersect_unittest.cc",
+        "src/trace_processor/core/exec/key_encoder_unittest.cc",
         "src/trace_processor/core/exec/operator_unittest.cc",
         "src/trace_processor/core/exec/row_batch_unittest.cc",
         "src/trace_processor/core/exec/row_store_unittest.cc",

BUILD

@@ -2587,6 +2587,8 @@ perfetto_filegroup(
         "src/trace_processor/core/exec/dataframe_scan.h",
         "src/trace_processor/core/exec/interval_intersect.cc",
         "src/trace_processor/core/exec/interval_intersect.h",
+        "src/trace_processor/core/exec/key_encoder.cc",
+        "src/trace_processor/core/exec/key_encoder.h",
         "src/trace_processor/core/exec/operator.cc",
         "src/trace_processor/core/exec/operator.h",
         "src/trace_processor/core/exec/pipeline.cc",

include/perfetto/base/compiler.h

@@ -61,6 +61,17 @@
 #define PERFETTO_NORETURN __declspec(noreturn)
 #endif
 
+// Promises that what a pointer reaches is reached through nothing else in its
+// scope, so writes through it cannot change what other pointers read. Not
+// `__restrict`: the macOS SDK defines that away to nothing in C++.
+#if defined(__GNUC__) || defined(__clang__)
+#define PERFETTO_RESTRICT __restrict__
+#elif defined(_MSC_VER)
+#define PERFETTO_RESTRICT __restrict
+#else
+#define PERFETTO_RESTRICT
+#endif
+
 #if defined(__GNUC__) || defined(__clang__)
 #define PERFETTO_DEBUG_FUNCTION_IDENTIFIER() __PRETTY_FUNCTION__
 #elif defined(_MSC_VER)

src/trace_processor/core/common/row_layout.h

@@ -19,6 +19,7 @@
 
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <type_traits>
 #include <vector>
 
@@ -39,21 +40,26 @@ class RowLayout {
     bool nullable = false;
     bool descending = false;
   };
+  // Small enough to pass in one register: Write() takes it by value, so its
+  // fields stay in registers while row bytes are written.
   struct Slot {
-    uint32_t offset;
-    uint32_t stride;
+    uint16_t offset;
+    uint16_t stride;
     bool nullable;
     bool descending;
   };
+  static_assert(sizeof(Slot) <= 8);
 
   RowLayout() = default;
   explicit RowLayout(const std::vector<Column>& columns) {
     for (const Column& column : columns) {
-      slots_.push_back({stride_, 0, column.nullable, column.descending});
+      slots_.push_back({static_cast<uint16_t>(stride_), 0, column.nullable,
+                        column.descending});
       stride_ += (column.nullable ? 1u : 0u) + ValueSize(column.type);
     }
+    PERFETTO_CHECK(stride_ <= std::numeric_limits<uint16_t>::max());
     for (Slot& slot : slots_) {
-      slot.stride = stride_;
+      slot.stride = static_cast<uint16_t>(stride_);
     }
   }
 
@@ -63,10 +69,10 @@ class RowLayout {
 
   // `get(i, &value)` returns false if row i has no value.
   template <typename T, typename Get>
-  PERFETTO_ALWAYS_INLINE static void Write(const Slot& slot,
+  PERFETTO_ALWAYS_INLINE static void Write(Slot slot,
                                            uint32_t count,
                                            Get get,
-                                           uint8_t* rows) {
+                                           uint8_t* PERFETTO_RESTRICT rows) {
     // Any other type would be implicitly converted to one of these by
     // Encode, writing a value of a different size to its slot.
     static_assert(std::is_same_v<T, uint32_t> || std::is_same_v<T, int32_t> ||

src/trace_processor/core/exec/BUILD.gn

@@ -32,6 +32,8 @@ source_set("exec") {
     "dataframe_scan.h",
     "interval_intersect.cc",
     "interval_intersect.h",
+    "key_encoder.cc",
+    "key_encoder.h",
     "operator.cc",
     "operator.h",
     "pipeline.cc",
@@ -80,6 +82,7 @@ perfetto_unittest_source_set("unittests") {
     "dataframe_scan_unittest.cc",
     "executor_contract_unittest.cc",
     "interval_intersect_unittest.cc",
+    "key_encoder_unittest.cc",
     "operator_unittest.cc",
     "row_batch_unittest.cc",
     "row_store_unittest.cc",
@@ -105,6 +108,7 @@ if (enable_perfetto_benchmarks) {
   source_set("benchmarks") {
     testonly = true
     sources = [
+      "key_encoder_benchmark.cc",
       "tree_number_nodes_benchmark.cc",
       "tree_order_benchmark.cc",
     ]
@@ -115,6 +119,7 @@ if (enable_perfetto_benchmarks) {
       "../../containers",
       "../common",
       "../dataframe",
+      "../util",
     ]
   }
 }

src/trace_processor/core/exec/column_view.h

@@ -131,6 +131,34 @@ class ColumnView {
   const BitVector* validity_ = nullptr;
 };
 
+// Reads a flat column's values through its selection and validity.
+template <typename T>
+class FlatColumnReader {
+ public:
+  explicit FlatColumnReader(const ColumnView& column)
+      : data_(static_cast<const T*>(column.data())),
+        selection_(column.selection()),
+        validity_(column.validity()) {
+    PERFETTO_DCHECK(column.kind() == ColumnView::Kind::kFlat);
+    PERFETTO_DCHECK(column.type().Is<typename TypeTagFor<T>::type>());
+  }
+
+  // False if the row holds no value.
+  PERFETTO_ALWAYS_INLINE bool Read(uint32_t row, T* out) const {
+    uint32_t index = selection_.GetIndex(row);
+    if (validity_ && !validity_->is_set(index)) {
+      return false;
+    }
+    *out = data_[index];
+    return true;
+  }
+
+ private:
+  const T* data_;
+  RowSelection selection_;
+  const BitVector* validity_;
+};
+
 // Whether two batches' views of a column can be combined. An implicit Id and
 // a stored Uint32 hold the same values: gathering turns one into the other.
 inline bool SameLogicalType(const ColumnView& a, const ColumnView& b) {

src/trace_processor/core/exec/interval_intersect.cc

@@ -21,7 +21,9 @@
 #include <cstring>
 #include <memory>
 #include <numeric>
+#include <optional>
 #include <string>
+#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -35,6 +37,7 @@
 #include "src/trace_processor/core/common/storage_types.h"
 #include "src/trace_processor/core/exec/column_chunk.h"
 #include "src/trace_processor/core/exec/column_view.h"
+#include "src/trace_processor/core/exec/key_encoder.h"
 #include "src/trace_processor/core/exec/operator.h"
 #include "src/trace_processor/core/exec/row_batch.h"
 #include "src/trace_processor/core/exec/row_selection.h"
@@ -85,27 +88,6 @@ struct Group {
   bool nonoverlapping = true;
 };
 
-// Reads one flat, non-null Int64 column of a batch.
-struct Int64Reader {
-  explicit Int64Reader(const ColumnView& column)
-      : data(static_cast<const int64_t*>(column.data())),
-        selection(column.selection()),
-        validity(column.validity()) {}
-
-  bool Read(uint32_t row, int64_t* out) const {
-    uint32_t index = selection.GetIndex(row);
-    if (validity && !validity->is_set(index)) {
-      return false;
-    }
-    *out = data[index];
-    return true;
-  }
-
-  const int64_t* data;
-  RowSelection selection;
-  const BitVector* validity;
-};
-
 base::Status ValidateOperand(const RowBatch& batch,
                              const IntervalIntersectOperand& operand,
                              uint32_t which) {
@@ -114,39 +96,25 @@ base::Status ValidateOperand(const RowBatch& batch,
     return column.kind() == ColumnView::Kind::kFlat &&
            column.type().Is<Int64>();
   };
-  bool ok = is_int64(operand.ts_column) && is_int64(operand.dur_column) &&
-            std::all_of(operand.key_columns.begin(), operand.key_columns.end(),
-                        is_int64);
+  bool ok = is_int64(operand.ts_column) && is_int64(operand.dur_column);
   return ok ? base::OkStatus()
             : base::ErrStatus(
-                  "INTERVAL INTERSECTION: operand %u's ts, dur and PER "
-                  "columns must be Int64",
+                  "INTERVAL INTERSECTION: operand %u's ts and dur columns "
+                  "must be Int64",
                   which + 1);
 }
 
-// A key column takes this many bytes: whether it holds a value, then the
-// value itself. Two rows which hold no value there agree on it, as they
-// would under GROUP BY.
-constexpr size_t kKeyColumnBytes = 1 + sizeof(int64_t);
-
-void WriteKeyColumn(std::string& key,
-                    uint32_t at,
-                    bool present,
-                    int64_t value) {
-  char* to = key.data() + at * kKeyColumnBytes;
-  to[0] = present ? 1 : 0;
-  memcpy(to + 1, &value, sizeof(value));
-}
-
 class IntersectState : public OperatorState {
  public:
   ~IntersectState() override;
 
   std::vector<std::unique_ptr<OperatorState>> operand_states;
   std::vector<std::unique_ptr<RowStore>> stores;
+  // Shared by the operands, so their keys' types must agree.
+  KeyEncoder keys;
   // Per operand, its rows by key. A key with no entry in some operand covers
   // nothing, so only keys every operand has produce regions.
-  std::vector<base::FlatHashMap<std::string, Group>> groups;
+  std::vector<base::FlatHashMapV2<std::string, Group>> groups;
 
   bool computed = false;
   Regions regions;
@@ -181,17 +149,19 @@ base::Status Collect(const IntervalIntersectOperand& operand,
                      uint32_t which,
                      OperatorState& state,
                      RowStore& store,
-                     base::FlatHashMap<std::string, Group>& groups) {
+                     KeyEncoder& keys,
+                     base::FlatHashMapV2<std::string, Group>& groups) {
   RowBatch batch;
   RowBatch retained;
-  std::string key(operand.key_columns.size() * kKeyColumnBytes, '\0');
   while (operand.source->GetData(batch, state)) {
     RETURN_IF_ERROR(ValidateOperand(batch, operand, which));
-    Int64Reader ts(batch.column(operand.ts_column));
-    Int64Reader dur(batch.column(operand.dur_column));
-    std::vector<Int64Reader> keys;
-    for (uint32_t column : operand.key_columns) {
-      keys.emplace_back(batch.column(column));
+    FlatColumnReader<int64_t> ts(batch.column(operand.ts_column));
+    FlatColumnReader<int64_t> dur(batch.column(operand.dur_column));
+    if (std::optional<uint32_t> bad = keys.Encode(batch, operand.key_columns)) {
+      return base::ErrStatus(
+          "INTERVAL INTERSECTION: operand %u's PER column %u must hold one "
+          "type, the same in every operand",
+          which + 1, *bad + 1);
     }
     for (uint32_t row = 0; row < batch.size(); ++row) {
       int64_t start;
@@ -205,17 +175,14 @@ base::Status Collect(const IntervalIntersectOperand& operand,
             "below zero",
             which + 1);
       }
-      for (uint32_t i = 0; i < keys.size(); ++i) {
-        // Read before writing: argument evaluation order is unspecified, so
-        // passing both `Read(&value)` and `value` could copy it unread.
-        int64_t value = 0;
-        bool present = keys[i].Read(row, &value);
-        WriteKeyColumn(key, i, present, value);
+      std::string_view key = keys.Key(row);
+      Group* group = groups.Find(key);
+      if (!group) {
+        group = groups.Insert(std::string(key), Group{}).first;
       }
-      Group& group = groups[key];
-      group.intervals.push_back({static_cast<Ts>(start),
-                                 static_cast<Ts>(start + length),
-                                 store.size() + row});
+      group->intervals.push_back({static_cast<Ts>(start),
+                                  static_cast<Ts>(start + length),
+                                  store.size() + row});
     }
     retained.Reset();
     for (uint32_t column : operand.retained_columns) {
@@ -356,7 +323,7 @@ bool IntervalIntersect::GetData(RowBatch& out, OperatorState& state) const {
     s.scratch.narrowed.operands = count;
     for (uint32_t i = 0; i < count; ++i) {
       s.status = Collect(operands_[i], i, *s.operand_states[i], *s.stores[i],
-                         s.groups[i]);
+                         s.keys, s.groups[i]);
       if (!s.status.ok()) {
         return false;
       }

src/trace_processor/core/exec/interval_intersect.h

@@ -32,7 +32,7 @@ namespace perfetto::trace_processor::core::exec {
 struct IntervalIntersectOperand {
   // Read in full before any region is found.
   const Source* source = nullptr;
-  // All must be flat Int64 columns of `source`'s batches.
+  // Flat Int64 columns of `source`'s batches.
   uint32_t ts_column = 0;
   uint32_t dur_column = 0;
   // Compared pairwise across the operands, in this order.