From 866745d25943ba051b05f0c9baf98b0237b26d29 Mon Sep 17 00:00:00 2001 From: Lalit Maganti Date: Wed, 30 Sep 2026 00:52:17 +0100 Subject: [PATCH] tp: add Sort and GroupBy operators to the executor Operators which need their input ordered should not order it themselves: the plan should, so it can skip the work when the input is already in order. These are the two operators it uses. Sort sorts rows stably by numeric keys through RowLayout. GroupBy puts rows with equal keys next to each other with a counting sort, keeping their order within each group, and numbers the groups. --- Android.bp | 4 + BUILD | 5 + src/trace_processor/core/exec/BUILD.gn | 7 + src/trace_processor/core/exec/group_by.cc | 120 ++++++++++++ src/trace_processor/core/exec/group_by.h | 71 +++++++ .../core/exec/group_by_unittest.cc | 150 +++++++++++++++ src/trace_processor/core/exec/key_encoder.cc | 58 ++---- src/trace_processor/core/exec/layout_column.h | 53 ++++++ src/trace_processor/core/exec/sort.cc | 174 ++++++++++++++++++ src/trace_processor/core/exec/sort.h | 76 ++++++++ .../core/exec/sort_unittest.cc | 163 ++++++++++++++++ 11 files changed, 842 insertions(+), 39 deletions(-) create mode 100644 src/trace_processor/core/exec/group_by.cc create mode 100644 src/trace_processor/core/exec/group_by.h create mode 100644 src/trace_processor/core/exec/group_by_unittest.cc create mode 100644 src/trace_processor/core/exec/layout_column.h create mode 100644 src/trace_processor/core/exec/sort.cc create mode 100644 src/trace_processor/core/exec/sort.h create mode 100644 src/trace_processor/core/exec/sort_unittest.cc diff --git a/Android.bp b/Android.bp index 48456d4146..7ab674c583 100644 --- a/Android.bp +++ b/Android.bp @@ -17822,6 +17822,7 @@ filegroup { "src/trace_processor/core/exec/column_chunk.cc", "src/trace_processor/core/exec/column_view.cc", "src/trace_processor/core/exec/dataframe_scan.cc", + "src/trace_processor/core/exec/group_by.cc", "src/trace_processor/core/exec/interval_intersect.cc", "src/trace_processor/core/exec/key_encoder.cc", "src/trace_processor/core/exec/operator.cc", @@ -17829,6 +17830,7 @@ filegroup { "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/sort.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", @@ -17848,11 +17850,13 @@ filegroup { "src/trace_processor/core/exec/breaker_unittest.cc", "src/trace_processor/core/exec/dataframe_scan_unittest.cc", "src/trace_processor/core/exec/executor_contract_unittest.cc", + "src/trace_processor/core/exec/group_by_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", + "src/trace_processor/core/exec/sort_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/BUILD b/BUILD index d92ec14054..942f005435 100644 --- a/BUILD +++ b/BUILD @@ -2585,11 +2585,14 @@ perfetto_filegroup( "src/trace_processor/core/exec/column_view.h", "src/trace_processor/core/exec/dataframe_scan.cc", "src/trace_processor/core/exec/dataframe_scan.h", + "src/trace_processor/core/exec/group_by.cc", + "src/trace_processor/core/exec/group_by.h", "src/trace_processor/core/exec/interval_columns.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/layout_column.h", "src/trace_processor/core/exec/operator.cc", "src/trace_processor/core/exec/operator.h", "src/trace_processor/core/exec/pipeline.cc", @@ -2601,6 +2604,8 @@ perfetto_filegroup( "src/trace_processor/core/exec/row_selection.h", "src/trace_processor/core/exec/row_store.cc", "src/trace_processor/core/exec/row_store.h", + "src/trace_processor/core/exec/sort.cc", + "src/trace_processor/core/exec/sort.h", "src/trace_processor/core/exec/tree_accumulate.cc", "src/trace_processor/core/exec/tree_accumulate.h", "src/trace_processor/core/exec/tree_number_nodes.cc", diff --git a/src/trace_processor/core/exec/BUILD.gn b/src/trace_processor/core/exec/BUILD.gn index 71f9047c60..f40e2b50e5 100644 --- a/src/trace_processor/core/exec/BUILD.gn +++ b/src/trace_processor/core/exec/BUILD.gn @@ -30,11 +30,14 @@ source_set("exec") { "column_view.h", "dataframe_scan.cc", "dataframe_scan.h", + "group_by.cc", + "group_by.h", "interval_columns.h", "interval_intersect.cc", "interval_intersect.h", "key_encoder.cc", "key_encoder.h", + "layout_column.h", "operator.cc", "operator.h", "pipeline.cc", @@ -46,6 +49,8 @@ source_set("exec") { "row_selection.h", "row_store.cc", "row_store.h", + "sort.cc", + "sort.h", "tree_accumulate.cc", "tree_accumulate.h", "tree_number_nodes.cc", @@ -82,11 +87,13 @@ perfetto_unittest_source_set("unittests") { "breaker_unittest.cc", "dataframe_scan_unittest.cc", "executor_contract_unittest.cc", + "group_by_unittest.cc", "interval_intersect_unittest.cc", "key_encoder_unittest.cc", "operator_unittest.cc", "row_batch_unittest.cc", "row_store_unittest.cc", + "sort_unittest.cc", "tree_accumulate_unittest.cc", "tree_number_nodes_unittest.cc", "tree_order_unittest.cc", diff --git a/src/trace_processor/core/exec/group_by.cc b/src/trace_processor/core/exec/group_by.cc new file mode 100644 index 0000000000..80597b3a55 --- /dev/null +++ b/src/trace_processor/core/exec/group_by.cc @@ -0,0 +1,120 @@ +/* + * 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/group_by.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "perfetto/base/status.h" +#include "src/trace_processor/core/common/storage_types.h" +#include "src/trace_processor/core/exec/column_view.h" +#include "src/trace_processor/core/util/span.h" + +namespace perfetto::trace_processor::core::exec { + +GroupBy::GroupBy(std::vector key_columns) + : key_columns_(std::move(key_columns)) {} +GroupBy::~GroupBy() = default; +GroupBy::State::~State() = default; + +std::unique_ptr GroupBy::CreateState() const { + return std::make_unique(); +} + +bool GroupBy::Consume(const RowBatch& in, Breaker::State& state) const { + auto& s = static_cast(state); + uint32_t count = in.size(); + if (count == 0) { + return true; + } + if (std::optional bad = s.keys.Encode(in, key_columns_)) { + s.status = base::ErrStatus( + "GROUP BY: key %u must be a column holding one type", *bad + 1); + return false; + } + for (uint32_t row = 0; row < count; ++row) { + std::string_view key = s.keys.Key(row); + uint32_t group; + if (const uint32_t* found = s.group_of.Find(key)) { + group = *found; + } else { + group = static_cast(s.group_of.size()); + s.group_of.Insert(std::string(key), group); + } + s.groups.push_back(group); + } + s.status = s.rows.Append(in); + return s.status.ok(); +} + +bool GroupBy::Finalize(Breaker::State& state) const { + auto& s = static_cast(state); + // A counting sort on the group. + auto groups = static_cast(s.group_of.size()); + auto next = FlexVector::CreateFilled(groups, 0); + for (uint32_t group : s.groups) { + ++next[group]; + } + uint32_t start = 0; + for (uint32_t& at : next) { + uint32_t rows = at; + at = start; + start += rows; + } + s.order.resize(s.groups.size()); + s.ordered_groups.resize(s.groups.size()); + for (uint32_t row = 0; row < s.groups.size(); ++row) { + uint32_t at = next[s.groups[row]]++; + s.order[at] = row; + s.ordered_groups[at] = s.groups[row]; + } + return true; +} + +bool GroupBy::Serve(RowBatch& out, Breaker::State& state) const { + auto& s = static_cast(state); + if (s.emitted == s.rows.size()) { + return false; + } + uint32_t count = std::min(kMaxBatchRows, s.rows.size() - s.emitted); + const uint32_t* begin = s.order.data() + s.emitted; + count = s.rows.View(&out, Span(begin, begin + count)); + ColumnView group = + ColumnView::Reference(StorageType{Uint32{}}, s.ordered_groups.data()); + group.SetRange(s.emitted); + out.AddColumn(group); + s.emitted += count; + return true; +} + +void GroupBy::Reset(Breaker::State& state) const { + auto& s = static_cast(state); + s.group_of.Clear(); + s.groups.clear(); + s.rows.Clear(); + s.order.clear(); + s.ordered_groups.clear(); + s.emitted = 0; +} + +} // namespace perfetto::trace_processor::core::exec diff --git a/src/trace_processor/core/exec/group_by.h b/src/trace_processor/core/exec/group_by.h new file mode 100644 index 0000000000..74ceed6e95 --- /dev/null +++ b/src/trace_processor/core/exec/group_by.h @@ -0,0 +1,71 @@ +/* + * 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_GROUP_BY_H_ +#define SRC_TRACE_PROCESSOR_CORE_EXEC_GROUP_BY_H_ + +#include +#include +#include +#include + +#include "perfetto/ext/base/flat_hash_map.h" +#include "src/trace_processor/core/exec/breaker.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_store.h" +#include "src/trace_processor/core/util/flex_vector.h" + +namespace perfetto::trace_processor::core::exec { + +// Puts rows sharing their keys next to each other. Stable, so an order the +// input had holds within each group. Appends each row's group number as a +// Uint32 column, so what follows finds group boundaries without comparing +// keys. +class GroupBy : public Breaker { + public: + explicit GroupBy(std::vector key_columns); + ~GroupBy() override; + BatchPreference batch_preference() const override { + return BatchPreference::kThroughput; + } + + private: + struct State : Breaker::State { + ~State() override; + + KeyEncoder keys; + base::FlatHashMapV2 group_of; + FlexVector groups; + RowStore rows; + FlexVector order; + FlexVector ordered_groups; + uint32_t emitted = 0; + }; + + std::unique_ptr CreateState() const override; + bool Consume(const RowBatch& in, Breaker::State& state) const override; + bool Finalize(Breaker::State& state) const override; + bool Serve(RowBatch& out, Breaker::State& state) const override; + void Reset(Breaker::State& state) const override; + + std::vector key_columns_; +}; + +} // namespace perfetto::trace_processor::core::exec + +#endif // SRC_TRACE_PROCESSOR_CORE_EXEC_GROUP_BY_H_ diff --git a/src/trace_processor/core/exec/group_by_unittest.cc b/src/trace_processor/core/exec/group_by_unittest.cc new file mode 100644 index 0000000000..2267192cae --- /dev/null +++ b/src/trace_processor/core/exec/group_by_unittest.cc @@ -0,0 +1,150 @@ +/* + * 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/group_by.h" + +#include +#include +#include +#include +#include +#include + +#include "perfetto/base/status.h" +#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/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 "test/gtest_and_gmock.h" + +namespace perfetto::trace_processor::core::exec { +namespace { + +using testing::ElementsAre; + +class StringSource final : public Source { + public: + StringSource(std::vector keys, uint32_t chunk) + : keys_(std::move(keys)), chunk_(chunk) {} + + std::unique_ptr MakeState() const override { + return std::make_unique(); + } + void Rewind(OperatorState& state) const override { + state.Cast().emitted = 0; + } + bool GetData(RowBatch& out, OperatorState& state) const override { + State& s = state.Cast(); + auto rows = static_cast(keys_.size()); + if (s.emitted == rows) { + return false; + } + uint32_t count = std::min(chunk_, rows - s.emitted); + out.Reset(); + out.AddColumn(ColumnView::Reference(StorageType{Id{}}, nullptr)); + out.AddColumn(ColumnView::Reference(StorageType{String{}}, keys_.data())); + out.Compose(RowSelection::Range(s.emitted), count); + out.SetCardinality(count); + s.emitted += count; + return true; + } + + private: + struct State : OperatorState { + uint32_t emitted = 0; + }; + + std::vector keys_; + uint32_t chunk_; +}; + +struct Output { + std::vector ids; + std::vector groups; + base::Status status; +}; + +Output Drain(const Source& source) { + std::unique_ptr state = source.MakeState(); + RowBatch batch; + Output out; + while (source.GetData(batch, *state)) { + for (uint32_t id : test::ReadColumn(batch, 0)) { + out.ids.push_back(id); + } + for (uint32_t group : test::ReadColumn(batch, 2)) { + out.groups.push_back(group); + } + } + out.status = source.status(*state); + return out; +} + +std::vector> GroupOn(uint32_t column) { + std::vector> ops; + ops.push_back(std::make_unique(std::vector{column})); + return ops; +} + +TEST(GroupByTest, GroupsInOrderOfFirstRowKeepingTheirRowsInOrder) { + test::ArraySource source({7, 3, 7, 5, 3, 7}); + Pipeline grouped(source, GroupOn(1), {}); + Output out = Drain(grouped); + ASSERT_TRUE(out.status.ok()) << out.status.message(); + EXPECT_THAT(out.ids, ElementsAre(0, 2, 5, 1, 4, 3)); + EXPECT_THAT(out.groups, ElementsAre(0, 0, 0, 1, 1, 2)); +} + +TEST(GroupByTest, KeepsOrderWithinGroupsAcrossManyBatches) { + std::vector values(3 * kMaxBatchRows + 17); + std::mt19937 rng(3); + for (int64_t& value : values) { + value = static_cast(rng() % 97); + } + test::ArraySource source(values); + Pipeline grouped(source, GroupOn(1), {}); + Output out = Drain(grouped); + ASSERT_TRUE(out.status.ok()) << out.status.message(); + ASSERT_EQ(out.ids.size(), values.size()); + for (uint32_t i = 1; i < out.ids.size(); ++i) { + if (out.groups[i] == out.groups[i - 1]) { + EXPECT_EQ(values[out.ids[i]], values[out.ids[i - 1]]); + EXPECT_LT(out.ids[i - 1], out.ids[i]); + } else { + EXPECT_EQ(out.groups[i], out.groups[i - 1] + 1); + } + } +} + +TEST(GroupByTest, GroupsStrings) { + StringPool pool; + StringPool::Id a = pool.InternString("a"); + StringPool::Id b = pool.InternString("b"); + StringSource source({b, StringPool::Id::Null(), a, b, StringPool::Id::Null()}, + 2); + Pipeline grouped(source, GroupOn(1), {}); + Output out = Drain(grouped); + ASSERT_TRUE(out.status.ok()) << out.status.message(); + EXPECT_THAT(out.ids, ElementsAre(0, 3, 1, 4, 2)); + EXPECT_THAT(out.groups, ElementsAre(0, 0, 1, 1, 2)); +} + +} // namespace +} // namespace perfetto::trace_processor::core::exec diff --git a/src/trace_processor/core/exec/key_encoder.cc b/src/trace_processor/core/exec/key_encoder.cc index 65c3601bdf..a1cd13cf47 100644 --- a/src/trace_processor/core/exec/key_encoder.cc +++ b/src/trace_processor/core/exec/key_encoder.cc @@ -26,36 +26,12 @@ #include "src/trace_processor/core/common/row_layout.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/layout_column.h" #include "src/trace_processor/core/exec/row_batch.h" #include "src/trace_processor/core/exec/row_selection.h" #include "src/trace_processor/core/util/bit_vector.h" namespace perfetto::trace_processor::core::exec { -namespace { - -template -void Write(const ColumnView& column, - const RowLayout::Slot& slot, - uint32_t count, - uint8_t* rows, - Value value) { - RowSelection selection = column.selection(); - const BitVector* validity = column.validity(); - RowLayout::Write( - slot, count, - [&](uint32_t row, T* out) { - uint32_t index = selection.GetIndex(row); - if (validity && !validity->is_set(index)) { - return false; - } - *out = value(index); - return true; - }, - rows); -} - -} // namespace - std::optional KeyEncoder::Encode( const RowBatch& batch, const std::vector& columns) { @@ -118,31 +94,35 @@ std::optional KeyEncoder::Encode( const ColumnView& column = batch.column(columns[k]); const RowLayout::Slot& slot = layout_.slot(k); if (column.kind() == ColumnView::Kind::kSequence) { - Write(column, slot, count, rows, - [](uint32_t index) { return int64_t{index}; }); + WriteLayoutColumn(column, slot, count, rows, + [](uint32_t index) { return int64_t{index}; }); continue; } const void* data = column.data(); switch (column.type().index()) { case StorageType::GetTypeIndex(): - Write(column, slot, count, rows, [data](uint32_t index) { - return int64_t{static_cast(data)[index]}; - }); + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return int64_t{static_cast(data)[index]}; + }); break; case StorageType::GetTypeIndex(): - Write(column, slot, count, rows, [data](uint32_t index) { - return int64_t{static_cast(data)[index]}; - }); + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return int64_t{static_cast(data)[index]}; + }); break; case StorageType::GetTypeIndex(): - Write(column, slot, count, rows, [data](uint32_t index) { - return static_cast(data)[index]; - }); + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); break; case StorageType::GetTypeIndex(): - Write(column, slot, count, rows, [data](uint32_t index) { - return static_cast(data)[index]; - }); + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); break; case StorageType::GetTypeIndex(): { // Not nullable: the null id is 0. diff --git a/src/trace_processor/core/exec/layout_column.h b/src/trace_processor/core/exec/layout_column.h new file mode 100644 index 0000000000..309c1dc93c --- /dev/null +++ b/src/trace_processor/core/exec/layout_column.h @@ -0,0 +1,53 @@ +/* + * 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_LAYOUT_COLUMN_H_ +#define SRC_TRACE_PROCESSOR_CORE_EXEC_LAYOUT_COLUMN_H_ + +#include + +#include "src/trace_processor/core/common/row_layout.h" +#include "src/trace_processor/core/exec/column_view.h" +#include "src/trace_processor/core/exec/row_selection.h" +#include "src/trace_processor/core/util/bit_vector.h" + +namespace perfetto::trace_processor::core::exec { + +// `value(index)` converts the value at physical row `index` into a T. +template +void WriteLayoutColumn(const ColumnView& column, + const RowLayout::Slot& slot, + uint32_t count, + uint8_t* rows, + Value value) { + RowSelection selection = column.selection(); + const BitVector* validity = column.validity(); + RowLayout::Write( + slot, count, + [&](uint32_t row, T* out) { + uint32_t index = selection.GetIndex(row); + if (validity && !validity->is_set(index)) { + return false; + } + *out = value(index); + return true; + }, + rows); +} + +} // namespace perfetto::trace_processor::core::exec + +#endif // SRC_TRACE_PROCESSOR_CORE_EXEC_LAYOUT_COLUMN_H_ diff --git a/src/trace_processor/core/exec/sort.cc b/src/trace_processor/core/exec/sort.cc new file mode 100644 index 0000000000..9d83a29960 --- /dev/null +++ b/src/trace_processor/core/exec/sort.cc @@ -0,0 +1,174 @@ +/* + * 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/sort.h" + +#include +#include +#include +#include +#include +#include + +#include "perfetto/base/logging.h" +#include "perfetto/base/status.h" +#include "src/trace_processor/core/common/row_layout.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/layout_column.h" +#include "src/trace_processor/core/util/ops.h" +#include "src/trace_processor/core/util/span.h" + +namespace perfetto::trace_processor::core::exec { +namespace { + +std::optional LayoutType(const ColumnView& column) { + switch (column.kind()) { + case ColumnView::Kind::kSequence: + return RowLayout::Type::kUint32; + case ColumnView::Kind::kFlat: + switch (column.type().index()) { + case StorageType::GetTypeIndex(): + return RowLayout::Type::kUint32; + case StorageType::GetTypeIndex(): + return RowLayout::Type::kInt32; + case StorageType::GetTypeIndex(): + return RowLayout::Type::kInt64; + case StorageType::GetTypeIndex(): + return RowLayout::Type::kDouble; + default: + return std::nullopt; + } + case ColumnView::Kind::kVariant: + return std::nullopt; + } + PERFETTO_FATAL("Unknown column kind"); +} + +} // namespace + +Sort::Sort(SortSpec spec) : spec_(std::move(spec)) {} +Sort::~Sort() = default; +Sort::State::~State() = default; + +std::unique_ptr Sort::CreateState() const { + auto state = std::make_unique(); + state->types.resize(spec_.keys.size()); + return state; +} + +bool Sort::Consume(const RowBatch& in, Breaker::State& state) const { + auto& s = static_cast(state); + uint32_t count = in.size(); + if (count == 0) { + return true; + } + bool first = !s.types.empty() && !s.types[0]; + for (uint32_t k = 0; k < spec_.keys.size(); ++k) { + std::optional type = + LayoutType(in.column(spec_.keys[k].column)); + if (!type || (s.types[k] && *s.types[k] != *type)) { + s.status = base::ErrStatus( + "SORT: key %u must be a column of numbers of one type", k + 1); + return false; + } + s.types[k] = type; + } + if (first) { + // Nullability is only known batch by batch. + std::vector columns; + for (uint32_t k = 0; k < spec_.keys.size(); ++k) { + columns.push_back({*s.types[k], true, spec_.keys[k].descending}); + } + s.layout = RowLayout(columns); + } + + uint64_t at = uint64_t{s.rows.size()} * s.layout.stride(); + s.keys.resize(at + uint64_t{count} * s.layout.stride()); + uint8_t* rows = s.keys.data() + at; + for (uint32_t k = 0; k < spec_.keys.size(); ++k) { + const ColumnView& column = in.column(spec_.keys[k].column); + const RowLayout::Slot& slot = s.layout.slot(k); + if (column.kind() == ColumnView::Kind::kSequence) { + WriteLayoutColumn(column, slot, count, rows, + [](uint32_t index) { return index; }); + continue; + } + const void* data = column.data(); + switch (*s.types[k]) { + case RowLayout::Type::kUint32: + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); + break; + case RowLayout::Type::kInt32: + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); + break; + case RowLayout::Type::kInt64: + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); + break; + case RowLayout::Type::kDouble: + WriteLayoutColumn( + column, slot, count, rows, [data](uint32_t index) { + return static_cast(data)[index]; + }); + break; + } + } + s.status = s.rows.Append(in); + return s.status.ok(); +} + +bool Sort::Finalize(Breaker::State& state) const { + auto& s = static_cast(state); + s.order.resize(s.rows.size()); + for (uint32_t row = 0; row < s.rows.size(); ++row) { + s.order[row] = row; + } + Span order(s.order.data(), s.order.data() + s.order.size()); + core::ops::SortRowLayout( + Span(s.keys.data(), s.keys.data() + s.keys.size()), + s.layout.stride(), &order); + return true; +} + +bool Sort::Serve(RowBatch& out, Breaker::State& state) const { + auto& s = static_cast(state); + if (s.emitted == s.rows.size()) { + return false; + } + uint32_t count = std::min(kMaxBatchRows, s.rows.size() - s.emitted); + const uint32_t* begin = s.order.data() + s.emitted; + s.emitted += s.rows.View(&out, Span(begin, begin + count)); + return true; +} + +void Sort::Reset(Breaker::State& state) const { + auto& s = static_cast(state); + s.keys.clear(); + s.rows.Clear(); + s.order.clear(); + s.emitted = 0; +} + +} // namespace perfetto::trace_processor::core::exec diff --git a/src/trace_processor/core/exec/sort.h b/src/trace_processor/core/exec/sort.h new file mode 100644 index 0000000000..9e637f362a --- /dev/null +++ b/src/trace_processor/core/exec/sort.h @@ -0,0 +1,76 @@ +/* + * 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_SORT_H_ +#define SRC_TRACE_PROCESSOR_CORE_EXEC_SORT_H_ + +#include +#include +#include +#include + +#include "src/trace_processor/core/common/row_layout.h" +#include "src/trace_processor/core/exec/breaker.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_store.h" +#include "src/trace_processor/core/util/flex_vector.h" + +namespace perfetto::trace_processor::core::exec { + +struct SortSpec { + struct Key { + uint32_t column = 0; + bool descending = false; + }; + std::vector keys; +}; + +// Stably sorts rows by their keys. Nulls sort as in SQLite: first ascending, +// last descending. +// TODO(lalitm): support string keys, which needs the string pool for ranks. +class Sort : public Breaker { + public: + explicit Sort(SortSpec); + ~Sort() override; + BatchPreference batch_preference() const override { + return BatchPreference::kThroughput; + } + + private: + struct State : Breaker::State { + ~State() override; + + std::vector> types; + RowLayout layout; + FlexVector keys; + RowStore rows; + FlexVector order; + uint32_t emitted = 0; + }; + + std::unique_ptr CreateState() const override; + bool Consume(const RowBatch& in, Breaker::State& state) const override; + bool Finalize(Breaker::State& state) const override; + bool Serve(RowBatch& out, Breaker::State& state) const override; + void Reset(Breaker::State& state) const override; + + SortSpec spec_; +}; + +} // namespace perfetto::trace_processor::core::exec + +#endif // SRC_TRACE_PROCESSOR_CORE_EXEC_SORT_H_ diff --git a/src/trace_processor/core/exec/sort_unittest.cc b/src/trace_processor/core/exec/sort_unittest.cc new file mode 100644 index 0000000000..0d375a7eea --- /dev/null +++ b/src/trace_processor/core/exec/sort_unittest.cc @@ -0,0 +1,163 @@ +/* + * 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/sort.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "perfetto/base/status.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/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/util/bit_vector.h" +#include "test/gtest_and_gmock.h" + +namespace perfetto::trace_processor::core::exec { +namespace { + +using testing::ElementsAre; + +class KeysSource final : public Source { + public: + KeysSource(std::vector> first, + std::vector second, + uint32_t chunk) + : second_(std::move(second)), chunk_(chunk) { + validity_ = BitVector::CreateWithSize(first.size()); + for (uint32_t i = 0; i < first.size(); ++i) { + first_.push_back(first[i].value_or(0)); + if (first[i]) { + validity_.set(i); + } + } + } + + std::unique_ptr MakeState() const override { + return std::make_unique(); + } + void Rewind(OperatorState& state) const override { + state.Cast().emitted = 0; + } + bool GetData(RowBatch& out, OperatorState& state) const override { + State& s = state.Cast(); + auto rows = static_cast(first_.size()); + if (s.emitted == rows) { + return false; + } + uint32_t count = std::min(chunk_, rows - s.emitted); + out.Reset(); + out.AddColumn(ColumnView::Reference(StorageType{Id{}}, nullptr)); + out.AddColumn( + ColumnView::Reference(StorageType{Int64{}}, first_.data(), &validity_)); + out.AddColumn(ColumnView::Reference(StorageType{Double{}}, second_.data())); + out.Compose(RowSelection::Range(s.emitted), count); + out.SetCardinality(count); + s.emitted += count; + return true; + } + + private: + struct State : OperatorState { + uint32_t emitted = 0; + }; + + std::vector first_; + BitVector validity_; + std::vector second_; + uint32_t chunk_; +}; + +std::vector Ids(const Source& source, base::Status* status) { + std::unique_ptr state = source.MakeState(); + RowBatch batch; + std::vector ids; + while (source.GetData(batch, *state)) { + for (uint32_t id : test::ReadColumn(batch, 0)) { + ids.push_back(id); + } + } + *status = source.status(*state); + return ids; +} + +std::vector> SortBy(std::vector keys) { + std::vector> ops; + ops.push_back(std::make_unique(SortSpec{std::move(keys)})); + return ops; +} + +TEST(SortTest, OrdersManyBatchesStably) { + // Few distinct values across several batches, so most rows tie. + std::vector values(3 * kMaxBatchRows + 17); + std::mt19937 rng(7); + for (int64_t& value : values) { + value = static_cast(rng() % 50) - 25; + } + test::ArraySource source(values); + Pipeline sorted(source, SortBy({{1, false}}), {}); + + base::Status status; + std::vector ids = Ids(sorted, &status); + ASSERT_TRUE(status.ok()) << status.message(); + std::vector expected(values.size()); + std::iota(expected.begin(), expected.end(), 0u); + std::stable_sort( + expected.begin(), expected.end(), + [&](uint32_t a, uint32_t b) { return values[a] < values[b]; }); + EXPECT_EQ(ids, expected); +} + +TEST(SortTest, NullsSortAsInSqlite) { + KeysSource source({3, std::nullopt, 1, std::nullopt}, {0, 0, 0, 0}, 3); + base::Status status; + { + Pipeline sorted(source, SortBy({{1, false}}), {}); + EXPECT_THAT(Ids(sorted, &status), ElementsAre(1, 3, 2, 0)); + } + { + Pipeline sorted(source, SortBy({{1, true}}), {}); + EXPECT_THAT(Ids(sorted, &status), ElementsAre(0, 2, 1, 3)); + } +} + +TEST(SortTest, LaterKeysBreakTies) { + KeysSource source({2, 1, 2, 1}, {0.5, -1.5, -0.5, 2.5}, 2); + Pipeline sorted(source, SortBy({{1, false}, {2, true}}), {}); + base::Status status; + EXPECT_THAT(Ids(sorted, &status), ElementsAre(3, 1, 0, 2)); + EXPECT_TRUE(status.ok()) << status.message(); +} + +TEST(SortTest, IdsSortAsTheirValues) { + test::ArraySource source({5, 4, 3}); + Pipeline sorted(source, SortBy({{0, true}}), {}); + base::Status status; + EXPECT_THAT(Ids(sorted, &status), ElementsAre(2, 1, 0)); +} + +} // namespace +} // namespace perfetto::trace_processor::core::exec