trace_processor: add special handling for CPU in window operator

This allows for fast queries from the UI when span joining with other
tables.

Change-Id: I4e49a5f0fb5fc379e88972edc84c4f4a2f04fb25
diff --git a/src/trace_processor/window_operator_table.cc b/src/trace_processor/window_operator_table.cc
index e10b14d..ae85dc4 100644
--- a/src/trace_processor/window_operator_table.cc
+++ b/src/trace_processor/window_operator_table.cc
@@ -133,24 +133,44 @@
   current_cpu_ = 0;
   quantum_ts_ = 0;
   row_id_ = 0;
-  return_first = qc.constraints().size() == 1 &&
-                 qc.constraints()[0].iColumn == Column::kRowId &&
-                 sqlite3_value_int(v[0]) == 0;
+
+  // Set return first if there is a equals constraint on the row id asking to
+  // return the first row.
+  bool return_first = qc.constraints().size() == 1 &&
+                      qc.constraints()[0].iColumn == Column::kRowId &&
+                      IsOpEq(qc.constraints()[0].op) &&
+                      sqlite3_value_int(v[0]) == 0;
+  // Set return CPU if there is an equals constraint on the CPU column.
+  bool return_cpu = qc.constraints().size() == 1 &&
+                    qc.constraints()[0].iColumn == Column::kCpu &&
+                    IsOpEq(qc.constraints()[0].op);
+  if (return_first) {
+    filter_type_ = FilterType::kReturnFirst;
+  } else if (return_cpu) {
+    filter_type_ = FilterType::kReturnCpu;
+    current_cpu_ = static_cast<uint32_t>(sqlite3_value_int(v[0]));
+  } else {
+    filter_type_ = FilterType::kReturnAll;
+  }
   return SQLITE_OK;
 }
 
 int WindowOperatorTable::Cursor::Next() {
-  // If we're only returning the first row, set the values to EOF.
-  if (return_first) {
-    current_cpu_ = base::kMaxCpus;
-    current_ts_ = window_end_;
-    return SQLITE_OK;
-  }
-
-  if (++current_cpu_ == base::kMaxCpus && current_ts_ < window_end_) {
-    current_cpu_ = 0;
-    current_ts_ += step_size_;
-    quantum_ts_++;
+  switch (filter_type_) {
+    case FilterType::kReturnFirst:
+      current_ts_ = window_end_;
+      break;
+    case FilterType::kReturnCpu:
+      current_ts_ += step_size_;
+      quantum_ts_++;
+      break;
+    case FilterType::kReturnAll:
+      if (++current_cpu_ == base::kMaxCpus && current_ts_ < window_end_) {
+        current_cpu_ = 0;
+        current_ts_ += step_size_;
+        quantum_ts_++;
+      }
+      break;
   }
   row_id_++;
   return SQLITE_OK;
diff --git a/src/trace_processor/window_operator_table.h b/src/trace_processor/window_operator_table.h
index db1b165..8b2ac93 100644
--- a/src/trace_processor/window_operator_table.h
+++ b/src/trace_processor/window_operator_table.h
@@ -65,6 +65,16 @@
     int Column(sqlite3_context*, int N) override;
 
    private:
+    // Defines the data to be generated by the table.
+    enum FilterType {
+      // Returns all the spans for all CPUs.
+      kReturnAll = 0,
+      // Only returns the first span of the table. Useful for UPDATE operations.
+      kReturnFirst = 1,
+      // Only returns all the spans for a chosen CPU.
+      kReturnCpu = 2,
+    };
+
     uint64_t const window_start_;
     uint64_t const window_end_;
     uint64_t const step_size_;
@@ -74,7 +84,8 @@
     uint32_t current_cpu_ = 0;
     uint64_t quantum_ts_ = 0;
     uint64_t row_id_ = 0;
-    bool return_first = false;
+
+    FilterType filter_type_ = FilterType::kReturnAll;
   };
 
   uint64_t quantum_ = 0;