Simple TabletReaderTest setup (facebookincubator#452)

Summary:

Adding this simple standalone tablet reader test which can be useful for later experimentations

Differential Revision: D91945755

Author: srsuryadev

dwio/nimble/tablet/tests/CMakeLists.txt

@@ -31,6 +31,8 @@ target_link_libraries(
 
 add_executable(nimble_tablet_tests TabletTest.cpp)
 
+add_executable(nimble_tablet_reader_tests TabletReaderTest.cpp)
+
 add_test(nimble_tablet_tests nimble_tablet_tests)
 
 target_link_libraries(

dwio/nimble/tablet/tests/TabletReaderTest.cpp

@@ -0,0 +1,187 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * 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 <gtest/gtest.h>
+
+#include "dwio/nimble/common/Buffer.h"
+#include "dwio/nimble/common/tests/TestUtils.h"
+#include "dwio/nimble/index/tests/TabletIndexTestUtils.h"
+#include "dwio/nimble/tablet/TabletReader.h"
+#include "dwio/nimble/tablet/TabletWriter.h"
+#include "dwio/nimble/tablet/tests/TabletTestUtils.h"
+#include "velox/common/file/File.h"
+#include "velox/common/memory/Memory.h"
+
+using namespace facebook;
+
+namespace {
+
+class TabletReaderTest : public ::testing::Test {
+ protected:
+  using ChunkSpec = nimble::index::test::ChunkSpec;
+  using KeyChunkSpec = nimble::index::test::KeyChunkSpec;
+  using StreamSpec = nimble::index::test::StreamSpec;
+
+  static void SetUpTestCase() {
+    // Initialize a fresh global MemoryManager instance for tests to ensure
+    // deterministic behavior and isolation across unit tests.
+    velox::memory::MemoryManager::testingSetInstance({});
+  }
+
+  void SetUp() override {
+    // Per-test setup hook. No-op for now, reserved for future initialization.
+  }
+
+  // Test case describing the expected outcome of a key lookup against the
+  // tablet index. If expectedStripeIndex is std::nullopt, no matching stripe
+  // should be found for the given key.
+  struct LookupTestCase {
+    std::string key; // lookup key (serialized)
+    std::optional<uint32_t>
+        expectedStripeIndex; // expected stripe index or nullopt if no match
+  };
+
+  // Test case for verifying StripeIndexGroup::lookupChunk. This validates both
+  // which stripe should contain the key and the global file row ID where the
+  // key-range begins (relative to the start of the file across all stripes).
+  struct KeyLookupTestCase {
+    std::string key; // key to look up (matches key encoding in KeyStream)
+    uint32_t expectedStripeIndex; // stripe index expected to contain the key
+    uint32_t
+        expectedFileRowId; // absolute row id in the file where the chunk starts
+  };
+
+  // Create a single data Stream from a specification using shared test
+  // utilities. The spec includes stream offset and chunk metadata (rowCount,
+  // size).
+  static nimble::Stream createStream(
+      nimble::Buffer& buffer,
+      const StreamSpec& spec) {
+    return nimble::index::test::createStream(buffer, spec);
+  }
+
+  // Create multiple data Streams from a list of specifications. Reserves
+  // capacity to avoid reallocations and uses createStream(...) for each entry.
+  static std::vector<nimble::Stream> createStreams(
+      nimble::Buffer& buffer,
+      const std::vector<StreamSpec>& specs) {
+    std::vector<nimble::Stream> streams;
+    streams.reserve(specs.size());
+    for (const auto& spec : specs) {
+      streams.push_back(createStream(buffer, spec));
+    }
+    return streams;
+  }
+
+  // Create a KeyStream from chunk specifications. Each KeyChunkSpec defines the
+  // row count and [minKey, maxKey] range for a chunk used by the index.
+  static nimble::KeyStream createKeyStream(
+      nimble::Buffer& buffer,
+      const std::vector<KeyChunkSpec>& chunkSpecs) {
+    return nimble::index::test::createKeyStream(buffer, chunkSpecs);
+  }
+
+  // Memory pools used by writer/reader and test buffers. rootPool_ owns pool_.
+  std::shared_ptr<velox::memory::MemoryPool> rootPool_{
+      velox::memory::memoryManager()->addRootPool("TabletTest")};
+  std::shared_ptr<velox::memory::MemoryPool> pool_{
+      rootPool_->addLeafChild("TabletTest")};
+};
+
+TEST_F(TabletReaderTest, SimpleStreamReader) {
+  std::string file;
+  // Write-side in-memory file; backed by a std::string to capture the
+  // serialized tablet.
+  velox::InMemoryWriteFile writeFile(&file);
+
+  // Configure an index over two columns with ascending order and nulls-first.
+  // enforceKeyOrder ensures input rows respect the sorted key order per stripe.
+  nimble::TabletIndexConfig indexConfig{
+      .columns = {"col1", "col2"},
+      .sortOrders = {velox::core::kAscNullsFirst, velox::core::kAscNullsFirst},
+      .enforceKeyOrder = true,
+  };
+
+  // Create a tablet writer. metadataFlushThreshold=0 forces metadata to flush
+  // after each stripe for easier testing. Stream deduplication disabled to
+  // keep stream layout predictable.
+  auto tabletWriter = nimble::TabletWriter::create(
+      &writeFile,
+      *pool_,
+      {
+          // Force metadata flush after every stripe to simplify verification.
+          .metadataFlushThreshold = 0,
+          // Keep each stream distinct (no dedup) for deterministic
+          // sizes/offsets.
+          .streamDeduplicationEnabled = false,
+          .indexConfig = indexConfig,
+      });
+
+  // Shared buffer used by test helpers to assemble streams/key streams.
+  nimble::Buffer buffer{*pool_};
+
+  // Write stripe 0 (Stripe Group 0)
+  // - Total rows in stripe: 100
+  // - Stream 0 has 2 chunks, each with 50 rows (sizes 10 and 12 bytes)
+
+  // - KeyStream encodes two key chunks: [aaa, bbb] for first 50 rows,
+  //   and [bbb, ccc] for the next 50 rows.
+  {
+    auto streams = createStreams(
+        buffer,
+        {
+            {.offset = 0,
+             .chunks =
+                 {
+                     {.rowCount = 50, .size = 10},
+                     {.rowCount = 50, .size = 12},
+                 }},
+        });
+
+    auto keyStream = createKeyStream(
+        buffer,
+        {
+            // First 50 rows span keys ["aaa", "bbb"]; second 50 span ["bbb",
+            // "ccc"].
+            {50, "aaa", "bbb"},
+            {50, "bbb", "ccc"},
+        });
+
+    // Persist stripe 0: write its streams and key metadata into the tablet.
+    tabletWriter->writeStripe(100, std::move(streams), std::move(keyStream));
+  }
+
+  // Finalize writer and seal the in-memory file contents.
+  tabletWriter->close();
+  writeFile.close();
+
+  // Read back the tablet from the in-memory buffer to validate invariants.
+  nimble::testing::InMemoryTrackableReadFile readFile(file, false);
+  auto tablet = nimble::TabletReader::create(&readFile, *pool_);
+
+  // Helper exposes internal details for test assertions (e.g., stripe groups).
+  nimble::test::TabletReaderTestHelper tabletHelper(tablet.get());
+  EXPECT_EQ(tabletHelper.numStripeGroups(), 1)
+      << "Expected a single stripe group";
+
+  // Basic sanity checks on tablet properties.
+  EXPECT_EQ(tablet->stripeCount(), 1) << "One stripe should have been written";
+  EXPECT_EQ(tablet->tabletRowCount(), 100)
+      << "Total row count across all stripes";
+  EXPECT_EQ(tablet->stripeRowCount(0), 100)
+      << "Stripe 0 should contain 100 rows";
+}
+
+} // namespace