feat: Implement IR cursor to allow searches over IR streams.

CMake/resolve_dependency_modules/clp.cmake

@@ -16,7 +16,7 @@ include_guard(GLOBAL)
 FetchContent_Declare(
   clp
   GIT_REPOSITORY https://github.com/y-scope/clp.git
-  GIT_TAG 19cd534e629d746395efc64343a60f768b0c9a2d)
+  GIT_TAG bfd4f60ffe9c5d69618cc8416ec6729c76ee9862)
 
 set(CLP_BUILD_CLP_REGEX_UTILS
     OFF

CMake/resolve_dependency_modules/ystdlib_cpp.cmake

@@ -21,3 +21,4 @@ FetchContent_Declare(
 FetchContent_Populate(ystdlib_cpp)
 
 set(CLP_YSTDLIB_SOURCE_DIRECTORY "${ystdlib_cpp_SOURCE_DIR}")
+include_directories(${ystdlib_cpp_SOURCE_DIR}/src)

velox/connectors/clp/ClpDataSource.cpp

@@ -22,6 +22,7 @@
 #include "velox/connectors/clp/ClpTableHandle.h"
 #include "velox/connectors/clp/search_lib/ClpS3AuthProviderBase.h"
 #include "velox/connectors/clp/search_lib/archive/ClpArchiveCursor.h"
+#include "velox/connectors/clp/search_lib/ir/ClpIrCursor.h"
 #include "velox/vector/FlatVector.h"
 
 namespace facebook::velox::connector::clp {
@@ -116,6 +117,9 @@ void ClpDataSource::addSplit(std::shared_ptr<ConnectorSplit> split) {
   if (ClpConnectorSplit::SplitType::kArchive == clpSplit->type_) {
     cursor_ =
         std::make_unique<search_lib::ClpArchiveCursor>(inputSource, splitPath);
+  } else if (ClpConnectorSplit::SplitType::kIr == clpSplit->type_) {
+    cursor_ =
+        std::make_unique<search_lib::ClpIrCursor>(inputSource, splitPath, true);
   } else {
     VELOX_UNSUPPORTED(
         "Unsupported split type: {}", static_cast<int>(clpSplit->type_));

velox/connectors/clp/search_lib/CMakeLists.txt

@@ -22,7 +22,11 @@ velox_add_library(
   ClpS3AuthProviderBase.h)
 
 add_subdirectory(archive)
+add_subdirectory(ir)
 
-velox_link_libraries(clp-s-search PUBLIC clp-s-archive-search)
+velox_link_libraries(
+  clp-s-search
+  PUBLIC clp-s-archive-search
+  PUBLIC clp-s-ir-search)
 
 target_compile_features(clp-s-search PRIVATE cxx_std_20)

velox/connectors/clp/search_lib/ir/CMakeLists.txt

@@ -0,0 +1,32 @@
+# Copyright (c) Facebook, Inc. and its 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.
+velox_add_library(
+  clp-s-ir-search
+  STATIC
+  ClpIrCursor.cpp
+  ClpIrCursor.h
+  ClpIrUnitHandler.cpp
+  ClpIrUnitHandler.h
+  ClpIrVectorLoader.cpp
+  ClpIrVectorLoader.h)
+
+velox_link_libraries(
+  clp-s-ir-search
+  PUBLIC clp_s::archive_reader
+  PRIVATE
+    clp_s::clp_dependencies
+    clp_s::io
+    clp_s::search
+    clp_s::search::kql
+    velox_vector)

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

@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) Facebook, Inc. and its 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 "clp_s/ColumnReader.hpp"
+#include "clp_s/InputConfig.hpp"
+
+#include "ffi/ir_stream/search/QueryHandler.hpp"
+#include "velox/connectors/clp/search_lib/ir/ClpIrCursor.h"
+#include "velox/connectors/clp/search_lib/ir/ClpIrVectorLoader.h"
+
+using namespace clp_s;
+
+namespace facebook::velox::connector::clp::search_lib {
+
+uint64_t ClpIrCursor::fetchNext(uint64_t numRows) {
+  readerIndex_ = 0;
+  if (ErrorCode::Success != errorCode_) {
+    return 0;
+  }
+
+  if (false == currentSplitLoaded_) {
+    errorCode_ = loadSplit();
+    if (ErrorCode::Success != errorCode_) {
+      return 0;
+    }
+  }
+
+  auto deserializeResult = deserialize(numRows);
+  if (deserializeResult.has_error()) {
+    auto error = deserializeResult.error();
+    VELOX_FAIL(
+        "IR file {} might be broken, failed to deserialize. {}: {}",
+        this->splitPath_,
+        error.category().name(),
+        error.message());
+  }
+  return irDeserializer_->get_num_log_events_deserialized();
+}
+
+size_t ClpIrCursor::getNumFilteredRows() const {
+  return irDeserializer_->get_ir_unit_handler().getFilteredLogEvents()->size();
+}
+
+VectorPtr ClpIrCursor::createVector(
+    memory::MemoryPool* pool,
+    const TypePtr& vectorType,
+    size_t vectorSize) {
+  VELOX_CHECK_EQ(
+      projectedColumnNameNodeIdMap_.size(),
+      outputColumns_.size(),
+      "Projected columns size {} does not match fields size {}",
+      projectedColumnNameNodeIdMap_.size(),
+      outputColumns_.size());
+  return createVectorHelper(pool, vectorType, vectorSize);
+}
+
+ErrorCode ClpIrCursor::loadSplit() {
+  auto networkAuthOption = inputSource_ == InputSource::Filesystem
+      ? NetworkAuthOption{.method = AuthMethod::None}
+      : NetworkAuthOption{.method = AuthMethod::S3PresignedUrlV4};
+
+  auto irHandler = ClpIrUnitHandler{};
+
+  auto projections = splitFieldsToNamesAndTypes();
+  auto queryHandlerResult{QueryHandlerType::create(
+      projectionResolutionCallback_,
+      std::move(expr_),
+      projections,
+      ignoreCase_)};
+  if (!queryHandlerResult) {
+    VLOG(2) << "Failed to create query handler for deserialization.";
+    return ErrorCode::InternalError;
+  }
+  auto queryHandler = std::move(queryHandlerResult).value();
+
+  auto irPath = Path{.source = inputSource_, .path = splitPath_};
+  irReader_ = try_create_reader(irPath, networkAuthOption);
+  if (nullptr == irReader_) {
+    VLOG(2) << "Failed to open kv-ir stream \"" << splitPath_
+            << "\" for reading.";
+    return ErrorCode::InternalError;
+  }
+
+  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.";
+    return ErrorCode::InternalError;
+  }
+  irDeserializer_ = std::make_shared<
+      ::clp::ffi::ir_stream::Deserializer<ClpIrUnitHandler, QueryHandlerType>>(
+      std::move(deserializerResult).value());
+
+  currentSplitLoaded_ = true;
+  return ErrorCode::Success;
+}
+
+std::vector<std::pair<std::string, search::ast::literal_type_bitmask_t>>
+ClpIrCursor::splitFieldsToNamesAndTypes() const {
+  auto result = std::vector<
+      std::pair<std::string, search::ast::literal_type_bitmask_t>>{};
+  for (size_t i{0}; i < outputColumns_.size(); ++i) {
+    auto column = outputColumns_[i];
+    search::ast::literal_type_bitmask_t literalType;
+    switch (column.type) {
+      case ColumnType::Array:
+        literalType = search::ast::LiteralType::ArrayT;
+        break;
+      case ColumnType::Boolean:
+        literalType = search::ast::LiteralType::BooleanT;
+        break;
+      case ColumnType::Float:
+        literalType = search::ast::LiteralType::FloatT;
+        break;
+      case ColumnType::Integer:
+        literalType = search::ast::LiteralType::IntegerT;
+        break;
+      case ColumnType::String:
+        literalType = search::ast::LiteralType::VarStringT;
+        break;
+      case ColumnType::Timestamp:
+        // TODO: IR timestamp support pending; constrain to Unknown to avoid
+        // mismatched projections.
+        literalType = search::ast::LiteralType::EpochDateT;
+        break;
+      default:
+        literalType = search::ast::LiteralType::UnknownT;
+        break;
+    }
+    result.emplace_back(column.name, literalType);
+  }
+  return result;
+}
+
+ystdlib::error_handling::Result<void> ClpIrCursor::deserialize(
+    uint64_t numRows) {
+  irDeserializer_->get_ir_unit_handler().clearFilteredLogEvents();
+  uint64_t cnt{0};
+  while (cnt < numRows) {
+    auto deserializeResult =
+        irDeserializer_->deserialize_next_ir_unit(*irReader_);
+    if (deserializeResult.has_error()) {
+      auto error = deserializeResult.error();
+      if (std::errc::result_out_of_range == error ||
+          irDeserializer_->is_stream_completed()) {
+        break;
+      }
+      return error;
+    }
+    cnt++;
+  }
+  return ystdlib::error_handling::success();
+}
+
+VectorPtr ClpIrCursor::createVectorHelper(
+    memory::MemoryPool* pool,
+    const TypePtr& vectorType,
+    size_t vectorSize) {
+  if (vectorType->kind() == TypeKind::ROW) {
+    std::vector<VectorPtr> children;
+    auto& rowType = vectorType->as<TypeKind::ROW>();
+    for (uint32_t i = 0; i < rowType.size(); ++i) {
+      children.push_back(
+          createVectorHelper(pool, rowType.childAt(i), vectorSize));
+    }
+    return std::make_shared<RowVector>(
+        pool, vectorType, nullptr, vectorSize, std::move(children));
+  }
+  auto vector = BaseVector::create(vectorType, vectorSize, pool);
+  vector->setNulls(allocateNulls(vectorSize, pool, bits::kNull));
+  VELOX_CHECK_LT(
+      readerIndex_,
+      projectedColumnNameNodeIdMap_.size(),
+      "Reader index out of bounds");
+  auto projectedColumn = outputColumns_[readerIndex_];
+  auto projectedColumnType = projectedColumn.type;
+  auto it = projectedColumnNameNodeIdMap_.find(projectedColumn.name);
+  VELOX_CHECK(
+      it != projectedColumnNameNodeIdMap_.end(),
+      "Projected column '{}' not found in node id map",
+      projectedColumn.name);
+  auto projectedColumnNodeId = it->second;
+  readerIndex_++;
+  return std::make_shared<LazyVector>(
+      pool,
+      vectorType,
+      vectorSize,
+      std::make_unique<ClpIrVectorLoader>(
+          projectedColumnType,
+          projectedColumnNodeId,
+          irDeserializer_->get_ir_unit_handler().getFilteredLogEvents()),
+      std::move(vector));
+}
+
+} // namespace facebook::velox::connector::clp::search_lib

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

@@ -0,0 +1,89 @@
+/*
+ * Copyright (c) Facebook, Inc. and its 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 "ffi/ir_stream/Deserializer.hpp"
+#include "velox/connectors/clp/search_lib/BaseClpCursor.h"
+#include "velox/connectors/clp/search_lib/ir/ClpIrUnitHandler.h"
+
+namespace facebook::velox::connector::clp::search_lib {
+
+class ClpIrCursor final : public BaseClpCursor {
+ public:
+  explicit ClpIrCursor(
+      clp_s::InputSource inputSource,
+      std::string_view splitPath,
+      bool ignoreCase)
+      : BaseClpCursor(inputSource, splitPath), ignoreCase_(ignoreCase) {}
+
+  ClpIrCursor(const ClpIrCursor&) = delete;
+  ClpIrCursor& operator=(const ClpIrCursor&) = delete;
+  ClpIrCursor(ClpIrCursor&&) = delete;
+  ClpIrCursor& operator=(ClpIrCursor&&) = delete;
+
+  uint64_t fetchNext(uint64_t numRows) override;
+
+  size_t getNumFilteredRows() const override;
+
+  VectorPtr createVector(
+      memory::MemoryPool* pool,
+      const TypePtr& vectorType,
+      size_t vectorSize) override;
+
+ protected:
+  ErrorCode loadSplit() override;
+
+ private:
+  std::function<ystdlib::error_handling::Result<void>(
+      bool,
+      ::clp::ffi::SchemaTree::Node::id_t,
+      std::pair<std::string_view, size_t>)>
+      projectionResolutionCallback_ =
+          [this](
+              [[maybe_unused]] bool isAutoGenerated,
+              [[maybe_unused]] ::clp::ffi::SchemaTree::Node::id_t nodeId,
+              [[maybe_unused]] std::pair<std::string_view, size_t>
+                  projected_key_and_index)
+      -> ystdlib::error_handling::Result<void> {
+    projectedColumnNameNodeIdMap_.insert(
+        {std::string(projected_key_and_index.first), nodeId});
+    return ystdlib::error_handling::success();
+  };
+  using QueryHandlerType = ::clp::ffi::ir_stream::search::QueryHandler<
+      decltype(projectionResolutionCallback_)>;
+  bool ignoreCase_;
+  std::shared_ptr<
+      ::clp::ffi::ir_stream::Deserializer<ClpIrUnitHandler, QueryHandlerType>>
+      irDeserializer_;
+  std::shared_ptr<::clp::ReaderInterface> irReader_{nullptr};
+  std::unordered_map<std::string, ::clp::ffi::SchemaTree::Node::id_t>
+      projectedColumnNameNodeIdMap_;
+  size_t readerIndex_{0};
+
+  std::vector<
+      std::pair<std::string, clp_s::search::ast::literal_type_bitmask_t>>
+  splitFieldsToNamesAndTypes() const;
+
+  ystdlib::error_handling::Result<void> deserialize(uint64_t numRows);
+
+  VectorPtr createVectorHelper(
+      memory::MemoryPool* pool,
+      const TypePtr& vectorType,
+      size_t vectorSize);
+};
+
+} // namespace facebook::velox::connector::clp::search_lib