tp: scan SQL queries into batches Queries which filter, join, group, or compute still need SQLite. SqlScan steps a prepared statement and fills RowBatch instances instead of exposing the row-at-a-time cursor to the rest of the pipeline. The shared relation analyzer decides how each result column is stored. Columns traced to a dataframe come out flat at the proven type. Other columns use Variant because SQLite may return a different type on each row. Preparing, stepping, and rewinding the statement all follow the Source status contract. Type validation is left to the operator which knows how a column will be used.
diff --git a/Android.bp b/Android.bp index 3dd43eb..5287f42 100644 --- a/Android.bp +++ b/Android.bp
@@ -18876,6 +18876,22 @@ ], } +// GN: //src/trace_processor/perfetto_sql/exec:exec +filegroup { + name: "perfetto_src_trace_processor_perfetto_sql_exec_exec", + srcs: [ + "src/trace_processor/perfetto_sql/exec/sql_scan.cc", + ], +} + +// GN: //src/trace_processor/perfetto_sql/exec:unittests +filegroup { + name: "perfetto_src_trace_processor_perfetto_sql_exec_unittests", + srcs: [ + "src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc", + ], +} + // GN: //src/trace_processor/perfetto_sql/generator:gen_cc_perfetto_sql_descriptor genrule { name: "perfetto_src_trace_processor_perfetto_sql_generator_gen_cc_perfetto_sql_descriptor", @@ -24170,6 +24186,8 @@ ":perfetto_src_trace_processor_metrics_unittests", ":perfetto_src_trace_processor_perfetto_sql_engine_engine", ":perfetto_src_trace_processor_perfetto_sql_engine_unittests", + ":perfetto_src_trace_processor_perfetto_sql_exec_exec", + ":perfetto_src_trace_processor_perfetto_sql_exec_unittests", ":perfetto_src_trace_processor_perfetto_sql_generator_generator", ":perfetto_src_trace_processor_perfetto_sql_generator_unittests", ":perfetto_src_trace_processor_perfetto_sql_intrinsics_types_types",
diff --git a/src/perfetto_sql/syntaqlite/BUILD.gn b/src/perfetto_sql/syntaqlite/BUILD.gn index 7e89d9f..7e09769 100644 --- a/src/perfetto_sql/syntaqlite/BUILD.gn +++ b/src/perfetto_sql/syntaqlite/BUILD.gn
@@ -36,6 +36,7 @@ assert_no_deps = [ "../../trace_processor/*" ] visibility = [ "..:intrinsic_macro_expansion", + "../../trace_processor/perfetto_sql/exec:*", "../../trace_processor/perfetto_sql/lineage:*", "../../trace_processor/perfetto_sql/parser:*", "../../trace_processor/perfetto_sql/tokenizer",
diff --git a/src/trace_processor/BUILD.gn b/src/trace_processor/BUILD.gn index 73cd62e..683e69a 100644 --- a/src/trace_processor/BUILD.gn +++ b/src/trace_processor/BUILD.gn
@@ -470,6 +470,7 @@ deps += [ "../perfetto_sql/analysis:unittests", "perfetto_sql/engine:unittests", + "perfetto_sql/exec:unittests", "perfetto_sql/lineage:unittests", "perfetto_sql/parser:unittests", "perfetto_sql/tokenizer:unittests",
diff --git a/src/trace_processor/perfetto_sql/exec/BUILD.gn b/src/trace_processor/perfetto_sql/exec/BUILD.gn new file mode 100644 index 0000000..3dd9917 --- /dev/null +++ b/src/trace_processor/perfetto_sql/exec/BUILD.gn
@@ -0,0 +1,54 @@ +# 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. + +import("../../../../gn/test.gni") + +source_set("exec") { + sources = [ + "sql_scan.cc", + "sql_scan.h", + ] + deps = [ + "../../../../gn:default_deps", + "../../../../gn:sqlite", + "../../../base", + "../../../perfetto_sql/analysis", + "../../../perfetto_sql/syntaqlite", + "../../containers", + "../../core/common", + "../../core/exec", + "../../core/util", + "../../sqlite", + "../lineage:connection_catalog", + ] +} + +perfetto_unittest_source_set("unittests") { + testonly = true + sources = [ "sql_scan_unittest.cc" ] + deps = [ + ":exec", + "../../../../gn:default_deps", + "../../../../gn:gtest_and_gmock", + "../../../base", + "../../../perfetto_sql/analysis", + "../../containers", + "../../core/common", + "../../core/exec", + "../../core/exec:test_utils", + "../../core/util", + "../../sqlite", + "../lineage:connection_catalog", + ] +}
diff --git a/src/trace_processor/perfetto_sql/exec/sql_scan.cc b/src/trace_processor/perfetto_sql/exec/sql_scan.cc new file mode 100644 index 0000000..1a6b043 --- /dev/null +++ b/src/trace_processor/perfetto_sql/exec/sql_scan.cc
@@ -0,0 +1,334 @@ +/* + * 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/perfetto_sql/exec/sql_scan.h" + +#include <sqlite3.h> + +#include <cstdint> +#include <memory> +#include <optional> +#include <string> +#include <utility> +#include <vector> + +#include "perfetto/base/status.h" +#include "perfetto/ext/base/status_macros.h" +#include "perfetto/ext/base/status_or.h" +#include "src/perfetto_sql/analysis/relation.h" +#include "src/perfetto_sql/syntaqlite/syntaqlite_perfetto.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/row_batch.h" +#include "src/trace_processor/core/exec/row_selection.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" +#include "src/trace_processor/sqlite/sql_source.h" +#include "src/trace_processor/sqlite/sqlite_connection.h" + +namespace perfetto::trace_processor::exec { +namespace { + +using core::Double; +using core::Int64; +using core::StorageType; +using core::String; +using core::exec::ColumnView; +using core::exec::kMaxBatchRows; +using core::exec::RowBatch; +using core::exec::RowSelection; +using core::exec::Variant; +struct ParserDeleter { + void operator()(SyntaqliteParser* parser) const { + syntaqlite_parser_destroy(parser); + } +}; +using ScopedParser = std::unique_ptr<SyntaqliteParser, ParserDeleter>; + +// The types lineage established, lined up with the query's columns. If the two +// disagree on the number of columns they are not describing the same query, so +// no type is claimed for any of them. +std::vector<std::optional<StorageType>> Prove(const SqlSource& sql, + uint32_t count, + const lineage::Catalog* catalog) { + std::vector<std::optional<StorageType>> types(count); + if (!catalog) { + return types; + } + + ScopedParser parser(syntaqlite_parser_create_perfetto(nullptr)); + syntaqlite_parser_reset(parser.get(), sql.sql().data(), + static_cast<uint32_t>(sql.sql().size())); + if (syntaqlite_parser_next(parser.get()) != SYNTAQLITE_PARSE_OK) { + return types; + } + lineage::analysis::RelationAnalyzer analyzer(*catalog); + auto resolved = analyzer.AnalyzeQuery( + {parser.get(), syntaqlite_result_root(parser.get())}); + if (!resolved.ok() || resolved->columns().size() != count) { + return types; + } + for (uint32_t i = 0; i < count; ++i) { + std::optional<StorageType> type = + catalog->ColumnType(resolved->columns()[i]); + // An Id has no storage of its own: its value is the row it sits at. A + // query result has no such rows to point at, so materialise it at the + // narrowest width which holds one. + if (type && type->Is<core::Id>()) { + type = StorageType{core::Uint32{}}; + } + types[i] = type; + } + return types; +} + +} // namespace + +base::StatusOr<std::unique_ptr<SqlScan>> SqlScan::Create( + SqliteConnection* connection, + SqlSource sql, + StringPool* pool, + const lineage::Catalog* catalog) { + // Prepared here only to read the column names, then discarded: a statement + // belongs to one execution, but the columns belong to the query. + SqliteConnection::PreparedStatement statement = + connection->PrepareStatement(sql); + RETURN_IF_ERROR(statement.status()); + + sqlite3_stmt* stmt = statement.sqlite_stmt(); + auto count = static_cast<uint32_t>(sqlite3_column_count(stmt)); + std::vector<std::string> names; + names.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + const char* name = sqlite3_column_name(stmt, static_cast<int>(i)); + names.emplace_back(name ? name : ""); + } + std::vector<std::optional<StorageType>> types = Prove(sql, count, catalog); + return std::unique_ptr<SqlScan>(new SqlScan( + connection, std::move(sql), std::move(names), std::move(types), pool)); +} + +SqlScan::SqlScan(SqliteConnection* connection, + SqlSource sql, + std::vector<std::string> names, + std::vector<std::optional<StorageType>> types, + StringPool* pool) + : connection_(connection), + sql_(std::move(sql)), + names_(std::move(names)), + types_(std::move(types)), + pool_(pool) {} + +SqlScan::~SqlScan() = default; +SqlScan::State::~State() = default; + +std::unique_ptr<core::exec::OperatorState> SqlScan::MakeState() const { + auto state = std::make_unique<State>(); + Prepare(*state); + state->columns.reserve(names_.size()); + for (uint32_t i = 0; i < names_.size(); ++i) { + auto column = std::make_shared<Column>(); + if (!types_[i]) { + column->variants.resize(kMaxBatchRows); + } else if (types_[i]->Is<core::Uint32>()) { + column->uint32s.resize(kMaxBatchRows); + } else if (types_[i]->Is<core::Int32>()) { + column->int32s.resize(kMaxBatchRows); + } else if (types_[i]->Is<Int64>()) { + column->ints.resize(kMaxBatchRows); + } else if (types_[i]->Is<Double>()) { + column->doubles.resize(kMaxBatchRows); + } else { + column->strings.resize(kMaxBatchRows); + } + if (types_[i]) { + column->validity = core::BitVector::CreateWithSize(kMaxBatchRows); + } + state->columns.push_back(std::move(column)); + } + return state; +} + +void SqlScan::Prepare(State& state) const { + state.statement.emplace(connection_->PrepareStatement(sql_)); + state.status = state.statement->status(); + state.done = false; + if (!state.status.ok()) { + return; + } + sqlite3_stmt* stmt = state.statement->sqlite_stmt(); + uint32_t count = static_cast<uint32_t>(sqlite3_column_count(stmt)); + if (count != names_.size()) { + state.status = + base::ErrStatus("SQL source: result shape changed between executions"); + return; + } + for (uint32_t i = 0; i < count; ++i) { + const char* name = sqlite3_column_name(stmt, static_cast<int>(i)); + if (names_[i] != (name ? name : "")) { + state.status = base::ErrStatus( + "SQL source: result shape changed between executions"); + return; + } + } +} + +base::Status SqlScan::status(const core::exec::OperatorState& state) const { + return state.Cast<const State>().status; +} + +void SqlScan::Rewind(core::exec::OperatorState& state) const { + Prepare(state.Cast<State>()); +} + +bool SqlScan::ReadValue(State& s, + sqlite3_stmt* stmt, + uint32_t index, + uint32_t row) const { + auto col = static_cast<int>(index); + Column& column = *s.columns[index]; + int type = sqlite3_column_type(stmt, col); + if (!types_[index]) { + switch (type) { + case SQLITE_INTEGER: + column.variants[row] = Variant::Int64(sqlite3_column_int64(stmt, col)); + return true; + case SQLITE_FLOAT: + column.variants[row] = + Variant::Double(sqlite3_column_double(stmt, col)); + return true; + case SQLITE_TEXT: + column.variants[row] = Variant::String(pool_->InternString( + reinterpret_cast<const char*>(sqlite3_column_text(stmt, col)))); + return true; + case SQLITE_NULL: + column.variants[row] = Variant::Null(); + return true; + default: + s.status = base::ErrStatus( + "SQL source: column '%s' holds a blob, which a pipeline cannot " + "carry", + names_[index].c_str()); + return false; + } + } + StorageType proven = *types_[index]; + if (type == SQLITE_NULL) { + // The row is null, but write the slot anyway. A flat column's storage is + // readable at every row, so a reader summing it needs no per-row branch and + // never sees a value left over from the previous batch. + if (proven.Is<core::Uint32>()) { + column.uint32s[row] = 0; + } else if (proven.Is<core::Int32>()) { + column.int32s[row] = 0; + } else if (proven.Is<Int64>()) { + column.ints[row] = 0; + } else if (proven.Is<Double>()) { + column.doubles[row] = 0; + } else { + column.strings[row] = StringPool::Id::Null(); + } + return true; + } + bool integer = proven.Is<Int64>() || proven.Is<core::Uint32>() || + proven.Is<core::Int32>(); + bool matches = (type == SQLITE_INTEGER && integer) || + (type == SQLITE_FLOAT && proven.Is<Double>()) || + (type == SQLITE_TEXT && proven.Is<String>()); + if (!matches) { + // Only reachable if the type lineage established turned out to be wrong. + s.status = base::ErrStatus( + "SQL source: column '%s' does not hold what it was traced back to", + names_[index].c_str()); + return false; + } + if (proven.Is<core::Uint32>()) { + column.uint32s[row] = + static_cast<uint32_t>(sqlite3_column_int64(stmt, col)); + } else if (proven.Is<core::Int32>()) { + column.int32s[row] = static_cast<int32_t>(sqlite3_column_int64(stmt, col)); + } else if (proven.Is<Int64>()) { + column.ints[row] = sqlite3_column_int64(stmt, col); + } else if (proven.Is<Double>()) { + column.doubles[row] = sqlite3_column_double(stmt, col); + } else { + column.strings[row] = pool_->InternString( + reinterpret_cast<const char*>(sqlite3_column_text(stmt, col))); + } + column.validity.set(row); + return true; +} + +bool SqlScan::GetData(RowBatch& out, core::exec::OperatorState& state) const { + State& s = state.Cast<State>(); + if (s.done || !s.status.ok()) { + return false; + } + for (const std::shared_ptr<Column>& column : s.columns) { + if (column->validity.size() != 0) { + column->validity.ClearAllBits(); + } + } + sqlite3_stmt* stmt = s.statement->sqlite_stmt(); + uint32_t count = 0; + while (count < kMaxBatchRows && s.statement->Step()) { + for (uint32_t i = 0; i < s.columns.size(); ++i) { + if (!ReadValue(s, stmt, i, count)) { + return false; + } + } + ++count; + } + if (!s.statement->status().ok()) { + s.status = s.statement->status(); + return false; + } + s.done = count < kMaxBatchRows; + if (count == 0) { + return false; + } + + out.Reset(); + for (uint32_t i = 0; i < s.columns.size(); ++i) { + const std::shared_ptr<Column>& column = s.columns[i]; + if (!types_[i]) { + out.AddColumn(ColumnView::Variants(column->variants.data()), column); + continue; + } + const void* data = nullptr; + if (types_[i]->Is<core::Uint32>()) { + data = column->uint32s.data(); + } else if (types_[i]->Is<core::Int32>()) { + data = column->int32s.data(); + } else if (types_[i]->Is<Int64>()) { + data = column->ints.data(); + } else if (types_[i]->Is<Double>()) { + data = column->doubles.data(); + } else { + data = column->strings.data(); + } + out.AddColumn(ColumnView::Reference(*types_[i], data, &column->validity), + column); + } + out.Compose(RowSelection::Range(0), count); + out.SetCardinality(count); + return true; +} + +} // namespace perfetto::trace_processor::exec
diff --git a/src/trace_processor/perfetto_sql/exec/sql_scan.h b/src/trace_processor/perfetto_sql/exec/sql_scan.h new file mode 100644 index 0000000..1536d97 --- /dev/null +++ b/src/trace_processor/perfetto_sql/exec/sql_scan.h
@@ -0,0 +1,117 @@ +/* + * 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_PERFETTO_SQL_EXEC_SQL_SCAN_H_ +#define SRC_TRACE_PROCESSOR_PERFETTO_SQL_EXEC_SQL_SCAN_H_ + +#include <cstdint> +#include <memory> +#include <optional> +#include <string> +#include <vector> + +#include "perfetto/base/status.h" +#include "perfetto/ext/base/status_or.h" +#include "src/trace_processor/containers/string_pool.h" +#include "src/trace_processor/core/common/storage_types.h" +#include "src/trace_processor/core/exec/operator.h" +#include "src/trace_processor/core/exec/row_batch.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" +#include "src/trace_processor/sqlite/sql_source.h" +#include "src/trace_processor/sqlite/sqlite_connection.h" + +struct sqlite3_stmt; + +namespace perfetto::trace_processor::exec { + +// Reads a pipeline's rows from a SQL query. +// +// Promises nothing about the order the rows arrive in, because SQLite does +// not. +// +// Each column carries its type per row unless the query can be traced back to +// a dataframe column, which is the only way to establish a type. SQLite's +// declared types establish nothing: an INTEGER column holds text if something +// puts text in it. +class SqlScan : public core::exec::Source { + public: + // Without a catalog no column can be traced back, so all of them are + // variants. + static base::StatusOr<std::unique_ptr<SqlScan>> Create( + SqliteConnection*, + SqlSource, + StringPool*, + const lineage::Catalog* = nullptr); + ~SqlScan() override; + + // The query's columns, in the order a batch carries them. + const std::vector<std::string>& column_names() const { return names_; } + + // The type of column `i`, or nothing when the column carries a type per + // row. + std::optional<core::StorageType> column_type(uint32_t i) const { + return types_[i]; + } + + std::unique_ptr<core::exec::OperatorState> MakeState() const override; + bool GetData(core::exec::RowBatch& out, + core::exec::OperatorState&) const override; + void Rewind(core::exec::OperatorState&) const override; + base::Status status(const core::exec::OperatorState&) const override; + + private: + // One column's values. Only the buffer matching the column's type is ever + // filled. + struct Column { + std::vector<core::exec::Variant> variants; + std::vector<uint32_t> uint32s; + std::vector<int32_t> int32s; + std::vector<int64_t> ints; + std::vector<double> doubles; + std::vector<StringPool::Id> strings; + core::BitVector validity; + }; + + struct State : core::exec::OperatorState { + ~State() override; + std::optional<SqliteConnection::PreparedStatement> statement; + // Shared so a batch can keep the values alive. + std::vector<std::shared_ptr<Column>> columns; + bool done = false; + base::Status status = base::OkStatus(); + }; + + SqlScan(SqliteConnection*, + SqlSource, + std::vector<std::string>, + std::vector<std::optional<core::StorageType>>, + StringPool*); + void Prepare(State&) const; + + bool ReadValue(State&, sqlite3_stmt*, uint32_t index, uint32_t row) const; + + SqliteConnection* connection_; + SqlSource sql_; + std::vector<std::string> names_; + std::vector<std::optional<core::StorageType>> types_; + StringPool* pool_; +}; + +} // namespace perfetto::trace_processor::exec + +#endif // SRC_TRACE_PROCESSOR_PERFETTO_SQL_EXEC_SQL_SCAN_H_
diff --git a/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc b/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc new file mode 100644 index 0000000..1266e5b --- /dev/null +++ b/src/trace_processor/perfetto_sql/exec/sql_scan_unittest.cc
@@ -0,0 +1,457 @@ +/* + * 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/perfetto_sql/exec/sql_scan.h" + +#include <algorithm> +#include <cstdint> +#include <map> +#include <memory> +#include <optional> +#include <string> +#include <string_view> +#include <utility> +#include <vector> + +#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/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/variant.h" +#include "src/trace_processor/core/util/bit_vector.h" +#include "src/trace_processor/perfetto_sql/lineage/connection_catalog.h" +#include "src/trace_processor/sqlite/sql_source.h" +#include "src/trace_processor/sqlite/sqlite_connection.h" +#include "test/gtest_and_gmock.h" + +namespace perfetto::trace_processor::exec { +namespace { + +using core::BitVector; +using core::Double; +using core::Int64; +using core::StorageType; +using core::String; +using core::exec::ColumnView; +using core::exec::kMaxBatchRows; +using core::exec::RowBatch; +using core::exec::RowCursor; +using core::exec::Variant; + +// Drives a plan the way an executor does: creates the state, owns the batch. +class Execution { + public: + explicit Execution(const core::exec::Source& source) + : source_(source), state_(source.MakeState()) {} + + RowBatch* Next() { + return source_.GetData(batch_, *state_) ? &batch_ : nullptr; + } + void Rewind() { source_.Rewind(*state_); } + base::Status status() const { return source_.status(*state_); } + + private: + const core::exec::Source& source_; + std::unique_ptr<core::exec::OperatorState> state_; + RowBatch batch_; +}; + +std::vector<int64_t> ReadInts(const RowBatch& batch, uint32_t index) { + std::vector<int64_t> out; + for (const Variant& cell : + core::exec::test::ReadColumn<Variant>(batch, index)) { + out.push_back(cell.AsInt64()); + } + return out; +} + +struct TestColumn { + std::string name; + core::StorageType type; +}; + +TestColumn Typed(std::string name, core::StorageType type) { + return {std::move(name), type}; +} + +class TestCatalog : public lineage::Catalog { + public: + void Add(std::string name, std::vector<TestColumn> columns) { + dataframes_[std::move(name)] = std::move(columns); + } + + std::optional<lineage::analysis::LeafRelation> FindLeafRelation( + std::string_view name) const override { + auto dataframe = dataframes_.find(std::string(name)); + if (dataframe == dataframes_.end()) { + return std::nullopt; + } + lineage::analysis::LeafRelation relation; + relation.name = name; + for (const TestColumn& column : dataframe->second) { + relation.columns.push_back(column.name); + } + return relation; + } + + std::optional<std::string> FindViewSql(std::string_view) const override { + return std::nullopt; + } + + std::optional<core::StorageType> ColumnType( + const lineage::analysis::ColumnLineage& column) const override { + std::optional<core::StorageType> type; + for (const auto& origin : column.origins) { + auto dataframe = dataframes_.find(std::string(origin.relation_name)); + if (dataframe == dataframes_.end()) { + return std::nullopt; + } + auto found = + std::find_if(dataframe->second.begin(), dataframe->second.end(), + [&](const TestColumn& candidate) { + return candidate.name == origin.column_name; + }); + if (found == dataframe->second.end() || + (type && !(*type == found->type))) { + return std::nullopt; + } + type = found->type; + } + return type; + } + + private: + std::map<std::string, std::vector<TestColumn>> dataframes_; +}; + +class SqlScanTest : public ::testing::Test { + protected: + SqlScanTest() + : connection_(SqliteConnection::CreateConnectionToNewDatabase()) {} + + void Exec(const std::string& sql) { + auto statement = + connection_->PrepareStatement(SqlSource::FromExecuteQuery(sql)); + ASSERT_TRUE(statement.status().ok()) << statement.status().c_message(); + while (statement.Step()) { + } + ASSERT_TRUE(statement.status().ok()) << statement.status().c_message(); + } + + base::StatusOr<std::unique_ptr<SqlScan>> Scan( + const std::string& sql, + const lineage::Catalog* catalog = nullptr) { + return SqlScan::Create(connection_.get(), SqlSource::FromExecuteQuery(sql), + &pool_, catalog); + } + + StringPool pool_; + std::unique_ptr<SqliteConnection> connection_; +}; + +TEST_F(SqlScanTest, AQuerysColumnsAreKnownBeforeItsRows) { + auto scan = Scan("SELECT 1 AS a, 'x' AS b"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + EXPECT_THAT((*scan)->column_names(), testing::ElementsAre("a", "b")); +} + +TEST_F(SqlScanTest, AQuerysRowsArriveAsABatch) { + auto scan = Scan( + "SELECT a FROM (" + "SELECT 2 AS ord, 20 AS a UNION ALL " + "SELECT 3, 30 UNION ALL SELECT 1, 10) ORDER BY ord"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + EXPECT_EQ(batch->size(), 3u); + EXPECT_THAT(ReadInts(*batch, 0), testing::ElementsAre(10, 20, 30)); + EXPECT_EQ(run.Next(), nullptr); + EXPECT_TRUE(run.status().ok()); +} + +// One column holding three different types and a null, which SQLite allows. +TEST_F(SqlScanTest, OneColumnCanHoldMoreThanOneType) { + auto scan = Scan( + "SELECT value FROM (" + "SELECT 1 AS ord, 7 AS value UNION ALL SELECT 2, 1.5 " + "UNION ALL SELECT 3, 'hello' UNION ALL SELECT 4, NULL) ORDER BY ord"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + std::vector<Variant> cells = core::exec::test::ReadColumn<Variant>(*batch, 0); + ASSERT_EQ(cells.size(), 4u); + EXPECT_EQ(cells[0].AsInt64(), 7); + EXPECT_EQ(cells[1].AsDouble(), 1.5); + EXPECT_EQ(pool_.Get(cells[2].AsString()).ToStdString(), "hello"); + EXPECT_EQ(cells[3].type, Variant::Type::kNull); + EXPECT_TRUE(run.status().ok()); +} + +// A declared type is not binding in SQLite, so it is not trusted here. +TEST_F(SqlScanTest, ADeclaredTypeIsNotBelieved) { + 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(); + Execution run(**scan); + + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + std::vector<Variant> cells = core::exec::test::ReadColumn<Variant>(*batch, 0); + ASSERT_EQ(cells.size(), 2u); + EXPECT_EQ(cells[0].AsInt64(), 1); + EXPECT_EQ(pool_.Get(cells[1].AsString()).ToStdString(), "not a number"); + EXPECT_TRUE(run.status().ok()); +} + +TEST_F(SqlScanTest, AColumnWhichIsNeverAnythingIsAColumnOfNulls) { + auto scan = Scan("SELECT NULL AS a UNION ALL SELECT NULL"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + for (const Variant& cell : core::exec::test::ReadColumn<Variant>(*batch, 0)) { + EXPECT_EQ(cell.type, Variant::Type::kNull); + } +} + +TEST_F(SqlScanTest, MoreRowsThanFitInABatchArriveInSeveral) { + Exec( + "CREATE TABLE t AS WITH RECURSIVE r(x) AS (" + " SELECT 0 UNION ALL SELECT x + 1 FROM r WHERE x < 4999" + ") SELECT x FROM r"); + auto scan = Scan("SELECT x FROM t ORDER BY x"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + std::vector<int64_t> seen; + std::vector<uint32_t> sizes; + while (RowBatch* batch = run.Next()) { + sizes.push_back(batch->size()); + for (int64_t value : ReadInts(*batch, 0)) { + seen.push_back(value); + } + } + ASSERT_TRUE(run.status().ok()) << run.status().c_message(); + EXPECT_THAT(sizes, testing::ElementsAre(kMaxBatchRows, kMaxBatchRows, 904u)); + ASSERT_EQ(seen.size(), 5000u); + for (uint32_t i = 0; i < seen.size(); ++i) { + ASSERT_EQ(seen[i], int64_t{i}); + } +} + +TEST_F(SqlScanTest, AQueryWhichDoesNotRunIsReported) { + auto scan = Scan("SELECT * FROM not_a_table"); + EXPECT_FALSE(scan.ok()); +} + +TEST_F(SqlScanTest, AQueryWhichFailsPartWayThroughIsReported) { + auto scan = + Scan("SELECT 1 AS a UNION ALL SELECT abs(-9223372036854775807 - 1)"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + while (run.Next()) { + } + EXPECT_FALSE(run.status().ok()); +} + +TEST_F(SqlScanTest, ABlobIsReportedRatherThanCarried) { + auto scan = Scan("SELECT x'0102' AS a"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + EXPECT_EQ(run.Next(), nullptr); + EXPECT_FALSE(run.status().ok()); + EXPECT_THAT(run.status().message(), testing::HasSubstr("blob")); +} + +TEST_F(SqlScanTest, RewindClearsThePreviousExecutionError) { + auto scan = Scan("SELECT x'0102' AS a"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + ASSERT_EQ(run.Next(), nullptr); + ASSERT_FALSE(run.status().ok()); + + run.Rewind(); + EXPECT_TRUE(run.status().ok()); + EXPECT_EQ(run.Next(), nullptr); + EXPECT_FALSE(run.status().ok()); +} + +TEST_F(SqlScanTest, EachExecutionValidatesItsResultShape) { + Exec("CREATE TABLE t(a INTEGER)"); + Exec("INSERT INTO t VALUES(1)"); + auto scan = Scan("SELECT * FROM t"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Exec("ALTER TABLE t ADD COLUMN b TEXT"); + + Execution run(**scan); + EXPECT_EQ(run.Next(), nullptr); + EXPECT_FALSE(run.status().ok()); + EXPECT_THAT(run.status().message(), testing::HasSubstr("shape changed")); +} + +TEST_F(SqlScanTest, AScanCanBeRunAgain) { + auto scan = Scan("SELECT 1 AS a UNION ALL SELECT 2"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + + auto drain = [&] { + std::vector<int64_t> out; + while (RowBatch* batch = run.Next()) { + for (int64_t value : ReadInts(*batch, 0)) { + out.push_back(value); + } + } + return out; + }; + std::vector<int64_t> first = drain(); + run.Rewind(); + EXPECT_EQ(drain(), first); +} + +TEST_F(SqlScanTest, AQueryReachesARowCursor) { + auto scan = Scan( + "SELECT a FROM (SELECT 2 AS ord, 6 AS a UNION ALL " + "SELECT 3, 7 UNION ALL SELECT 1, 5) ORDER BY ord"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + + RowCursor cursor(**scan); + std::vector<int64_t> values; + for (bool more = cursor.Open(); more; more = cursor.Next()) { + values.push_back(cursor.Value<Variant>(0).AsInt64()); + } + EXPECT_THAT(values, testing::ElementsAre(5, 6, 7)); +} + +// 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) { + Exec("CREATE TABLE df(id INTEGER, name TEXT)"); + Exec("INSERT INTO df VALUES(7, 'hello'), (8, NULL)"); + TestCatalog catalog; + catalog.Add("df", {Typed("id", core::StorageType{core::Int64{}}), + Typed("name", core::StorageType{core::String{}})}); + + auto scan = Scan("SELECT id, name FROM df ORDER BY id", &catalog); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + ASSERT_TRUE((*scan)->column_type(0).has_value()); + EXPECT_TRUE((*scan)->column_type(0)->Is<core::Int64>()); + EXPECT_TRUE((*scan)->column_type(1)->Is<core::String>()); + + Execution run(**scan); + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + EXPECT_EQ(batch->column(0).kind(), ColumnView::Kind::kFlat); + const auto* ids = static_cast<const int64_t*>(batch->column(0).data()); + EXPECT_EQ(ids[0], 7); + EXPECT_EQ(ids[1], 8); + const BitVector* validity = batch->column(1).validity(); + ASSERT_NE(validity, nullptr); + EXPECT_TRUE(validity->is_set(0)); + EXPECT_FALSE(validity->is_set(1)); + std::vector<StringPool::Id> names = + core::exec::test::ReadColumn<StringPool::Id>(*batch, 1); + EXPECT_EQ(pool_.Get(names[0]).ToStdString(), "hello"); +} + +TEST_F(SqlScanTest, MixedNumericCompoundResultsStayVariants) { + Exec("CREATE TABLE ints(value INTEGER)"); + Exec("INSERT INTO ints VALUES(7)"); + Exec("CREATE TABLE doubles(value REAL)"); + Exec("INSERT INTO doubles VALUES(1.5)"); + TestCatalog catalog; + catalog.Add("ints", {Typed("value", StorageType{Int64{}})}); + catalog.Add("doubles", {Typed("value", StorageType{Double{}})}); + + auto scan = Scan( + "SELECT value FROM (" + "SELECT 1 AS ord, value FROM ints UNION ALL " + "SELECT 2, value FROM doubles) ORDER BY ord", + &catalog); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + EXPECT_FALSE((*scan)->column_type(0).has_value()); + Execution run(**scan); + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + std::vector<Variant> values = + core::exec::test::ReadColumn<Variant>(*batch, 0); + ASSERT_EQ(values.size(), 2u); + EXPECT_EQ(values[0].AsInt64(), 7); + EXPECT_EQ(values[1].AsDouble(), 1.5); +} + +// An expression cannot be traced back, so it stays a variant even when the +// column beside it does not. +TEST_F(SqlScanTest, OnlyTheColumnsWhichCanBeFollowedComeOutFlat) { + Exec("CREATE TABLE df(id INTEGER)"); + Exec("INSERT INTO df VALUES(7)"); + TestCatalog catalog; + catalog.Add("df", {Typed("id", core::StorageType{core::Int64{}})}); + + auto scan = Scan("SELECT id, id * 2 AS doubled FROM df", &catalog); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + EXPECT_TRUE((*scan)->column_type(0).has_value()); + EXPECT_FALSE((*scan)->column_type(1).has_value()); + + Execution run(**scan); + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + EXPECT_EQ(batch->column(0).kind(), ColumnView::Kind::kFlat); + EXPECT_EQ(batch->column(1).kind(), ColumnView::Kind::kVariant); +} + +// Without a catalog nothing can be traced back. +TEST_F(SqlScanTest, WithoutACatalogEveryColumnIsAVariant) { + Exec("CREATE TABLE df(id INTEGER)"); + auto scan = Scan("SELECT id FROM df"); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + EXPECT_FALSE((*scan)->column_type(0).has_value()); +} + +// A flat column's storage is readable at every row, so a reader which sums it +// without checking validity gets zero rather than a value left over from the +// previous batch. +TEST_F(SqlScanTest, ANullSlotOfAFlatColumnHoldsZero) { + Exec("CREATE TABLE df(id INTEGER)"); + Exec("INSERT INTO df VALUES(7), (NULL)"); + TestCatalog catalog; + catalog.Add("df", {Typed("id", core::StorageType{core::Int64{}})}); + + auto scan = Scan("SELECT id FROM df ORDER BY id IS NULL", &catalog); + ASSERT_TRUE(scan.ok()) << scan.status().c_message(); + Execution run(**scan); + RowBatch* batch = run.Next(); + ASSERT_NE(batch, nullptr); + ASSERT_EQ(batch->size(), 2u); + + const auto* ids = static_cast<const int64_t*>(batch->column(0).data()); + EXPECT_EQ(ids[0], 7); + EXPECT_EQ(ids[1], 0); + EXPECT_FALSE(batch->column(0).validity()->is_set(1)); +} + +} // namespace +} // namespace perfetto::trace_processor::exec
diff --git a/src/trace_processor/perfetto_sql/lineage/connection_catalog.cc b/src/trace_processor/perfetto_sql/lineage/connection_catalog.cc index 4076704..86d40c4 100644 --- a/src/trace_processor/perfetto_sql/lineage/connection_catalog.cc +++ b/src/trace_processor/perfetto_sql/lineage/connection_catalog.cc
@@ -48,6 +48,8 @@ } // namespace +Catalog::~Catalog() = default; + ConnectionCatalog::ConnectionCatalog(PerfettoSqlConnection* connection) : connection_(connection) {}
diff --git a/src/trace_processor/perfetto_sql/lineage/connection_catalog.h b/src/trace_processor/perfetto_sql/lineage/connection_catalog.h index 55eaaa6..b9480d9 100644 --- a/src/trace_processor/perfetto_sql/lineage/connection_catalog.h +++ b/src/trace_processor/perfetto_sql/lineage/connection_catalog.h
@@ -29,9 +29,18 @@ namespace analysis = ::perfetto::perfetto_sql::analysis; -// Adapts the dataframes and SQLite views of a live connection to the reusable -// PerfettoSQL relation analyzer. -class ConnectionCatalog final : public analysis::Catalog { +// Supplies live relation metadata and maps semantic origins back to dataframe +// storage. +class Catalog : public analysis::Catalog { + public: + ~Catalog() override; + + virtual std::optional<core::StorageType> ColumnType( + const analysis::ColumnLineage&) const = 0; +}; + +// Adapts one trace processor connection to semantic analysis. +class ConnectionCatalog final : public Catalog { public: explicit ConnectionCatalog(PerfettoSqlConnection*); @@ -42,7 +51,7 @@ // Maps all origins of a result column back to dataframe storage. Returns // nothing unless every origin has the same storage type. std::optional<core::StorageType> ColumnType( - const analysis::ColumnLineage&) const; + const analysis::ColumnLineage&) const override; private: PerfettoSqlConnection* connection_;