Add PageFile Reader/Writer

velox/common/file/FileInputStream.cpp

@@ -19,7 +19,7 @@
 namespace facebook::velox::common {
 
 FileInputStream::FileInputStream(
-    std::unique_ptr<ReadFile>&& file,
+    const std::shared_ptr<ReadFile>& file,
     uint64_t bufferSize,
     memory::MemoryPool* pool)
     : file_(std::move(file)),

velox/common/file/FileInputStream.h

@@ -28,7 +28,7 @@ namespace facebook::velox::common {
 class FileInputStream : public ByteInputStream {
  public:
   FileInputStream(
-      std::unique_ptr<ReadFile>&& file,
+      const std::shared_ptr<ReadFile>& file,
       uint64_t bufferSize,
       memory::MemoryPool* pool);
 
@@ -108,7 +108,7 @@ class FileInputStream : public ByteInputStream {
 
   void updateStats(uint64_t readBytes, uint64_t readTimeNs);
 
-  const std::unique_ptr<ReadFile> file_;
+  const std::shared_ptr<ReadFile> file_;
   const uint64_t fileSize_;
   const uint64_t bufferSize_;
   memory::MemoryPool* const pool_;

velox/dwio/CMakeLists.txt

@@ -32,6 +32,7 @@ velox_link_libraries(
 
 add_subdirectory(common)
 add_subdirectory(catalog)
+add_subdirectory(cte)
 add_subdirectory(dwrf)
 add_subdirectory(orc)
 add_subdirectory(parquet)

velox/dwio/common/Options.cpp

@@ -39,6 +39,8 @@ FileFormat toFileFormat(std::string_view s) {
     return FileFormat::ORC;
   } else if (s == "sst") {
     return FileFormat::SST;
+  } else if (s == "pagefile") {
+    return FileFormat::PAGEFILE;
   }
   return FileFormat::UNKNOWN;
 }
@@ -65,6 +67,8 @@ std::string_view toString(FileFormat fmt) {
       return "orc";
     case FileFormat::SST:
       return "sst";
+    case FileFormat::PAGEFILE:
+      return "pagefile";
     default:
       return "unknown";
   }

velox/dwio/common/Options.h

@@ -51,6 +51,7 @@ enum class FileFormat {
   NIMBLE = 8,
   ORC = 9,
   SST = 10, // rocksdb sst format
+  PAGEFILE = 11,
 };
 
 FileFormat toFileFormat(std::string_view s);

velox/dwio/cte/CMakeLists.txt

@@ -0,0 +1,20 @@
+# 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.
+
+add_subdirectory(reader)
+add_subdirectory(writer)
+
+if(${VELOX_BUILD_TESTING})
+  add_subdirectory(tests)
+endif()

velox/dwio/cte/RegisterCtePageReader.h

@@ -0,0 +1,25 @@
+/*
+ * 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
+
+namespace facebook::velox::pagefile {
+
+void registerCtePageReaderFactory();
+
+void unregisterCtePageReaderFactory();
+
+} // namespace facebook::velox::pagefile

velox/dwio/cte/RegisterCtePageWriter.h

@@ -0,0 +1,25 @@
+/*
+ * 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
+
+namespace facebook::velox::pagefile {
+
+void registerCtePageWriterFactory();
+
+void unregisterCtePageWriterFactory();
+
+} // namespace facebook::velox::pagefile

velox/dwio/cte/reader/CMakeLists.txt

@@ -0,0 +1,16 @@
+# 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(velox_dwio_cte_page_reader CtePageReader.cpp)
+velox_link_libraries(velox_dwio_cte_page_reader velox_dwio_common fmt::fmt)

velox/dwio/cte/reader/CtePageReader.cpp

@@ -0,0 +1,128 @@
+/*
+ * 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 "velox/dwio/cte/reader/CtePageReader.h"
+
+namespace facebook::velox::pagefile {
+using namespace facebook::velox::common;
+using namespace dwio::common;
+
+/// Metadata and options for reading Parquet.
+class ReaderBase {
+ public:
+  ReaderBase(
+      std::unique_ptr<dwio::common::BufferedInput>,
+      const dwio::common::ReaderOptions& options);
+
+  virtual ~ReaderBase() = default;
+
+  memory::MemoryPool& getMemoryPool() const {
+    return pool_;
+  }
+
+  FileInputStream* fileInputStream() const {
+    return input_.get();
+  }
+
+  const dwio::common::ReaderOptions& getReaderOptions() const {
+    return options_;
+  }
+
+ private:
+  memory::MemoryPool& pool_;
+  // Copy of options. Must be owned by 'this'.
+  const dwio::common::ReaderOptions options_;
+  std::unique_ptr<FileInputStream> input_;
+};
+
+ReaderBase::ReaderBase(
+    std::unique_ptr<dwio::common::BufferedInput> input,
+    const dwio::common::ReaderOptions& options)
+    : pool_{options.memoryPool()}, options_{options} {
+  input_ = std::make_unique<facebook::velox::common::FileInputStream>(
+      input->getReadFile(), 1 << 20, &(options.memoryPool()));
+}
+
+CtePageReader::CtePageReader(
+    const dwio::common::ReaderOptions& options,
+    std::unique_ptr<dwio::common::BufferedInput> input)
+    : readerBase_(std::make_shared<ReaderBase>(std::move(input), options)),
+      options_(options){};
+
+std::optional<uint64_t> CtePageReader::numberOfRows() const {
+  return readerBase_->fileInputStream()->size();
+}
+
+std::unique_ptr<RowReader> CtePageReader::createRowReader(
+    const RowReaderOptions& opts) const {
+  return std::make_unique<CtePageRowReader>(readerBase_, opts);
+}
+
+std::unique_ptr<CtePageReader> CtePageReader::create(
+    std::unique_ptr<dwio::common::BufferedInput> input,
+    const dwio::common::ReaderOptions& options) {
+  return std::make_unique<CtePageReader>(options, std::move(input));
+}
+
+CtePageRowReader::CtePageRowReader(
+    const std::shared_ptr<ReaderBase>& readerBase,
+    const dwio::common::RowReaderOptions& options)
+    : readerBase_{readerBase},
+      options_{options},
+      schema_{readerBase_->getReaderOptions().fileSchema()},
+      serde_(getNamedVectorSerde(VectorSerde::Kind::kPresto)){};
+
+uint64_t CtePageRowReader::next(
+    uint64_t size,
+    velox::VectorPtr& result,
+    const dwio::common::Mutation*) {
+  if (readerBase_->fileInputStream()->atEnd()) {
+    return 0;
+  }
+
+  RowVectorPtr rowVectorPtr = vectorPtrToRowVector(size, result);
+  VectorStreamGroup::read(
+      readerBase_->fileInputStream(),
+      &(readerBase_->getMemoryPool()),
+      schema_,
+      serde_,
+      &(rowVectorPtr),
+      &readOptions_);
+  return size;
+}
+
+const RowVectorPtr CtePageRowReader::vectorPtrToRowVector(
+    uint64_t size,
+    velox::VectorPtr result) {
+  std::vector<velox::VectorPtr> vectors;
+  vectors.emplace_back(result);
+
+  return std::make_shared<RowVector>(
+      &(readerBase_->getMemoryPool()),
+      schema_,
+      BufferPtr(nullptr),
+      size,
+      vectors);
+}
+
+void registerCtePageReaderFactory() {
+  dwio::common::registerReaderFactory(std::make_shared<CtePageReaderFactory>());
+}
+
+void unregisterCtePageReaderFactory() {
+  dwio::common::unregisterReaderFactory(dwio::common::FileFormat::PAGEFILE);
+}
+} // namespace facebook::velox::pagefile