tp: keep a pipeline's rows when SQLite reads it again

Android.bp

@@ -17811,6 +17811,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/memoize.cc",
         "src/trace_processor/core/exec/operator.cc",
         "src/trace_processor/core/exec/pipeline.cc",
         "src/trace_processor/core/exec/row_batch.cc",
@@ -17836,6 +17837,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/memoize_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

@@ -2584,6 +2584,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/memoize.cc",
+        "src/trace_processor/core/exec/memoize.h",
         "src/trace_processor/core/exec/operator.cc",
         "src/trace_processor/core/exec/operator.h",
         "src/trace_processor/core/exec/pipeline.cc",

src/trace_processor/core/exec/BUILD.gn

@@ -30,6 +30,8 @@ source_set("exec") {
     "dataframe_scan.h",
     "interval_intersect.cc",
     "interval_intersect.h",
+    "memoize.cc",
+    "memoize.h",
     "operator.cc",
     "operator.h",
     "pipeline.cc",
@@ -78,6 +80,7 @@ perfetto_unittest_source_set("unittests") {
     "dataframe_scan_unittest.cc",
     "executor_contract_unittest.cc",
     "interval_intersect_unittest.cc",
+    "memoize_unittest.cc",
     "operator_unittest.cc",
     "row_batch_unittest.cc",
     "row_store_unittest.cc",

src/trace_processor/core/exec/memoize.cc

@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2026 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/core/exec/memoize.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+
+#include "perfetto/base/status.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"
+#include "src/trace_processor/core/exec/row_store.h"
+
+namespace perfetto::trace_processor::core::exec {
+namespace {
+
+struct MemoizeState : OperatorState {
+  ~MemoizeState() override;
+
+  std::unique_ptr<OperatorState> source;
+  RowBatch batch;
+  RowStore kept;
+  // Whether `kept` holds all the source's rows.
+  bool complete = false;
+  // Whether this read is replaying `kept`, and how far it has got.
+  bool replaying = false;
+  uint32_t served = 0;
+  // Whether this read is keeping what it reads.
+  bool keeping = true;
+};
+
+MemoizeState::~MemoizeState() = default;
+
+}  // namespace
+
+Memoize::Memoize(const Source& source) : source_(source) {}
+Memoize::~Memoize() = default;
+
+std::unique_ptr<OperatorState> Memoize::MakeState() const {
+  auto state = std::make_unique<MemoizeState>();
+  state->source = source_.MakeState();
+  return state;
+}
+
+bool Memoize::GetData(RowBatch& out, OperatorState& state) const {
+  auto& s = state.Cast<MemoizeState>();
+  if (s.replaying) {
+    uint32_t left = s.kept.size() - s.served;
+    if (left == 0) {
+      return false;
+    }
+    s.served += s.kept.View(&out, s.served, std::min(left, kMaxBatchRows));
+    return true;
+  }
+  if (!source_.GetData(s.batch, *s.source)) {
+    s.complete = s.keeping && source_.status(*s.source).ok();
+    return false;
+  }
+  // Batches which cannot be kept are still passed on: this read just keeps
+  // nothing.
+  if (s.keeping && !s.kept.Append(s.batch).ok()) {
+    s.keeping = false;
+    s.kept.Clear();
+  }
+  out.CopyFrom(s.batch);
+  return true;
+}
+
+void Memoize::Rewind(OperatorState& state) const {
+  auto& s = state.Cast<MemoizeState>();
+  if (s.complete) {
+    s.replaying = true;
+    s.served = 0;
+    return;
+  }
+  source_.Rewind(*s.source);
+  s.kept.Clear();
+  s.keeping = true;
+}
+
+base::Status Memoize::status(const OperatorState& state) const {
+  const auto& s = state.Cast<const MemoizeState>();
+  return s.replaying ? base::OkStatus() : source_.status(*s.source);
+}
+
+}  // namespace perfetto::trace_processor::core::exec

src/trace_processor/core/exec/memoize.h

@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2026 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.
+ */
+
+#ifndef SRC_TRACE_PROCESSOR_CORE_EXEC_MEMOIZE_H_
+#define SRC_TRACE_PROCESSOR_CORE_EXEC_MEMOIZE_H_
+
+#include <memory>
+
+#include "perfetto/base/status.h"
+#include "src/trace_processor/core/exec/operator.h"
+#include "src/trace_processor/core/exec/row_batch.h"
+
+namespace perfetto::trace_processor::core::exec {
+
+// Passes on the batches of another source, keeping them as it goes. Once a
+// read has reached the end, every read after replays what was kept rather
+// than reading the source again.
+//
+// A read stopped part way keeps nothing, and neither does one whose batches
+// cannot be kept: the next read goes back to the source.
+class Memoize : public Source {
+ public:
+  // `source` must outlive this.
+  explicit Memoize(const Source& source);
+  Memoize(Source&&) = delete;
+  ~Memoize() override;
+
+  std::unique_ptr<OperatorState> MakeState() const override;
+  bool GetData(RowBatch& out, OperatorState& state) const override;
+  void Rewind(OperatorState& state) const override;
+  base::Status status(const OperatorState& state) const override;
+
+ private:
+  const Source& source_;
+};
+
+}  // namespace perfetto::trace_processor::core::exec
+
+#endif  // SRC_TRACE_PROCESSOR_CORE_EXEC_MEMOIZE_H_

src/trace_processor/core/exec/memoize_unittest.cc

@@ -0,0 +1,106 @@
+/*
+ * Copyright (C) 2026 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/core/exec/memoize.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "src/trace_processor/core/exec/operator.h"
+#include "src/trace_processor/core/exec/row_batch.h"
+#include "src/trace_processor/core/exec/row_cursor.h"
+#include "src/trace_processor/core/exec/test_utils.h"
+#include "test/gtest_and_gmock.h"
+
+namespace perfetto::trace_processor::core::exec {
+namespace {
+
+using test::ArraySource;
+using test::FailingSource;
+using test::Sequence;
+
+// Counts how many reads of `source` were started.
+class CountingSource final : public Source {
+ public:
+  explicit CountingSource(const Source& source) : source_(source) {}
+
+  std::unique_ptr<OperatorState> MakeState() const override {
+    return source_.MakeState();
+  }
+  void Rewind(OperatorState& state) const override {
+    ++reads;
+    source_.Rewind(state);
+  }
+  bool GetData(RowBatch& out, OperatorState& state) const override {
+    return source_.GetData(out, state);
+  }
+  base::Status status(const OperatorState& state) const override {
+    return source_.status(state);
+  }
+
+  mutable uint32_t reads = 0;
+
+ private:
+  const Source& source_;
+};
+
+// The values of column 1 from `cursor`, from its first row.
+std::vector<int64_t> ReadAll(RowCursor& cursor) {
+  std::vector<int64_t> values;
+  for (bool row = cursor.Open(); row; row = cursor.Next()) {
+    values.push_back(cursor.Value<int64_t>(1));
+  }
+  return values;
+}
+
+TEST(MemoizeTest, ReadsAfterTheFirstFullReadAreReplayed) {
+  ArraySource array(Sequence(5000));
+  CountingSource source(array);
+  Memoize memoize(source);
+  RowCursor cursor(memoize);
+  EXPECT_EQ(ReadAll(cursor), Sequence(5000));
+  EXPECT_EQ(ReadAll(cursor), Sequence(5000));
+  EXPECT_EQ(ReadAll(cursor), Sequence(5000));
+  EXPECT_TRUE(cursor.status().ok());
+  EXPECT_EQ(source.reads, 1u);
+}
+
+TEST(MemoizeTest, APartReadKeepsNothing) {
+  ArraySource array(Sequence(5000));
+  CountingSource source(array);
+  Memoize memoize(source);
+  RowCursor cursor(memoize);
+  ASSERT_TRUE(cursor.Open());
+  EXPECT_EQ(ReadAll(cursor), Sequence(5000));
+  EXPECT_EQ(ReadAll(cursor), Sequence(5000));
+  EXPECT_EQ(source.reads, 2u);
+}
+
+TEST(MemoizeTest, AFailedReadKeepsNothing) {
+  FailingSource failing;
+  CountingSource source(failing);
+  Memoize memoize(source);
+  RowCursor cursor(memoize);
+  ReadAll(cursor);
+  EXPECT_FALSE(cursor.status().ok());
+  ReadAll(cursor);
+  EXPECT_FALSE(cursor.status().ok());
+  EXPECT_EQ(source.reads, 2u);
+}
+
+}  // namespace
+}  // namespace perfetto::trace_processor::core::exec

src/trace_processor/perfetto_sql/engine/perfetto_sql_connection_unittest.cc

@@ -1239,6 +1239,22 @@ TEST_F(PerfettoSqlConnectionPipelineTest, PipelinesAsIntersectionOperands) {
   EXPECT_THAT(*rows, testing::ElementsAre("12,3", "20,2"));
 }
 
+// SQLite reads the inner side of a join once per outer row. A pipeline read
+// again with the same plan is not run again each time: `random()` in its
+// source would otherwise give each read different rows.
+TEST_F(PerfettoSqlConnectionPipelineTest, PipelinesReadAgainAreNotRunAgain) {
+  ASSERT_TRUE(Rows("CREATE TABLE picks(k); "
+                   "INSERT INTO picks VALUES (1), (2), (3), (4), (5)")
+                  .ok());
+  auto rows = Rows(
+      "SELECT count(DISTINCT p.r) FROM picks CROSS JOIN "
+      "(FROM (SELECT 1 AS one, random() AS r)) p");
+  ASSERT_TRUE(rows.ok()) << rows.status().message();
+  // The second read runs the pipeline once more to keep its rows; every read
+  // after replays them.
+  EXPECT_THAT(*rows, testing::ElementsAre("2"));
+}
+
 TEST_F(PerfettoSqlConnectionPipelineTest, PipelineSubqueriesNeedPipelines) {
   ASSERT_TRUE(Rows("PERFETTO PRAGMA pipelines = 0").ok());
   EXPECT_THAT(Rows("SELECT * FROM (FROM tree)").status().message(),

src/trace_processor/perfetto_sql/engine/pipeline_module.cc

@@ -34,6 +34,7 @@
 #include "src/trace_processor/containers/string_pool.h"
 #include "src/trace_processor/core/common/storage_types.h"
 #include "src/trace_processor/core/exec/column_view.h"
+#include "src/trace_processor/core/exec/memoize.h"
 #include "src/trace_processor/core/exec/row_cursor.h"
 #include "src/trace_processor/core/exec/variant.h"
 #include "src/trace_processor/core/util/bit_vector.h"
@@ -159,6 +160,14 @@ void CacheColumnReaders(PipelineModule::Cursor* c) {
   }
 }
 
+// Drops the plan the cursor is running and everything reading it.
+void ResetPlan(PipelineModule::Cursor* c) {
+  c->rows.reset();
+  c->memoize.reset();
+  c->plan.reset();
+  c->serialized.clear();
+}
+
 // Surfaces the executor's error, if any, once rows stop.
 int CheckStatus(PipelineModule::Cursor* cursor) {
   // status() walks every node so only check it once rows stop.
@@ -245,20 +254,26 @@ int PipelineModule::Filter(sqlite3_vtab_cursor* cursor,
   std::string_view serialized(
       static_cast<const char*>(sqlite3_value_blob(argv[0])),
       static_cast<size_t>(sqlite3_value_bytes(argv[0])));
-  // A cursor filtered again, as the inner side of a join is, runs the same
-  // pipeline each time.
   if (!c->plan || c->serialized != serialized) {
     Context* context = GetVtab(cursor->pVtab)->context;
     auto plan = context->connection->LoadPipeline(serialized);
     if (!plan.ok()) {
       return sqlite::utils::SetError(cursor->pVtab, plan.status());
     }
-    c->rows.reset();
+    ResetPlan(c);
     c->plan = std::move(*plan);
     c->serialized = serialized;
     c->pool = context->pool;
   }
-  c->rows = std::make_unique<core::exec::RowCursor>(c->plan->source());
+  // SQLite reads the inner side of a join again for each outer row. The
+  // second read of a plan keeps what it produces, and later reads replay it.
+  if (!c->rows) {
+    c->rows = std::make_unique<core::exec::RowCursor>(c->plan->source());
+  } else if (!c->memoize) {
+    c->rows.reset();
+    c->memoize = std::make_unique<core::exec::Memoize>(c->plan->source());
+    c->rows = std::make_unique<core::exec::RowCursor>(*c->memoize);
+  }
   c->rowid = 0;
   c->target_rowid.reset();
   if (idx_num && sqlite3_value_type(argv[1]) == SQLITE_INTEGER) {

src/trace_processor/perfetto_sql/engine/pipeline_module.h

@@ -27,6 +27,7 @@
 
 #include "perfetto/ext/base/status_or.h"
 #include "src/trace_processor/containers/string_pool.h"
+#include "src/trace_processor/core/exec/memoize.h"
 #include "src/trace_processor/core/exec/row_cursor.h"
 #include "src/trace_processor/perfetto_sql/pipeline/logical_plan.h"
 #include "src/trace_processor/perfetto_sql/pipeline/physical_plan.h"
@@ -71,6 +72,8 @@ struct PipelineModule : sqlite::Module<PipelineModule> {
     // The serialized plan `plan` was loaded from.
     std::string serialized;
     std::unique_ptr<pipeline::PhysicalPlan> plan;
+    // Keeps what `plan` produced once the cursor is read again.
+    std::unique_ptr<core::exec::Memoize> memoize;
     StringPool* pool = nullptr;
     std::unique_ptr<core::exec::RowCursor> rows;
     std::vector<ColumnReader> columns;