fix: support ZSTD-compressed IR and fix filteredLogEvents_ lifecycle. (#36)

Author: anlowee

velox/connectors/clp/search_lib/ir/ClpIrCursor.cpp

@@ -51,7 +51,7 @@ uint64_t ClpIrCursor::fetchNext(uint64_t numRows) {
 }
 
 size_t ClpIrCursor::getNumFilteredRows() const {
-  return irDeserializer_->get_ir_unit_handler().getFilteredLogEvents()->size();
+  return filteredLogEvents_->size();
 }
 
 VectorPtr ClpIrCursor::createVector(
@@ -72,8 +72,6 @@ ErrorCode ClpIrCursor::loadSplit() {
       ? NetworkAuthOption{.method = AuthMethod::None}
       : NetworkAuthOption{.method = AuthMethod::S3PresignedUrlV4};
 
-  auto irHandler = ClpIrUnitHandler{};
-
   auto projections = splitFieldsToNamesAndTypes();
   auto queryHandlerResult{QueryHandlerType::create(
       projectionResolutionCallback_,
@@ -89,15 +87,28 @@ ErrorCode ClpIrCursor::loadSplit() {
   auto irPath = Path{.source = inputSource_, .path = splitPath_};
   irReader_ = try_create_reader(irPath, networkAuthOption);
   if (nullptr == irReader_) {
+    VLOG(2) << "Failed to create IR reader";
+    return ErrorCode::InternalError;
+  }
+  irReaderZstdWrapper_ =
+      std::make_shared<::clp::streaming_compression::zstd::Decompressor>();
+  constexpr size_t cReaderBufferSize{64L * 1024L};
+  if (nullptr == irReaderZstdWrapper_) {
     VLOG(2) << "Failed to open kv-ir stream \"" << splitPath_
             << "\" for reading.";
     return ErrorCode::InternalError;
   }
+  irReaderZstdWrapper_->open(*irReader_, cReaderBufferSize);
 
+  filteredLogEvents_ = std::make_shared<
+      std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>();
+  auto irHandler = ClpIrUnitHandler{filteredLogEvents_};
   auto deserializerResult = ::clp::ffi::ir_stream::make_deserializer(
-      *irReader_, std::move(irHandler), std::move(queryHandler));
-  if (!deserializerResult) {
-    VLOG(2) << "Failed to create deserializer for deserialization.";
+      *irReaderZstdWrapper_, std::move(irHandler), std::move(queryHandler));
+  if (deserializerResult.has_error()) {
+    auto error = deserializerResult.error();
+    VLOG(2) << "Failed to create deserializer for deserialization, error: "
+            << error.message();
     return ErrorCode::InternalError;
   }
   irDeserializer_ = std::make_shared<
@@ -147,11 +158,11 @@ ClpIrCursor::splitFieldsToNamesAndTypes() const {
 
 ystdlib::error_handling::Result<void> ClpIrCursor::deserialize(
     uint64_t numRows) {
-  irDeserializer_->get_ir_unit_handler().clearFilteredLogEvents();
+  filteredLogEvents_->clear();
   uint64_t cnt{0};
   while (cnt < numRows) {
     auto deserializeResult =
-        irDeserializer_->deserialize_next_ir_unit(*irReader_);
+        irDeserializer_->deserialize_next_ir_unit(*irReaderZstdWrapper_);
     if (deserializeResult.has_error()) {
       auto error = deserializeResult.error();
       if (std::errc::result_out_of_range == error ||
@@ -201,7 +212,7 @@ VectorPtr ClpIrCursor::createVectorHelper(
       vectorType,
       vectorSize,
       std::make_unique<ClpIrVectorLoader>(
-          irDeserializer_->get_ir_unit_handler().getFilteredLogEvents(),
+          filteredLogEvents_,
           isResolved,
           std::move(projectedColumnNodeIds),
           projectedColumn.name,

velox/connectors/clp/search_lib/ir/ClpIrCursor.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "clp/streaming_compression/zstd/Decompressor.hpp"
 #include "ffi/ir_stream/Deserializer.hpp"
 #include "velox/connectors/clp/search_lib/BaseClpCursor.h"
 #include "velox/connectors/clp/search_lib/ir/ClpIrUnitHandler.h"
@@ -72,11 +73,16 @@ class ClpIrCursor final : public BaseClpCursor {
   };
   using QueryHandlerType = ::clp::ffi::ir_stream::search::QueryHandler<
       decltype(projectionResolutionCallback_)>;
+  std::shared_ptr<
+      std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>
+      filteredLogEvents_{nullptr};
   bool ignoreCase_;
   std::shared_ptr<
       ::clp::ffi::ir_stream::Deserializer<ClpIrUnitHandler, QueryHandlerType>>
-      irDeserializer_;
+      irDeserializer_{nullptr};
   std::shared_ptr<::clp::ReaderInterface> irReader_{nullptr};
+  std::shared_ptr<::clp::streaming_compression::zstd::Decompressor>
+      irReaderZstdWrapper_{nullptr};
   std::unordered_map<size_t, std::vector<::clp::ffi::SchemaTree::Node::id_t>>
       projectedColumnIdxNodeIdsMap_;
   size_t readerIndex_{0};

velox/connectors/clp/search_lib/ir/ClpIrUnitHandler.h

@@ -28,9 +28,13 @@ namespace facebook::velox::connector::clp::search_lib {
 
 class ClpIrUnitHandler {
  public:
-  ClpIrUnitHandler() {
-    filteredLogEvents_ = std::make_shared<
-        std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>();
+  ClpIrUnitHandler(
+      std::shared_ptr<
+          std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>
+          filteredLogEvents)
+      : filteredLogEvents_(filteredLogEvents) {
+    VELOX_CHECK_NOT_NULL(
+        filteredLogEvents_, "filteredLogEvents cannot be null");
   }
 
   // Destructor
@@ -67,16 +71,6 @@ class ClpIrUnitHandler {
     return ::clp::ffi::ir_stream::IRErrorCode::IRErrorCode_Success;
   }
 
-  std::shared_ptr<
-      const std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>
-  getFilteredLogEvents() const {
-    return filteredLogEvents_;
-  }
-
-  void clearFilteredLogEvents() {
-    filteredLogEvents_->clear();
-  }
-
  private:
   std::shared_ptr<
       std::vector<std::unique_ptr<::clp::ffi::KeyValuePairLogEvent>>>

velox/connectors/clp/tests/ClpConnectorTest.cpp

@@ -153,7 +153,7 @@ TEST_F(ClpConnectorTest, test1NoPushdown) {
   auto irOutput = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_1_ir.clps"),
+          getExampleFilePath("test_1_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   test::assertEqualVectors(irExpected, irOutput);
@@ -201,7 +201,7 @@ TEST_F(ClpConnectorTest, test1Pushdown) {
   auto irOutput = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_1_ir.clps"),
+          getExampleFilePath("test_1_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   test::assertEqualVectors(expected, irOutput);
@@ -260,7 +260,7 @@ TEST_F(ClpConnectorTest, test2NoPushdown) {
   auto irOutput = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_2_ir.clps"),
+          getExampleFilePath("test_2_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   test::assertEqualVectors(expected, irOutput);
@@ -317,7 +317,7 @@ TEST_F(ClpConnectorTest, test2Pushdown) {
   auto irOutput = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_2_ir.clps"),
+          getExampleFilePath("test_2_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   test::assertEqualVectors(expected, irOutput);
@@ -378,7 +378,7 @@ TEST_F(ClpConnectorTest, test2Hybrid) {
   auto irOutput = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_2_ir.clps"),
+          getExampleFilePath("test_2_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   test::assertEqualVectors(expected, irOutput);
@@ -433,7 +433,7 @@ TEST_F(ClpConnectorTest, test4IrTimestampNoPushdown) {
   auto output = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_4_ir.clps"),
+          getExampleFilePath("test_4_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   auto expected = makeRowVector({
@@ -445,6 +445,8 @@ TEST_F(ClpConnectorTest, test4IrTimestampNoPushdown) {
 }
 
 TEST_F(ClpConnectorTest, test4IrTimestampPushdown) {
+  // Only the second event meet the condition, the first event is a date string
+  // which is not supported yet so the value will be NULL.
   const std::shared_ptr<std::string> kqlQuery =
       std::make_shared<std::string>("(timestamp < 1756003005000000)");
   auto plan = PlanBuilder(pool_.get())
@@ -462,7 +464,7 @@ TEST_F(ClpConnectorTest, test4IrTimestampPushdown) {
   auto output = getResults(
       plan,
       {makeClpSplit(
-          getExampleFilePath("test_4_ir.clps"),
+          getExampleFilePath("test_4_ir.clp.zst"),
           ClpConnectorSplit::SplitType::kIr,
           kqlQuery)});
   auto expected = makeRowVector({