feat: Add RowRange and ChunkHeader utilities (facebookincubator#586)

Summary:
Pull Request resolved: facebookincubator#586

CONTEXT: NimbleIndexProjector needs lightweight value types for representing
row ranges and parsing chunk headers from raw tablet stream data.

WHAT:
- Add RowRange struct for representing [offset, length) row ranges with
  merge/overlap detection and fmt::formatter support.
- Add ChunkHeader for parsing chunk headers from raw tablet stream bytes,
  providing header size, data size, and total chunk size accessors.

Reviewed By: HuamengJiang, kewang1024, tanjialiang

Differential Revision: D97251759

fbshipit-source-id: 665f09aaf1a23a66fae3a7187891d84739a5f69a

Author: xiaoxmeng

dwio/nimble/common/ChunkHeader.h

@@ -0,0 +1,49 @@
+/*
+ * 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.
+ */
+#pragma once
+
+#include <cstdint>
+
+#include "dwio/nimble/common/EncodingPrimitives.h"
+#include "dwio/nimble/common/Types.h"
+
+namespace facebook::nimble {
+
+/// Chunk header: 4 bytes for compressed chunk length + 1 byte for compression
+/// type. This constant must remain 5 to maintain backward compatibility with
+/// existing Nimble files.
+constexpr int kChunkHeaderSize = 5;
+
+struct ChunkHeader {
+  uint32_t length;
+  CompressionType compressionType;
+};
+
+/// Reads a chunk header from 'pos', advancing 'pos' by kChunkHeaderSize bytes.
+inline ChunkHeader readChunkHeader(const char*& pos) {
+  return {
+      encoding::readUint32(pos),
+      static_cast<CompressionType>(encoding::readChar(pos))};
+}
+
+/// Writes a chunk header to 'pos', advancing 'pos' by kChunkHeaderSize bytes.
+inline void
+writeChunkHeader(uint32_t length, CompressionType compressionType, char*& pos) {
+  encoding::writeUint32(length, pos);
+  encoding::write(compressionType, pos);
+}
+
+} // namespace facebook::nimble

dwio/nimble/common/Constants.h

@@ -45,6 +45,4 @@ constexpr uint64_t kChunkingWriterMaxChunkSize{20 << 20}; // 20MB
 /// Used in place of kChunkingWriterMaxChunkSize for tables with large schemas.
 constexpr uint64_t kChunkingWriterWideSchemaMaxChunkSize{2 << 20}; // 2MB
 
-// Chunk header size: 4 bytes for length + 1 byte for compression type.
-constexpr int kChunkHeaderSize = 5;
 } // namespace facebook::nimble

dwio/nimble/common/tests/CMakeLists.txt

@@ -24,6 +24,7 @@ add_executable(
   nimble_common_tests
   BitEncoderTests.cpp
   BitsTests.cpp
+  ChunkHeaderTest.cpp
   ConstantsTest.cpp
   ExceptionHelperTest.cpp
   ExceptionTests.cpp

dwio/nimble/common/tests/ChunkHeaderTest.cpp

@@ -0,0 +1,63 @@
+/*
+ * 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/ChunkHeader.h"
+
+namespace facebook::nimble::test {
+
+TEST(ChunkHeaderTest, size) {
+  EXPECT_EQ(kChunkHeaderSize, 5);
+  EXPECT_EQ(kChunkHeaderSize, sizeof(uint32_t) + sizeof(CompressionType));
+}
+
+TEST(ChunkHeaderTest, writeAndRead) {
+  char buffer[kChunkHeaderSize];
+  auto* writePos = buffer;
+  writeChunkHeader(1'234, CompressionType::Zstd, writePos);
+  EXPECT_EQ(writePos - buffer, kChunkHeaderSize);
+
+  const char* readPos = buffer;
+  const auto header = readChunkHeader(readPos);
+  EXPECT_EQ(readPos - buffer, kChunkHeaderSize);
+  EXPECT_EQ(header.length, 1'234);
+  EXPECT_EQ(header.compressionType, CompressionType::Zstd);
+}
+
+TEST(ChunkHeaderTest, uncompressed) {
+  char buffer[kChunkHeaderSize];
+  auto* writePos = buffer;
+  writeChunkHeader(42, CompressionType::Uncompressed, writePos);
+
+  const char* readPos = buffer;
+  const auto header = readChunkHeader(readPos);
+  EXPECT_EQ(header.length, 42);
+  EXPECT_EQ(header.compressionType, CompressionType::Uncompressed);
+}
+
+TEST(ChunkHeaderTest, structuredBinding) {
+  char buffer[kChunkHeaderSize];
+  auto* writePos = buffer;
+  writeChunkHeader(100'000, CompressionType::Zstd, writePos);
+
+  const char* readPos = buffer;
+  const auto [length, compressionType] = readChunkHeader(readPos);
+  EXPECT_EQ(length, 100'000);
+  EXPECT_EQ(compressionType, CompressionType::Zstd);
+}
+
+} // namespace facebook::nimble::test

dwio/nimble/common/tests/ConstantsTest.cpp

@@ -14,20 +14,8 @@
  * limitations under the License.
  */
 
-#include <gtest/gtest.h>
-
-#include "dwio/nimble/common/Constants.h"
-
 namespace facebook::nimble::test {
 
-TEST(ConstantsTest, kChunkHeaderSize) {
-  // Chunk header format:
-  //   - 4 bytes (uint32_t) for compressed chunk length
-  //   - 1 byte for compression type
-  // This constant must remain 5 to maintain backward compatibility with
-  // existing Nimble files.
-  EXPECT_EQ(kChunkHeaderSize, 5);
-  EXPECT_EQ(kChunkHeaderSize, sizeof(uint32_t) + 1);
-}
+// kChunkHeaderSize tests moved to ChunkHeaderTest.cpp.
 
 } // namespace facebook::nimble::test

dwio/nimble/velox/CMakeLists.txt

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 add_library(nimble_velox_common SchemaUtils.cpp)
-target_link_libraries(nimble_velox_common nimble_common velox_type)
+target_link_libraries(nimble_velox_common nimble_common velox_type Folly::folly fmt::fmt)
 
 add_library(nimble_velox_schema SchemaTypes.cpp)
 target_link_libraries(nimble_velox_schema nimble_common Folly::folly)

dwio/nimble/velox/ChunkedStream.cpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #include "dwio/nimble/velox/ChunkedStream.h"
-#include "dwio/nimble/common/EncodingPrimitives.h"
+#include "dwio/nimble/common/ChunkHeader.h"
 #include "dwio/nimble/common/Exceptions.h"
 #include "dwio/nimble/tablet/Compression.h"
 #include "folly/io/Cursor.h"
@@ -37,11 +37,10 @@ std::string_view InMemoryChunkedStream::nextChunk() {
   ensureLoaded();
   uncompressed_.clear();
   NIMBLE_CHECK_LE(
-      sizeof(uint32_t) + sizeof(char),
+      kChunkHeaderSize,
       stream_.size() - (pos_ - stream_.data()),
       "Read beyond end of stream");
-  auto length = encoding::readUint32(pos_);
-  auto compressionType = static_cast<CompressionType>(encoding::readChar(pos_));
+  const auto [length, compressionType] = readChunkHeader(pos_);
   NIMBLE_CHECK_LE(
       length,
       stream_.size() - (pos_ - stream_.data()),
@@ -70,10 +69,10 @@ std::string_view InMemoryChunkedStream::nextChunk() {
 CompressionType InMemoryChunkedStream::peekCompressionType() {
   ensureLoaded();
   NIMBLE_CHECK_LE(
-      sizeof(uint32_t) + sizeof(char),
+      kChunkHeaderSize,
       stream_.size() - (pos_ - stream_.data()),
       "Read beyond end of stream");
-  auto pos = pos_ + sizeof(uint32_t);
+  auto* pos = pos_ + sizeof(uint32_t);
   return static_cast<CompressionType>(encoding::readChar(pos));
 }
 

dwio/nimble/velox/ChunkedStreamWriter.cpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #include "dwio/nimble/velox/ChunkedStreamWriter.h"
-#include "dwio/nimble/common/EncodingPrimitives.h"
+#include "dwio/nimble/common/ChunkHeader.h"
 #include "dwio/nimble/common/Exceptions.h"
 #include "dwio/nimble/tablet/Compression.h"
 
@@ -34,25 +34,22 @@ ChunkedStreamWriter::ChunkedStreamWriter(
 
 std::vector<std::string_view> ChunkedStreamWriter::encode(
     std::string_view chunk) {
-  constexpr uint8_t headerSize = sizeof(uint32_t) + sizeof(CompressionType);
-  auto* header = buffer_.reserve(headerSize);
+  auto* header = buffer_.reserve(kChunkHeaderSize);
   auto* pos = header;
 
   if (compressionParams_.type == CompressionType::Zstd) {
     auto compressed = ZstdCompression::compress(
         buffer_.getMemoryPool(), chunk, compressionParams_.zstdLevel);
     if (compressed.has_value()) {
-      encoding::writeUint32(compressed->size(), pos);
-      encoding::write(CompressionType::Zstd, pos);
+      writeChunkHeader(compressed->size(), CompressionType::Zstd, pos);
       return {
-          {header, headerSize},
+          {header, kChunkHeaderSize},
           buffer_.takeOwnership(compressed->releaseOwnership())};
     }
   }
 
-  encoding::writeUint32(chunk.size(), pos);
-  encoding::write(CompressionType::Uncompressed, pos);
-  return {{header, headerSize}, chunk};
+  writeChunkHeader(chunk.size(), CompressionType::Uncompressed, pos);
+  return {{header, kChunkHeaderSize}, chunk};
 }
 
 } // namespace facebook::nimble

dwio/nimble/velox/RowRange.h

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <fmt/format.h>
+#include <folly/hash/Hash.h>
+
+#include "velox/vector/TypeAliases.h"
+
+namespace facebook::nimble {
+
+/// Row range within a stripe [startRow, endRow).
+struct RowRange {
+  velox::vector_size_t startRow{0};
+  velox::vector_size_t endRow{0};
+
+  RowRange() = default;
+  RowRange(velox::vector_size_t start, velox::vector_size_t end)
+      : startRow(start), endRow(end) {}
+
+  inline velox::vector_size_t numRows() const {
+    return endRow - startRow;
+  }
+
+  inline bool empty() const {
+    return startRow >= endRow;
+  }
+
+  inline bool contains(const RowRange& other) const {
+    return startRow <= other.startRow && other.endRow <= endRow;
+  }
+
+  inline bool operator==(const RowRange& other) const {
+    return startRow == other.startRow && endRow == other.endRow;
+  }
+
+  inline std::string toString() const {
+    return fmt::format("[{}, {})", startRow, endRow);
+  }
+};
+
+struct RowRangeHash {
+  inline size_t operator()(const RowRange& range) const {
+    return folly::hash::hash_combine(range.startRow, range.endRow);
+  }
+};
+
+} // namespace facebook::nimble

dwio/nimble/velox/selective/ChunkedDecoder.cpp

@@ -16,8 +16,7 @@
 
 #include "dwio/nimble/velox/selective/ChunkedDecoder.h"
 
-#include "dwio/nimble/common/Constants.h"
-#include "dwio/nimble/common/EncodingPrimitives.h"
+#include "dwio/nimble/common/ChunkHeader.h"
 #include "dwio/nimble/common/Types.h"
 #include "dwio/nimble/encodings/EncodingFactory.h"
 #include "velox/common/testutil/TestValue.h"
@@ -33,9 +32,7 @@ using namespace facebook::velox;
 void ChunkedDecoder::loadNextChunk() {
   auto ret = ensureInput(kChunkHeaderSize);
   NIMBLE_CHECK(ret, "Failed to read chunk header");
-  auto length = encoding::readUint32(inputData_);
-  const auto compressionType =
-      static_cast<CompressionType>(encoding::readChar(inputData_));
+  const auto [length, compressionType] = readChunkHeader(inputData_);
   inputSize_ -= kChunkHeaderSize;
   ret = ensureInput(length);
   NIMBLE_CHECK(ret);