tp: fold tree values in a streaming pipeline Once rows are in tree order, a fold only needs one running value per node. TreeAccumulateUp consumes child-first rows and adds each completed child into its parent. TreeAccumulateDown consumes parent-first rows and carries each parent total into its children. Both operators append the accumulated value and stream batches without keeping another copy of the input. The SqlScan test covers the complete path from SQLite variants through type validation, node numbering, ordering, and accumulation.
diff --git a/Android.bp b/Android.bp index a8255df..271af3a 100644 --- a/Android.bp +++ b/Android.bp
@@ -17710,6 +17710,7 @@ "src/trace_processor/core/exec/row_batch.cc", "src/trace_processor/core/exec/row_cursor.cc", "src/trace_processor/core/exec/row_store.cc", + "src/trace_processor/core/exec/tree_accumulate.cc", "src/trace_processor/core/exec/tree_number_nodes.cc", "src/trace_processor/core/exec/tree_order.cc", ], @@ -17728,6 +17729,7 @@ "src/trace_processor/core/exec/dataframe_scan_unittest.cc", "src/trace_processor/core/exec/operator_unittest.cc", "src/trace_processor/core/exec/row_store_unittest.cc", + "src/trace_processor/core/exec/tree_accumulate_unittest.cc", "src/trace_processor/core/exec/tree_number_nodes_unittest.cc", "src/trace_processor/core/exec/tree_order_unittest.cc", ],
diff --git a/src/trace_processor/core/exec/BUILD.gn b/src/trace_processor/core/exec/BUILD.gn index e381a33..8f12368 100644 --- a/src/trace_processor/core/exec/BUILD.gn +++ b/src/trace_processor/core/exec/BUILD.gn
@@ -33,6 +33,8 @@ "row_selection.h", "row_store.cc", "row_store.h", + "tree_accumulate.cc", + "tree_accumulate.h", "tree_number_nodes.cc", "tree_number_nodes.h", "tree_order.cc", @@ -62,6 +64,7 @@ "dataframe_scan_unittest.cc", "operator_unittest.cc", "row_store_unittest.cc", + "tree_accumulate_unittest.cc", "tree_number_nodes_unittest.cc", "tree_order_unittest.cc", ]
diff --git a/src/trace_processor/core/exec/tree_accumulate.cc b/src/trace_processor/core/exec/tree_accumulate.cc new file mode 100644 index 0000000..737d716 --- /dev/null +++ b/src/trace_processor/core/exec/tree_accumulate.cc
@@ -0,0 +1,228 @@ +/* + * 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/tree_accumulate.h" + +#include <cstdint> +#include <memory> +#include <vector> + +#include "perfetto/base/status.h" +#include "perfetto/ext/base/utils.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/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/tree_number_nodes.h" +#include "src/trace_processor/core/util/flex_vector.h" + +namespace perfetto::trace_processor::core::exec { +namespace { + +// The column's values laid out flat, gathering once if it has an index +// selection. +template <typename T> +const T* Flatten(const ColumnView& column, + uint32_t count, + std::vector<T>* scratch) { + const auto* data = static_cast<const T*>(column.data()); + RowSelection selection = column.selection(); + if (selection.is_range()) { + return data + selection.offset(); + } + scratch->resize(count); + const uint32_t* rows = selection.data(); + T* out = scratch->data(); + for (uint32_t i = 0; i < count; ++i) { + out[i] = data[rows[i]]; + } + return out; +} + +const int64_t* FlattenValues(const ColumnView& column, + uint32_t count, + std::vector<int64_t>* scratch) { + const auto* data = static_cast<const int64_t*>(column.data()); + RowSelection selection = column.selection(); + const BitVector* validity = column.validity(); + if (selection.is_range() && !validity) { + return data + selection.offset(); + } + scratch->resize(count); + for (uint32_t row = 0; row < count; ++row) { + uint32_t index = selection.GetIndex(row); + (*scratch)[row] = validity && !validity->is_set(index) ? 0 : data[index]; + } + return scratch->data(); +} + +base::Status Validate(const RowBatch& in, AccumulateSpec spec) { + const ColumnView& node = in.column(spec.node_column); + const ColumnView& parent = in.column(spec.parent_column); + const ColumnView& value = in.column(spec.value_column); + bool nodes_ok = node.kind() == ColumnView::Kind::kFlat && + node.type().Is<Uint32>() && node.validity() == nullptr && + parent.kind() == ColumnView::Kind::kFlat && + parent.type().Is<Uint32>() && parent.validity() == nullptr; + if (!nodes_ok) { + return base::ErrStatus( + "TREE ACCUMULATE: node columns must be non-null Uint32"); + } + if (value.kind() != ColumnView::Kind::kFlat || !value.type().Is<Int64>()) { + return base::ErrStatus("TREE ACCUMULATE: values must be Int64"); + } + return base::OkStatus(); +} + +bool Add(AccumulateState& state, int64_t a, int64_t b, int64_t* out) { + if (base::CheckedAdd(a, b, out)) { + return true; + } + state.status = base::ErrStatus("TREE ACCUMULATE: integer overflow"); + return false; +} + +void Grow(std::vector<int64_t>* by_node, uint32_t node) { + if (by_node->size() <= node) { + by_node->resize(node + 1, 0); + } +} + +// Fills `out` with the input columns plus a column of `totals`. +void Emit(const RowBatch& in, + RowBatch& out, + const std::shared_ptr<FlexVector<int64_t>>& totals) { + out.CopyFrom(in); + ColumnView column = + ColumnView::Reference(StorageType{Int64{}}, totals->data()); + out.AddColumn(column, totals); +} + +} // namespace + +AccumulateState::~AccumulateState() = default; + +TreeAccumulateUp::TreeAccumulateUp(AccumulateSpec spec) : spec_(spec) {} +TreeAccumulateUp::~TreeAccumulateUp() = default; + +TreeAccumulateDown::TreeAccumulateDown(AccumulateSpec spec) : spec_(spec) {} +TreeAccumulateDown::~TreeAccumulateDown() = default; + +std::unique_ptr<OperatorState> TreeAccumulateUp::MakeState() const { + return std::make_unique<AccumulateState>(); +} +std::unique_ptr<OperatorState> TreeAccumulateDown::MakeState() const { + return std::make_unique<AccumulateState>(); +} + +void TreeAccumulateUp::Rewind(OperatorState& state) const { + AccumulateState& s = state.Cast<AccumulateState>(); + s.by_node.clear(); + s.status = base::OkStatus(); +} +void TreeAccumulateDown::Rewind(OperatorState& state) const { + AccumulateState& s = state.Cast<AccumulateState>(); + s.by_node.clear(); + s.status = base::OkStatus(); +} + +base::Status TreeAccumulateUp::status(const OperatorState& state) const { + return state.Cast<const AccumulateState>().status; +} +base::Status TreeAccumulateDown::status(const OperatorState& state) const { + return state.Cast<const AccumulateState>().status; +} + +OpResult TreeAccumulateUp::Execute(const RowBatch& in, + RowBatch& out, + OperatorState& state) const { + AccumulateState& s = state.Cast<AccumulateState>(); + s.status = Validate(in, spec_); + if (!s.status.ok()) { + return OpResult::kError; + } + uint32_t count = in.size(); + const uint32_t* nodes = + Flatten<uint32_t>(in.column(spec_.node_column), count, &s.node_scratch); + const uint32_t* parents = Flatten<uint32_t>(in.column(spec_.parent_column), + count, &s.parent_scratch); + const int64_t* values = + FlattenValues(in.column(spec_.value_column), count, &s.value_scratch); + + s.totals->resize(count); + int64_t* totals = s.totals->data(); + for (uint32_t row = 0; row < count; ++row) { + uint32_t node = nodes[row]; + Grow(&s.by_node, node); + // Every descendant has already been seen and added its value here, so the + // total is final the moment the node arrives. + int64_t total; + if (!Add(s, values[row], s.by_node[node], &total)) { + return OpResult::kError; + } + totals[row] = total; + uint32_t parent = parents[row]; + if (parent != kNoNode) { + Grow(&s.by_node, parent); + if (!Add(s, s.by_node[parent], total, &s.by_node[parent])) { + return OpResult::kError; + } + } + } + Emit(in, out, s.totals); + return OpResult::kNeedMoreInput; +} + +OpResult TreeAccumulateDown::Execute(const RowBatch& in, + RowBatch& out, + OperatorState& state) const { + AccumulateState& s = state.Cast<AccumulateState>(); + s.status = Validate(in, spec_); + if (!s.status.ok()) { + return OpResult::kError; + } + uint32_t count = in.size(); + const uint32_t* nodes = + Flatten<uint32_t>(in.column(spec_.node_column), count, &s.node_scratch); + const uint32_t* parents = Flatten<uint32_t>(in.column(spec_.parent_column), + count, &s.parent_scratch); + const int64_t* values = + FlattenValues(in.column(spec_.value_column), count, &s.value_scratch); + + s.totals->resize(count); + int64_t* totals = s.totals->data(); + for (uint32_t row = 0; row < count; ++row) { + uint32_t parent = parents[row]; + int64_t above = 0; + if (parent != kNoNode) { + Grow(&s.by_node, parent); + above = s.by_node[parent]; + } + int64_t total; + if (!Add(s, values[row], above, &total)) { + return OpResult::kError; + } + uint32_t node = nodes[row]; + Grow(&s.by_node, node); + s.by_node[node] = total; + totals[row] = total; + } + Emit(in, out, s.totals); + return OpResult::kNeedMoreInput; +} + +} // namespace perfetto::trace_processor::core::exec
diff --git a/src/trace_processor/core/exec/tree_accumulate.h b/src/trace_processor/core/exec/tree_accumulate.h new file mode 100644 index 0000000..e193365 --- /dev/null +++ b/src/trace_processor/core/exec/tree_accumulate.h
@@ -0,0 +1,95 @@ +/* + * 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_TREE_ACCUMULATE_H_ +#define SRC_TRACE_PROCESSOR_CORE_EXEC_TREE_ACCUMULATE_H_ + +#include <cstdint> +#include <memory> +#include <vector> + +#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/util/flex_vector.h" + +namespace perfetto::trace_processor::core::exec { + +// The columns holding the tree structure and the values being summed. Node and +// parent columns must be flat, non-null Uint32 columns; values must be flat +// Int64. A null value contributes zero. +struct AccumulateSpec { + uint32_t node_column = 0; + uint32_t parent_column = 1; + uint32_t value_column = 2; +}; + +// What one execution carries between batches: a running total per node, and +// the totals computed for the current batch. +class AccumulateState : public OperatorState { + public: + ~AccumulateState() override; + + std::vector<int64_t> by_node; + std::vector<uint32_t> node_scratch; + std::vector<uint32_t> parent_scratch; + std::vector<int64_t> value_scratch; + std::shared_ptr<FlexVector<int64_t>> totals = + std::make_shared<FlexVector<int64_t>>(); + base::Status status = base::OkStatus(); +}; + +// Sums each node's value with the values of everything below it. +// +// Requires rows child first, so every descendant has been seen when the node +// arrives. Input columns are preserved and one flat Int64 total column is +// appended. +class TreeAccumulateUp : public Operator { + public: + explicit TreeAccumulateUp(AccumulateSpec); + ~TreeAccumulateUp() override; + + std::unique_ptr<OperatorState> MakeState() const override; + OpResult Execute(const RowBatch&, RowBatch&, OperatorState&) const override; + void Rewind(OperatorState&) const override; + base::Status status(const OperatorState&) const override; + + private: + AccumulateSpec spec_; +}; + +// Sums each node's value with the values of everything above it. +// +// Requires rows parent first, so every ancestor has been totalled when the node +// arrives. Input columns are preserved and one flat Int64 total column is +// appended. +class TreeAccumulateDown : public Operator { + public: + explicit TreeAccumulateDown(AccumulateSpec); + ~TreeAccumulateDown() override; + + std::unique_ptr<OperatorState> MakeState() const override; + OpResult Execute(const RowBatch&, RowBatch&, OperatorState&) const override; + void Rewind(OperatorState&) const override; + base::Status status(const OperatorState&) const override; + + private: + AccumulateSpec spec_; +}; + +} // namespace perfetto::trace_processor::core::exec + +#endif // SRC_TRACE_PROCESSOR_CORE_EXEC_TREE_ACCUMULATE_H_
diff --git a/src/trace_processor/core/exec/tree_accumulate_unittest.cc b/src/trace_processor/core/exec/tree_accumulate_unittest.cc new file mode 100644 index 0000000..555ba54 --- /dev/null +++ b/src/trace_processor/core/exec/tree_accumulate_unittest.cc
@@ -0,0 +1,430 @@ +/* + * 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/tree_accumulate.h" + +#include <algorithm> +#include <cstdint> +#include <limits> +#include <memory> +#include <optional> +#include <random> +#include <utility> +#include <vector> + +#include "src/trace_processor/core/common/storage_types.h" +#include "src/trace_processor/core/exec/column_view.h" +#include "src/trace_processor/core/exec/operator.h" +#include "src/trace_processor/core/exec/pipeline.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/test_utils.h" +#include "src/trace_processor/core/exec/tree_number_nodes.h" +#include "src/trace_processor/core/exec/tree_order.h" +#include "src/trace_processor/core/util/bit_vector.h" +#include "test/gtest_and_gmock.h" + +namespace perfetto::trace_processor::core::exec { +namespace { + +using ::testing::ElementsAre; + +// Emits id, parent id and a value, in whatever order the rows were given. +class RowSource final : public Source { + public: + // `order` is the order the rows are emitted in, which is not the order they + // are numbered in. + RowSource(std::vector<int64_t> parents, + std::vector<int64_t> values, + uint32_t chunk_rows, + std::vector<uint32_t> order = {}) + : parents_(std::move(parents)), + values_(std::move(values)), + chunk_rows_(chunk_rows), + order_(std::move(order)) { + if (order_.empty()) { + for (uint32_t i = 0; i < parents_.size(); ++i) { + order_.push_back(i); + } + } + } + + std::unique_ptr<OperatorState> MakeState() const override { + return std::make_unique<State>(); + } + void Rewind(OperatorState& state) const override { + state.Cast<State>().offset = 0; + } + + bool GetData(RowBatch& out, OperatorState& state) const override { + State& s = state.Cast<State>(); + auto total = static_cast<uint32_t>(parents_.size()); + if (s.offset == total) { + return false; + } + uint32_t count = std::min(chunk_rows_, total - s.offset); + s.ids.resize(count); + s.parents.resize(count); + s.values.resize(count); + s.validity = BitVector::CreateWithSize(count); + for (uint32_t i = 0; i < count; ++i) { + uint32_t row = order_[s.offset + i]; + s.ids[i] = row; + s.values[i] = values_[row]; + s.parents[i] = parents_[row] < 0 ? 0 : parents_[row]; + if (parents_[row] >= 0) { + s.validity.set(i); + } + } + out.Reset(); + out.AddColumn(ColumnView::Reference(StorageType{Int64{}}, s.ids.data())); + out.AddColumn(ColumnView::Reference(StorageType{Int64{}}, s.parents.data(), + &s.validity)); + out.AddColumn(ColumnView::Reference(StorageType{Int64{}}, s.values.data())); + out.Compose(RowSelection::Range(0), count); + out.SetCardinality(count); + s.offset += count; + return true; + } + + private: + struct State : OperatorState { + ~State() override; + uint32_t offset = 0; + std::vector<int64_t> ids; + std::vector<int64_t> parents; + std::vector<int64_t> values; + BitVector validity; + }; + + std::vector<int64_t> parents_; + std::vector<int64_t> values_; + uint32_t chunk_rows_; + std::vector<uint32_t> order_; +}; + +RowSource::State::~State() = default; + +// The two folds written out directly: up sums everything below a node, down +// sums everything above it. +std::vector<int64_t> ReferenceUp(const std::vector<int64_t>& parent, + const std::vector<int64_t>& value) { + std::vector<int64_t> totals(parent.size(), 0); + for (size_t row = 0; row < parent.size(); ++row) { + for (int64_t walk = static_cast<int64_t>(row); walk >= 0; + walk = parent[static_cast<size_t>(walk)]) { + totals[static_cast<size_t>(walk)] += value[row]; + } + } + return totals; +} + +std::vector<int64_t> ReferenceDown(const std::vector<int64_t>& parent, + const std::vector<int64_t>& value) { + std::vector<int64_t> totals(parent.size(), 0); + for (size_t row = 0; row < parent.size(); ++row) { + for (int64_t walk = static_cast<int64_t>(row); walk >= 0; + walk = parent[static_cast<size_t>(walk)]) { + totals[row] += value[static_cast<size_t>(walk)]; + } + } + return totals; +} + +// Runs the whole pipeline, returning the totals by id and the number of +// batches produced. +struct Result { + std::vector<int64_t> totals; + uint32_t batches = 0; +}; + +Result Accumulate(const std::vector<int64_t>& parent, + const std::vector<int64_t>& value, + uint32_t chunk_rows, + bool up, + std::vector<uint32_t> order = {}, + std::optional<TreeRowOrder> arriving = std::nullopt) { + RowSource source(parent, value, chunk_rows, std::move(order)); + // Ids become node numbers before anything else sees them. + std::vector<std::unique_ptr<Operator>> numbering; + numbering.push_back(std::make_unique<TreeNumberNodes>(0, 1)); + Pipeline numbered(source, std::move(numbering)); + + std::unique_ptr<TreeOrder> ordered; + if (up) { + ordered = std::make_unique<TreeChildFirst>(numbered, 3, 4, arriving); + } else { + ordered = std::make_unique<TreeParentFirst>(numbered, 3, 4, arriving); + } + AccumulateSpec spec{3, 4, 2}; + std::vector<std::unique_ptr<Operator>> ops; + if (up) { + ops.push_back(std::make_unique<TreeAccumulateUp>(spec)); + } else { + ops.push_back(std::make_unique<TreeAccumulateDown>(spec)); + } + Pipeline pipeline(*ordered, std::move(ops)); + + std::unique_ptr<OperatorState> state = pipeline.MakeState(); + RowBatch batch; + Result result; + result.totals.assign(parent.size(), 0); + while (pipeline.GetData(batch, *state)) { + ++result.batches; + std::vector<int64_t> ids = test::ReadColumn<int64_t>(batch, 0); + std::vector<int64_t> totals = test::ReadColumn<int64_t>(batch, 5); + for (uint32_t row = 0; row < batch.size(); ++row) { + result.totals[static_cast<size_t>(ids[row])] = totals[row]; + } + } + EXPECT_TRUE(pipeline.status(*state).ok()) + << pipeline.status(*state).message(); + return result; +} + +// A root, its two children and a grandchild. +std::vector<int64_t> Parents() { + return {-1, 0, 0, 1}; +} +std::vector<int64_t> Values() { + return {1, 2, 3, 4}; +} + +TEST(TreeAccumulateTest, UpIsEverythingBelowANode) { + EXPECT_THAT(Accumulate(Parents(), Values(), 8, /*up=*/true).totals, + ElementsAre(10, 6, 3, 4)); +} + +TEST(TreeAccumulateTest, DownIsEverythingAboveANode) { + EXPECT_THAT(Accumulate(Parents(), Values(), 8, /*up=*/false).totals, + ElementsAre(1, 3, 4, 7)); +} + +TEST(TreeAccumulateTest, UpReportsIntegerOverflow) { + std::vector<uint32_t> nodes = {1, 0}; + std::vector<uint32_t> parents = {0, kNoNode}; + std::vector<int64_t> values = {std::numeric_limits<int64_t>::max(), 1}; + RowBatch in; + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, nodes.data())); + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, parents.data())); + in.AddColumn(ColumnView::Reference(StorageType{Int64{}}, values.data())); + in.SetCardinality(2); + + TreeAccumulateUp op({0, 1, 2}); + std::unique_ptr<OperatorState> state = op.MakeState(); + RowBatch out; + EXPECT_EQ(op.Execute(in, out, *state), OpResult::kError); + EXPECT_THAT(op.status(*state).message(), testing::HasSubstr("overflow")); + op.Rewind(*state); + EXPECT_TRUE(op.status(*state).ok()); +} + +TEST(TreeAccumulateTest, DownReportsIntegerOverflow) { + std::vector<uint32_t> nodes = {0, 1}; + std::vector<uint32_t> parents = {kNoNode, 0}; + std::vector<int64_t> values = {std::numeric_limits<int64_t>::max(), 1}; + RowBatch in; + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, nodes.data())); + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, parents.data())); + in.AddColumn(ColumnView::Reference(StorageType{Int64{}}, values.data())); + in.SetCardinality(2); + + TreeAccumulateDown op({0, 1, 2}); + std::unique_ptr<OperatorState> state = op.MakeState(); + RowBatch out; + EXPECT_EQ(op.Execute(in, out, *state), OpResult::kError); + EXPECT_THAT(op.status(*state).message(), testing::HasSubstr("overflow")); +} + +TEST(TreeAccumulateTest, NullValuesContributeZero) { + std::vector<uint32_t> nodes = {0, 1}; + std::vector<uint32_t> parents = {kNoNode, 0}; + std::vector<int64_t> values = {123, 7}; + BitVector validity = BitVector::CreateWithSize(2); + validity.set(1); + RowBatch in; + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, nodes.data())); + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, parents.data())); + in.AddColumn( + ColumnView::Reference(StorageType{Int64{}}, values.data(), &validity)); + in.SetCardinality(2); + + TreeAccumulateDown op({0, 1, 2}); + std::unique_ptr<OperatorState> state = op.MakeState(); + RowBatch out; + ASSERT_EQ(op.Execute(in, out, *state), OpResult::kNeedMoreInput); + EXPECT_THAT(test::ReadColumn<int64_t>(out, 3), ElementsAre(0, 7)); +} + +TEST(TreeAccumulateTest, WrongColumnTypesAreReported) { + std::vector<int64_t> nodes = {0}; + std::vector<uint32_t> parents = {kNoNode}; + std::vector<int64_t> values = {1}; + RowBatch in; + in.AddColumn(ColumnView::Reference(StorageType{Int64{}}, nodes.data())); + in.AddColumn(ColumnView::Reference(StorageType{Uint32{}}, parents.data())); + in.AddColumn(ColumnView::Reference(StorageType{Int64{}}, values.data())); + in.SetCardinality(1); + + TreeAccumulateDown op({0, 1, 2}); + std::unique_ptr<OperatorState> state = op.MakeState(); + RowBatch out; + EXPECT_EQ(op.Execute(in, out, *state), OpResult::kError); + EXPECT_THAT(op.status(*state).message(), testing::HasSubstr("Uint32")); +} + +// Nothing is buffered when the rows already arrive the right way round: one +// batch goes out for every batch that comes in. Any buffering is the ordering +// operator's doing, not the fold's. +TEST(TreeAccumulateTest, NothingIsBuffered) { + std::vector<int64_t> parent(100, -1); + std::vector<int64_t> value(100, 1); + for (uint32_t i = 1; i < 100; ++i) { + parent[i] = i - 1; + } + std::vector<uint32_t> ascending(100); + std::vector<uint32_t> descending(100); + for (uint32_t i = 0; i < 100; ++i) { + ascending[i] = i; + descending[i] = 99 - i; + } + EXPECT_EQ(Accumulate(parent, value, 10, /*up=*/true, descending, + TreeRowOrder::kChildFirst) + .batches, + 10u); + EXPECT_EQ(Accumulate(parent, value, 10, /*up=*/false, ascending, + TreeRowOrder::kParentFirst) + .batches, + 10u); +} + +std::vector<int64_t> RandomParents(std::mt19937& rng, uint32_t rows) { + std::vector<int64_t> parent(rows, -1); + for (uint32_t i = 1; i < rows; ++i) { + if (std::uniform_int_distribution<int>(0, 3)(rng) == 0) { + continue; + } + parent[i] = std::uniform_int_distribution<int64_t>(0, i - 1)(rng); + } + return parent; +} + +TEST(TreeAccumulateTest, MatchesTheDefinitions) { + std::mt19937 rng(11); + for (int trial = 0; trial < 20; ++trial) { + uint32_t rows = std::uniform_int_distribution<uint32_t>(1, 200)(rng); + std::vector<int64_t> parent = RandomParents(rng, rows); + std::vector<int64_t> value(rows); + for (uint32_t i = 0; i < rows; ++i) { + value[i] = std::uniform_int_distribution<int64_t>(-50, 50)(rng); + } + EXPECT_EQ(Accumulate(parent, value, 16, /*up=*/true).totals, + ReferenceUp(parent, value)); + EXPECT_EQ(Accumulate(parent, value, 16, /*up=*/false).totals, + ReferenceDown(parent, value)); + } +} + +TEST(TreeAccumulateTest, TheChunkSizeDoesNotChangeTheAnswer) { + std::mt19937 rng(3); + std::vector<int64_t> parent = RandomParents(rng, 300); + std::vector<int64_t> value(300, 2); + for (uint32_t chunk : {1u, 2u, 7u, 64u, 1024u}) { + EXPECT_EQ(Accumulate(parent, value, chunk, /*up=*/true).totals, + ReferenceUp(parent, value)) + << "chunk " << chunk; + EXPECT_EQ(Accumulate(parent, value, chunk, /*up=*/false).totals, + ReferenceDown(parent, value)) + << "chunk " << chunk; + } +} + +// The running totals carried between batches have to be discarded when the +// plan is run again. +TEST(TreeAccumulateTest, RunningAgainStartsOver) { + RowSource source(Parents(), Values(), 2); + std::vector<std::unique_ptr<Operator>> numbering; + numbering.push_back(std::make_unique<TreeNumberNodes>(0, 1)); + Pipeline numbered(source, std::move(numbering)); + TreeChildFirst order(numbered, 3, 4); + AccumulateSpec spec{3, 4, 2}; + std::vector<std::unique_ptr<Operator>> ops; + ops.push_back(std::make_unique<TreeAccumulateUp>(spec)); + Pipeline pipeline(order, std::move(ops)); + + std::unique_ptr<OperatorState> state = pipeline.MakeState(); + RowBatch batch; + auto drain = [&] { + std::vector<int64_t> totals(4, 0); + while (pipeline.GetData(batch, *state)) { + std::vector<int64_t> ids = test::ReadColumn<int64_t>(batch, 0); + std::vector<int64_t> values = test::ReadColumn<int64_t>(batch, 5); + for (uint32_t row = 0; row < batch.size(); ++row) { + totals[static_cast<size_t>(ids[row])] = values[row]; + } + } + return totals; + }; + std::vector<int64_t> first = drain(); + pipeline.Rewind(*state); + EXPECT_EQ(drain(), first); +} + +// A tree spanning several of the store's chunks, shuffled so that the ordering +// operator has to buffer all of it and hand it back in an order which draws +// each batch's rows from more than one chunk. +TEST(TreeAccumulateTest, ATreeTooBigForOneChunkIsStillFoldedRight) { + std::mt19937 rng(23); + std::vector<int64_t> parent = RandomParents(rng, kMaxBatchRows * 2 + 137); + std::vector<int64_t> value(parent.size()); + for (uint32_t i = 0; i < value.size(); ++i) { + value[i] = std::uniform_int_distribution<int64_t>(-30, 30)(rng); + } + std::vector<uint32_t> order(parent.size()); + for (uint32_t i = 0; i < order.size(); ++i) { + order[i] = i; + } + std::shuffle(order.begin(), order.end(), rng); + + EXPECT_EQ(Accumulate(parent, value, 512, /*up=*/true, order).totals, + ReferenceUp(parent, value)); + EXPECT_EQ(Accumulate(parent, value, 512, /*up=*/false, order).totals, + ReferenceDown(parent, value)); +} + +// Rows in no particular order have to be sorted before either fold can read +// them, which is the ordering operator's job, not this one's. +TEST(TreeAccumulateTest, RowsInNoOrderAreStillFoldedRight) { + std::mt19937 rng(19); + std::vector<int64_t> parent = RandomParents(rng, 250); + std::vector<int64_t> value(parent.size()); + for (uint32_t i = 0; i < value.size(); ++i) { + value[i] = std::uniform_int_distribution<int64_t>(-30, 30)(rng); + } + std::vector<uint32_t> order(parent.size()); + for (uint32_t i = 0; i < order.size(); ++i) { + order[i] = i; + } + std::shuffle(order.begin(), order.end(), rng); + + EXPECT_EQ(Accumulate(parent, value, 32, /*up=*/true, order).totals, + ReferenceUp(parent, value)); + EXPECT_EQ(Accumulate(parent, value, 32, /*up=*/false, order).totals, + ReferenceDown(parent, value)); +} + +} // namespace +} // namespace perfetto::trace_processor::core::exec
diff --git a/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc b/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc index 1266e5b..0ed2fdb 100644 --- a/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc +++ b/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc
@@ -28,12 +28,17 @@ #include "src/trace_processor/containers/string_pool.h" #include "src/trace_processor/core/common/storage_types.h" +#include "src/trace_processor/core/exec/assert_type.h" #include "src/trace_processor/core/exec/column_view.h" #include "src/trace_processor/core/exec/operator.h" +#include "src/trace_processor/core/exec/pipeline.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/row_selection.h" #include "src/trace_processor/core/exec/test_utils.h" +#include "src/trace_processor/core/exec/tree_accumulate.h" +#include "src/trace_processor/core/exec/tree_number_nodes.h" +#include "src/trace_processor/core/exec/tree_order.h" #include "src/trace_processor/core/exec/variant.h" #include "src/trace_processor/core/util/bit_vector.h" #include "src/trace_processor/perfetto_sql/lineage/connection_catalog.h" @@ -346,6 +351,67 @@ EXPECT_THAT(values, testing::ElementsAre(5, 6, 7)); } +// The whole pipeline: a query, its columns asserted to be integers, put into a +// tree order and folded up the tree. +TEST_F(SqlScanTest, AQueryReachesTheTreeOperators) { + Exec( + "CREATE TABLE t AS " + "SELECT 0 AS id, NULL AS parent_id, 10 AS self " + "UNION ALL SELECT 1, 0, 20 " + "UNION ALL SELECT 2, 0, 30 " + "UNION ALL SELECT 3, 1, 40"); + auto scan = Scan("SELECT id, parent_id, self FROM t ORDER BY id"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + + std::vector<std::unique_ptr<core::exec::Operator>> ops; + ops.push_back(std::make_unique<core::exec::AssertType>( + 0, core::exec::AssertTypeTarget{core::Int64{}}, "id")); + ops.push_back(std::make_unique<core::exec::AssertType>( + 1, core::exec::AssertTypeTarget{core::Int64{}}, "parent_id")); + ops.push_back(std::make_unique<core::exec::AssertType>( + 2, core::exec::AssertTypeTarget{core::Int64{}}, "self")); + ops.push_back(std::make_unique<core::exec::TreeNumberNodes>(0, 1)); + core::exec::Pipeline typed(**scan, std::move(ops)); + core::exec::TreeChildFirst order(typed, 3, 4); + core::exec::AccumulateSpec spec{3, 4, 2}; + std::vector<std::unique_ptr<core::exec::Operator>> folds; + folds.push_back(std::make_unique<core::exec::TreeAccumulateUp>(spec)); + core::exec::Pipeline folded(order, std::move(folds)); + + std::unique_ptr<core::exec::OperatorState> state = folded.MakeState(); + RowBatch batch; + std::vector<int64_t> totals(4, 0); + while (folded.GetData(batch, *state)) { + std::vector<int64_t> ids = core::exec::test::ReadColumn<int64_t>(batch, 0); + std::vector<int64_t> values = + core::exec::test::ReadColumn<int64_t>(batch, 5); + for (uint32_t row = 0; row < batch.size(); ++row) { + totals[static_cast<size_t>(ids[row])] = values[row]; + } + } + ASSERT_TRUE(folded.status(*state).ok()) << folded.status(*state).message(); + EXPECT_THAT(totals, testing::ElementsAre(100, 60, 30, 40)); +} + +// A column which is not the type claimed for it fails the pipeline. +TEST_F(SqlScanTest, AColumnWhichIsNotWhatWasAssertedIsReported) { + Exec("CREATE TABLE t(i INTEGER)"); + Exec("INSERT INTO t VALUES(1), ('not a number')"); + auto scan = Scan("SELECT i FROM t"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + + std::vector<std::unique_ptr<core::exec::Operator>> ops; + ops.push_back(std::make_unique<core::exec::AssertType>( + 0, core::exec::AssertTypeTarget{core::Int64{}}, "i")); + core::exec::Pipeline typed(**scan, std::move(ops)); + + Execution run(typed); + while (run.Next()) { + } + EXPECT_FALSE(run.status().ok()); + EXPECT_THAT(run.status().message(), testing::HasSubstr("'i'")); +} + // A column which can be traced back to a dataframe needs neither a variant nor // an assertion: it comes out flat. TEST_F(SqlScanTest, AColumnFollowedBackToADataframeComesOutFlat) {