tp: preserve dataframe scan storage semantics

Direct scans must only borrow stable finalized storage and must retain nullable ID values instead of treating every ID column as a non-null sequence.

Author: LalitMaganti

src/trace_processor/core/exec/column_view.h

@@ -48,8 +48,9 @@ class ColumnView {
     ColumnView view;
     view.type_ = type;
     if (type.Is<Id>()) {
-      PERFETTO_DCHECK(data == nullptr && validity == nullptr);
+      PERFETTO_DCHECK(data == nullptr);
       view.kind_ = Kind::kSequence;
+      view.validity_ = validity;
       return view;
     }
     view.kind_ = Kind::kFlat;

src/trace_processor/core/exec/dataframe_scan.cc

@@ -19,9 +19,12 @@
 #include <algorithm>
 #include <cstdint>
 #include <memory>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
+#include "perfetto/base/logging.h"
+
 #include "src/trace_processor/containers/string_pool.h"
 #include "src/trace_processor/core/common/storage_types.h"
 #include "src/trace_processor/core/dataframe/dataframe.h"
@@ -70,7 +73,14 @@ class ExpanderImpl final : public DataframeScan::Expander {
     buffer_->validity.ClearAllBits();
     for (uint32_t row = 0; row < count; ++row) {
       if (bits_->is_set(from + row)) {
-        buffer_->values[row] = packed_[consumed_++];
+        if constexpr (std::is_same_v<T, uint32_t>) {
+          buffer_->values[row] =
+              packed_ ? packed_[consumed_] : static_cast<uint32_t>(consumed_);
+        } else {
+          PERFETTO_DCHECK(packed_);
+          buffer_->values[row] = packed_[consumed_];
+        }
+        ++consumed_;
         buffer_->validity.set(row);
       } else {
         // Written even for a null row, so the storage is readable everywhere.
@@ -134,9 +144,11 @@ void BuildColumn(const dataframe::Column& column,
 
 }  // namespace
 
-DataframeScan::DataframeScan(const dataframe::Dataframe* dataframe,
+DataframeScan::DataframeScan(const dataframe::Dataframe& dataframe,
                              std::vector<uint32_t> columns)
-    : dataframe_(dataframe), columns_(std::move(columns)) {}
+    : dataframe_(&dataframe), columns_(std::move(columns)) {
+  PERFETTO_CHECK(dataframe.finalized());
+}
 
 DataframeScan::~DataframeScan() = default;
 DataframeScan::State::~State() = default;
@@ -150,8 +162,18 @@ std::unique_ptr<OperatorState> DataframeScan::MakeState() const {
     uint32_t index = columns_[i];
     StorageType type = dataframe_->column_type(index);
     if (type.Is<Id>()) {
-      // No storage at all: the value is the row it sits at.
-      state->columns[i] = ColumnView::Reference(type, nullptr, nullptr);
+      const auto& nulls = dataframe_->column(index).null_storage;
+      if (nulls.nullability().Is<NonNull>()) {
+        state->columns[i] = ColumnView::Reference(type, nullptr, nullptr);
+      } else if (nulls.nullability().Is<DenseNull>()) {
+        state->columns[i] =
+            ColumnView::Reference(type, nullptr, &nulls.GetNullBitVector());
+      } else {
+        auto impl = std::make_unique<ExpanderImpl<uint32_t>>(
+            StorageType{Uint32{}}, nullptr, &nulls.GetNullBitVector());
+        state->owners[i] = impl->owner();
+        state->expanders[i] = std::move(impl);
+      }
       continue;
     }
     const dataframe::Column& column = dataframe_->column(index);

src/trace_processor/core/exec/dataframe_scan.h

@@ -41,9 +41,11 @@ namespace perfetto::trace_processor::core::exec {
 // execution, so a relation can be free for most of its columns and pay a
 // bounded amount for the rest. Nothing is materialised ahead of being asked
 // for, so a query which reads one batch and stops does one batch of work.
+//
+// The dataframe must be finalized and must outlive the scan.
 class DataframeScan : public Source {
  public:
-  DataframeScan(const dataframe::Dataframe*, std::vector<uint32_t> columns);
+  DataframeScan(const dataframe::Dataframe&, std::vector<uint32_t> columns);
   ~DataframeScan() override;
 
   std::unique_ptr<OperatorState> MakeState() const override;

src/trace_processor/core/exec/dataframe_scan_unittest.cc

@@ -18,7 +18,9 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
+#include <variant>
 #include <vector>
 
 #include "src/trace_processor/containers/string_pool.h"
@@ -28,12 +30,15 @@
 #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/test_utils.h"
 #include "test/gtest_and_gmock.h"
 
 namespace perfetto::trace_processor::core::exec {
 namespace {
 
 using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::Optional;
 
 // Bigger than an int32, so a column holding it is stored as an int64.
 constexpr int64_t kBig = int64_t{1} << 40;
@@ -62,23 +67,111 @@ class DataframeScanTest : public ::testing::Test {
     RowBatch batch;
     std::vector<int64_t> out;
     while (scan.GetData(batch, *state)) {
-      const ColumnView& view = batch.column(column);
-      const auto* data = static_cast<const int64_t*>(view.data());
-      for (uint32_t i = 0; i < batch.size(); ++i) {
-        out.push_back(data[view.selection().GetIndex(i)]);
-      }
+      std::vector<int64_t> values = test::ReadColumn<int64_t>(batch, column);
+      out.insert(out.end(), values.begin(), values.end());
     }
     return out;
   }
 
   StringPool pool_;
 };
 
+inline constexpr auto kAllTypes = dataframe::CreateTypedDataframeSpec(
+    {"id", "u32", "i32", "i64", "double", "string"},
+    dataframe::CreateTypedColumnSpec(Id{},
+                                     NonNull{},
+                                     IdSorted{},
+                                     NoDuplicates{}),
+    dataframe::CreateTypedColumnSpec(Uint32{}, NonNull{}, Unsorted{}),
+    dataframe::CreateTypedColumnSpec(Int32{}, NonNull{}, Unsorted{}),
+    dataframe::CreateTypedColumnSpec(Int64{}, NonNull{}, Unsorted{}),
+    dataframe::CreateTypedColumnSpec(Double{}, NonNull{}, Unsorted{}),
+    dataframe::CreateTypedColumnSpec(String{}, NonNull{}, Unsorted{}));
+
+inline constexpr auto kSparseId = dataframe::CreateTypedDataframeSpec(
+    {"id"},
+    dataframe::CreateTypedColumnSpec(Id{},
+                                     SparseNull{},
+                                     Unsorted{},
+                                     NoDuplicates{}));
+
+inline constexpr auto kDenseId = dataframe::CreateTypedDataframeSpec(
+    {"id"},
+    dataframe::CreateTypedColumnSpec(Id{},
+                                     DenseNull{},
+                                     Unsorted{},
+                                     NoDuplicates{}));
+
+TEST_F(DataframeScanTest, ReadsEveryStorageType) {
+  StringPool::Id first = pool_.InternString("first");
+  StringPool::Id second = pool_.InternString("second");
+  dataframe::Dataframe df =
+      dataframe::Dataframe::CreateFromTypedSpec(kAllTypes, &pool_);
+  df.InsertUnchecked(kAllTypes, std::monostate{}, uint32_t{7}, int32_t{-3},
+                     kBig, 1.5, first);
+  df.InsertUnchecked(kAllTypes, std::monostate{}, uint32_t{9}, int32_t{4},
+                     kBig + 1, 2.5, second);
+  df.Finalize();
+
+  DataframeScan scan(df, {0, 1, 2, 3, 4, 5});
+  std::unique_ptr<OperatorState> state = scan.MakeState();
+  RowBatch batch;
+  ASSERT_TRUE(scan.GetData(batch, *state));
+  EXPECT_TRUE(batch.column(0).type().Is<Id>());
+  EXPECT_TRUE(batch.column(1).type().Is<Uint32>());
+  EXPECT_TRUE(batch.column(2).type().Is<Int32>());
+  EXPECT_TRUE(batch.column(3).type().Is<Int64>());
+  EXPECT_TRUE(batch.column(4).type().Is<Double>());
+  EXPECT_TRUE(batch.column(5).type().Is<String>());
+  EXPECT_THAT(test::ReadColumn<uint32_t>(batch, 0), ElementsAre(0u, 1u));
+  EXPECT_THAT(test::ReadColumn<uint32_t>(batch, 1), ElementsAre(7u, 9u));
+  EXPECT_THAT(test::ReadColumn<int32_t>(batch, 2), ElementsAre(-3, 4));
+  EXPECT_THAT(test::ReadColumn<int64_t>(batch, 3), ElementsAre(kBig, kBig + 1));
+  EXPECT_THAT(test::ReadColumn<double>(batch, 4), ElementsAre(1.5, 2.5));
+  EXPECT_THAT(test::ReadColumn<StringPool::Id>(batch, 5),
+              ElementsAre(first, second));
+}
+
+TEST_F(DataframeScanTest, PreservesNullableIdSemantics) {
+  dataframe::Dataframe sparse =
+      dataframe::Dataframe::CreateFromTypedSpec(kSparseId, &pool_);
+  sparse.InsertUnchecked(kSparseId,
+                         std::optional<std::monostate>{std::monostate{}});
+  sparse.InsertUnchecked(kSparseId, std::optional<std::monostate>{});
+  sparse.InsertUnchecked(kSparseId,
+                         std::optional<std::monostate>{std::monostate{}});
+  sparse.Finalize();
+
+  DataframeScan sparse_scan(sparse, {0});
+  std::unique_ptr<OperatorState> sparse_state = sparse_scan.MakeState();
+  RowBatch batch;
+  ASSERT_TRUE(sparse_scan.GetData(batch, *sparse_state));
+  EXPECT_TRUE(batch.column(0).type().Is<Uint32>());
+  EXPECT_THAT(test::ReadNullableColumn<uint32_t>(batch, 0),
+              ElementsAre(Optional(0u), Eq(std::nullopt), Optional(1u)));
+
+  dataframe::Dataframe dense =
+      dataframe::Dataframe::CreateFromTypedSpec(kDenseId, &pool_);
+  dense.InsertUnchecked(kDenseId,
+                        std::optional<std::monostate>{std::monostate{}});
+  dense.InsertUnchecked(kDenseId, std::optional<std::monostate>{});
+  dense.InsertUnchecked(kDenseId,
+                        std::optional<std::monostate>{std::monostate{}});
+  dense.Finalize();
+
+  DataframeScan dense_scan(dense, {0});
+  std::unique_ptr<OperatorState> dense_state = dense_scan.MakeState();
+  ASSERT_TRUE(dense_scan.GetData(batch, *dense_state));
+  EXPECT_TRUE(batch.column(0).type().Is<Id>());
+  EXPECT_THAT(test::ReadNullableColumn<uint32_t>(batch, 0),
+              ElementsAre(Optional(0u), Eq(std::nullopt), Optional(2u)));
+}
+
 // The point of the operator: the batch reads the dataframe's own storage.
 TEST_F(DataframeScanTest, ReadsTheDataframesOwnStorage) {
   dataframe::Dataframe df =
       Build({kBig + 10, kBig + 20, kBig + 30}, {true, true, true});
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
 
   std::unique_ptr<OperatorState> state = scan.MakeState();
   RowBatch batch;
@@ -91,7 +184,7 @@ TEST_F(DataframeScanTest, ReadsTheDataframesOwnStorage) {
 TEST_F(DataframeScanTest, HandsBackEveryRow) {
   dataframe::Dataframe df =
       Build({kBig + 10, kBig + 20, kBig + 30}, {true, true, true});
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
   EXPECT_THAT(Drain(scan, 0), ElementsAre(kBig + 10, kBig + 20, kBig + 30));
 }
 
@@ -102,7 +195,7 @@ TEST_F(DataframeScanTest, SplitsIntoBatches) {
     values[i] = kBig + i;
   }
   dataframe::Dataframe df = Build(values, present);
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
   EXPECT_EQ(Drain(scan, 0).size(), values.size());
 }
 
@@ -111,7 +204,7 @@ TEST_F(DataframeScanTest, SplitsIntoBatches) {
 TEST_F(DataframeScanTest, AColumnWithoutASlotPerRowIsExpanded) {
   dataframe::Dataframe df =
       Build({kBig + 10, 0, kBig + 30}, {true, false, true});
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
 
   std::unique_ptr<OperatorState> state = scan.MakeState();
   RowBatch batch;
@@ -142,15 +235,15 @@ TEST_F(DataframeScanTest, AColumnWithoutASlotPerRowSpansBatches) {
     expected.push_back(values[i]);
   }
   dataframe::Dataframe df = Build(values, present);
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
   EXPECT_EQ(Drain(scan, 0), expected);
 }
 
 // Replaying has to wind the packed values back too, not just the row counter.
 TEST_F(DataframeScanTest, AColumnWithoutASlotPerRowIsReplayable) {
   dataframe::Dataframe df =
       Build({kBig + 10, 0, kBig + 30}, {true, false, true});
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
 
   std::unique_ptr<OperatorState> state = scan.MakeState();
   RowBatch batch;
@@ -167,7 +260,7 @@ TEST_F(DataframeScanTest, AColumnWithoutASlotPerRowIsReplayable) {
 
 TEST_F(DataframeScanTest, IsReplayable) {
   dataframe::Dataframe df = Build({kBig + 1, kBig + 2}, {true, true});
-  DataframeScan scan(&df, {0});
+  DataframeScan scan(df, {0});
 
   std::unique_ptr<OperatorState> state = scan.MakeState();
   RowBatch batch;