Added a check sort app and fixed metadata bug in chunk reader

Author: m09526

cpp/include/cmdline/data_sink.hpp

@@ -19,11 +19,15 @@ struct SinkInfoDetails
 [[nodiscard]] std::unique_ptr<cudf::io::data_sink> make_data_sink(std::convertible_to<std::string> auto const &path,
   std::shared_ptr<Aws::S3::S3Client> &s3client);
 
-cudf::io::chunked_parquet_writer_options_builder write_opts(cudf::io::sink_info const &sink) noexcept;
+cudf::io::chunked_parquet_writer_options_builder write_opts(cudf::io::sink_info const &sink,
+  cudf::io::table_metadata const &metadata) noexcept;
 
-[[no_discard]] SinkInfoDetails make_writer(std::string const &path, std::shared_ptr<Aws::S3::S3Client> &s3client);
+[[no_discard]] SinkInfoDetails make_writer(std::string const &path,
+  cudf::io::table_metadata const &metadata,
+  std::shared_ptr<Aws::S3::S3Client> &s3client);
 
 [[no_discard]] SinkInfoDetails make_writer(std::convertible_to<std::string> auto const &path,
+  cudf::io::table_metadata const &metadata,
   std::shared_ptr<Aws::S3::S3Client> &s3client) {
-    return make_writer(path, s3client);
+    return make_writer(path, metadata, s3client);
 }

cpp/src/cmdline/CMakeLists.txt

@@ -37,3 +37,18 @@ target_link_system_libraries(
 target_link_libraries(chunk_reader PRIVATE gpu_compact::io)
 target_include_directories(chunk_reader PRIVATE "${CMAKE_BINARY_DIR}/configured_files/include")
 target_include_directories(chunk_reader PRIVATE "${CMAKE_SOURCE_DIR}/include/cmdline")
+
+add_executable(
+  check_sort
+  check_sort.cpp
+  configure_logging.cpp
+  ../../include/cmdline/configure_logging.hpp)
+target_link_system_libraries(
+    check_sort
+  PRIVATE
+  CLI11::CLI11
+  fmt::fmt
+  spdlog::spdlog
+  cudf::cudf)
+target_include_directories(check_sort PRIVATE "${CMAKE_BINARY_DIR}/configured_files/include")
+target_include_directories(check_sort PRIVATE "${CMAKE_SOURCE_DIR}/include/cmdline")

cpp/src/cmdline/check_sort.cpp

@@ -0,0 +1,85 @@
+#include "configure_logging.hpp"
+
+#include <CLI/CLI.hpp>// NOLINT
+
+#include <cudf/concatenate.hpp>
+#include <cudf/io/datasource.hpp>
+#include <cudf/io/parquet.hpp>
+#include <cudf/io/types.hpp>
+#include <cudf/sorting.hpp>
+#include <cudf/table/table.hpp>
+#include <cudf/table/table_view.hpp>
+#include <cudf/types.hpp>
+#include <cudf/utilities/error.hpp>
+#include <internal_use_only/config.hpp>
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/owning_wrapper.hpp>
+#include <rmm/mr/device/pool_memory_resource.hpp>
+#ifdef SPDLOG_ACTIVE_LEVEL
+#undef SPDLOG_ACTIVE_LEVEL
+#endif
+#define SPDLOG_ACTIVE_LEVEL SPDLOG_LEVEL_DEBUG
+#include <spdlog/spdlog.h>
+
+#include <cstddef>
+#include <memory>
+#include <string>
+
+int main(int argc, char **argv) {
+    configure_logging();
+    // NOLINTNEXTLINE
+    CLI::App app{ "Simple program to check if Parquet file is sorted with cuDF", "check_sort" };
+    app.set_version_flag("--version", std::string{ gpu_compact::cmake::project_version });
+
+    std::string inputFile;
+    app.add_option("input", inputFile, "Input Parquet file")->required();
+    std::string colName;
+    app.add_option("Column name", colName, "Column to validate sort order")->required();
+    CLI11_PARSE(app, argc, argv);// NOLINT
+
+    // force gpu initialization so it's not included in the time
+    rmm::cuda_stream_default.synchronize();
+
+    auto cuda_mr = std::make_shared<rmm::mr::cuda_memory_resource>();
+    auto mr =
+      rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(cuda_mr, rmm::percent_of_free_device_memory(95));
+    rmm::mr::set_current_device_resource(mr.get());
+
+    auto opts =
+      cudf::io::parquet_reader_options::builder(cudf::io::source_info(inputFile)).columns({ colName }).build();
+    cudf::io::chunked_parquet_reader reader{ 500 * 1'048'576l, 500 * 1'048'576l, opts };
+
+    SPDLOG_INFO("Validating sort on column '{}' in file '{}'", colName, inputFile);
+
+    // Loop doing reads
+    // A sort problem may occur across a chunk boundary, so we use a slightly hacky workaround of keeping
+    // the last chunk stored and then concatenate the current chunk on to it and then check that for sorting.
+    ::size_t totalRowsRead = 0;
+    ::size_t chunkNo = 0;
+    cudf::io::table_with_metadata prevTable;
+    while (reader.has_next()) {
+        auto currentTable = reader.read_chunk();
+        ::size_t rowsRead = currentTable.tbl->num_rows();
+        SPDLOG_INFO("Checking chunk number {:d} has {:d} rows", chunkNo, rowsRead);
+
+        std::unique_ptr<cudf::table> checkTable;
+        cudf::table_view checkView = currentTable.tbl->view();
+        if (prevTable.tbl) {
+            checkTable = cudf::concatenate(std::vector<cudf::table_view>{ *prevTable.tbl, *currentTable.tbl });
+            checkView = std::move(checkTable->view());
+        }
+        bool chunkSorted = cudf::is_sorted(checkView, { cudf::order::ASCENDING }, { cudf::null_order::AFTER });
+        if (!chunkSorted) {
+            SPDLOG_ERROR("Chunk number {:d} contains an incorrect sort order between rows [{:d},{:d})",
+              chunkNo,
+              totalRowsRead,
+              totalRowsRead + rowsRead);
+            CUDF_FAIL("Incorrect sort detected");
+        }
+        chunkNo++;
+        totalRowsRead += rowsRead;
+        prevTable = std::move(currentTable);
+    }
+
+    SPDLOG_INFO("Finished, file is correctly sorted");
+}

cpp/src/cmdline/chunk_reader.cpp

@@ -42,14 +42,19 @@ ::size_t calcRowsWritten(auto const &readers) noexcept {
       });
 }
 
-[[nodiscard]] inline std::chrono::time_point<std::chrono::steady_clock> timestamp() noexcept {
+[[nodiscard]] std::chrono::time_point<std::chrono::steady_clock> timestamp() noexcept {
     return std::chrono::steady_clock::now();
 }
 
+[[nodiscard]] cudf::io::table_metadata grabMetaData(std::string const &file) {
+    auto opts = cudf::io::parquet_reader_options::builder(cudf::io::source_info(file)).build();
+    return cudf::io::read_parquet(opts).metadata;
+}
+
 int main(int argc, char **argv) {
     configure_logging();
     // NOLINTNEXTLINE
-    CLI::App app{ "Simple program to test spooling through Parquet files in chunks with cuDF", "chunk_reader" };
+    CLI::App app{ "Simple program to test chunking compaction algorithm with cuDF", "chunk_reader" };
     app.set_version_flag("--version", std::string{ gpu_compact::cmake::project_version });
 
     std::string outputFile;
@@ -98,8 +103,10 @@ int main(int argc, char **argv) {
             dynamic_cast<gpu_compact::io::PrefetchingSource *>(std::get<0>(readers.back()).get())->prefetch(true);
         }
 
+        // Grab metadata for schema from first file
+        auto const tableMetadata = grabMetaData(inputFiles[0]);
         // Make writer
-        SinkInfoDetails sinkDetails = make_writer(outputFile, s3client);
+        SinkInfoDetails sinkDetails = make_writer(outputFile, tableMetadata, s3client);
         auto &writer = *sinkDetails.writer;
 
         SPDLOG_INFO("Start reading files");

cpp/src/cmdline/data_sink.cpp

@@ -13,9 +13,11 @@
     }
 }
 
-cudf::io::chunked_parquet_writer_options_builder write_opts(cudf::io::sink_info const &sink) noexcept {
+cudf::io::chunked_parquet_writer_options_builder write_opts(cudf::io::sink_info const &sink,
+  cudf::io::table_metadata const &metadata) noexcept {
     return cudf::io::chunked_parquet_writer_options::builder(sink)
       .compression(cudf::io::compression_type::ZSTD)
+      .metadata(cudf::io::table_input_metadata{ metadata })
       .row_group_size_bytes(65 * 1'048'576)
       .row_group_size_rows(1'000'000)
       .max_page_size_bytes(512 * 1024)
@@ -25,9 +27,11 @@ cudf::io::chunked_parquet_writer_options_builder write_opts(cudf::io::sink_info
       .dictionary_policy(cudf::io::dictionary_policy::ADAPTIVE);
 }
 
-[[no_discard]] SinkInfoDetails make_writer(std::string const &path, std::shared_ptr<Aws::S3::S3Client> &s3client) {
+[[no_discard]] SinkInfoDetails make_writer(std::string const &path,
+  cudf::io::table_metadata const &metadata,
+  std::shared_ptr<Aws::S3::S3Client> &s3client) {
     auto data_sink = make_data_sink(path, s3client);
     cudf::io::sink_info sink{ &*data_sink };
-    auto wopts = write_opts(sink);
+    auto wopts = write_opts(sink, metadata);
     return { sink, std::move(data_sink), std::make_unique<cudf::io::parquet_chunked_writer>(wopts.build()) };
 }