Add filter program

Author: EthanSteinberg

build_helper.sh

@@ -1,9 +1,10 @@
 #!/bin/sh
 
 cd native
-bazel build -c opt _meds_reader.so meds_reader_convert
+bazel build -c opt _meds_reader.so meds_reader_convert meds_reader_filter
 cd ..
 
 rm -f src/meds_reader/_meds_reader* src/meds_reader/meds_reader_convert*
 cp native/bazel-bin/_meds_reader.so src/meds_reader/_meds_reader.so
 cp native/bazel-bin/meds_reader_convert src/meds_reader/meds_reader_convert
+cp native/bazel-bin/meds_reader_filter src/meds_reader/meds_reader_filter

native/BUILD

@@ -26,6 +26,16 @@ cc_library(
     ],
 )
 
+cc_library(
+    name="filter_database",
+    srcs=["filter_database.cc", "pdqsort.h"],
+    hdrs=["filter_database.hh"],
+    deps=[
+        ":mmap_file",
+    ],
+)
+
+
 cc_library(
     name="binary_version",
     hdrs=["binary_version.hh"],
@@ -84,6 +94,15 @@ cc_binary(
     ],
 )
 
+cc_binary(
+    name="meds_reader_filter",
+    srcs=["meds_reader_filter.cc"],
+    deps=[
+        ":filter_database",
+        "@CLI11//:cli11",
+    ],
+)
+
 pybind_extension(
     name="_meds_reader",
     srcs=[

native/filter_database.cc

@@ -0,0 +1,203 @@
+#include "filter_database.hh"
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <unistd.h>
+
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <string>
+#include <thread>
+
+#include "mmap_file.hh"
+#include "pdqsort.h"
+
+namespace {
+
+void copy_subset(const std::filesystem::path& source_path,
+                 const std::filesystem::path& destination_path,
+                 const std::vector<size_t>& offsets) {
+    MmapFile source(source_path);
+
+    absl::Span<const uint64_t> byte_offsets = source.data<uint64_t>();
+
+    std::string_view data = source.bytes();
+
+    std::ofstream destination(destination_path, std::ios_base::out |
+                                                    std::ios_base::binary |
+                                                    std::ios_base::trunc);
+
+    std::vector<const char*> result_pointers;
+    result_pointers.reserve(offsets.size());
+
+    std::vector<uint64_t> result_byte_offsets;
+    result_byte_offsets.reserve(offsets.size() + 1);
+
+    uint64_t current_offset = sizeof(uint64_t) * (offsets.size() + 1);
+
+    for (size_t offset : offsets) {
+        uint64_t start = byte_offsets[offset];
+        uint64_t end = byte_offsets[offset + 1];
+
+        result_byte_offsets.push_back(current_offset);
+        result_pointers.push_back(data.data() + start);
+
+        current_offset += (end - start);
+    }
+
+    result_byte_offsets.push_back(current_offset);
+
+    destination.write((const char*)result_byte_offsets.data(),
+                      sizeof(uint64_t) * result_byte_offsets.size());
+
+    for (size_t i = 0; i < result_pointers.size(); i++) {
+        uint64_t length = result_byte_offsets[i + 1] - result_byte_offsets[i];
+        destination.write(result_pointers[i], length);
+    }
+}
+
+void filter_database_property(const std::filesystem::path& source_path,
+                              const std::filesystem::path& destination_path,
+                              const std::vector<size_t>& offsets,
+                              const std::string& property_name) {
+    std::filesystem::path source_property_path = source_path / property_name;
+    std::filesystem::path destination_property_path =
+        destination_path / property_name;
+
+    std::filesystem::create_directory(destination_property_path);
+
+    std::filesystem::path source_dictionary_path =
+        source_property_path / "dictionary";
+
+    if (std::filesystem::exists(source_dictionary_path)) {
+        std::filesystem::copy(source_dictionary_path,
+                              destination_property_path / "dictionary");
+    }
+
+    std::filesystem::copy(source_property_path / "zdict",
+                          destination_property_path / "zdict");
+
+    copy_subset(source_property_path / "data",
+                destination_property_path / "data", offsets);
+}
+
+}  // namespace
+
+void filter_database(const char* source, const char* destination,
+                     const char* subject_ids_file, int num_threads) {
+    std::filesystem::path source_path(source);
+    std::filesystem::path destination_path(destination);
+
+    MmapFile subject_ids_data{std::string(subject_ids_file)};
+
+    absl::Span<const int64_t> unsorted_subject_ids =
+        subject_ids_data.data<int64_t>();
+
+    std::vector<int64_t> sorted_subject_ids(std::begin(unsorted_subject_ids),
+                                            std::end(unsorted_subject_ids));
+
+    pdqsort(std::begin(sorted_subject_ids), std::end(sorted_subject_ids));
+
+    absl::Span<const int64_t> subject_ids(sorted_subject_ids.data(),
+                                          sorted_subject_ids.size());
+
+    std::filesystem::create_directory(destination_path);
+
+    std::filesystem::copy(source_path / "metadata",
+                          destination_path / "metadata");
+
+    {
+        std::ofstream subject_ids_file(
+            destination_path / "subject_id",
+            std::ios_base::out | std::ios_base::binary | std::ios_base::trunc);
+
+        subject_ids_file.write((const char*)subject_ids.data(),
+                               sizeof(int64_t) * subject_ids.size());
+    }
+
+    std::vector<size_t> offsets;
+    offsets.reserve(subject_ids.size());
+
+    {
+        MmapFile source_subject_ids_file(source_path / "subject_id");
+        absl::Span<const int64_t> source_subject_ids =
+            source_subject_ids_file.data<int64_t>();
+
+        auto first =
+            std::lower_bound(std::begin(source_subject_ids),
+                             std::end(source_subject_ids), subject_ids.front());
+        auto last =
+            std::upper_bound(std::begin(source_subject_ids),
+                             std::end(source_subject_ids), subject_ids.back());
+
+        for (int64_t subject_id : subject_ids) {
+            auto iter = std::lower_bound(first, last, subject_id);
+            if (*iter != subject_id) {
+                throw std::runtime_error(
+                    std::string("Could not find subject_id ") +
+                    std::to_string(subject_id) + " in database " +
+                    std::to_string(*iter));
+            }
+
+            offsets.push_back(iter - std::begin(source_subject_ids));
+
+            first = ++iter;
+        }
+    }
+
+    {
+        MmapFile source_subject_lengths_file(source_path /
+                                             "meds_reader.length");
+        absl::Span<const uint32_t> source_subject_lengths =
+            source_subject_lengths_file.data<uint32_t>();
+        std::vector<uint32_t> subject_lengths;
+        subject_lengths.reserve(subject_ids.size());
+
+        for (size_t offset : offsets) {
+            subject_lengths.push_back(source_subject_lengths[offset]);
+        }
+
+        std::ofstream subject_lengths_file(
+            destination_path / "meds_reader.length",
+            std::ios_base::out | std::ios_base::binary | std::ios_base::trunc);
+
+        subject_lengths_file.write((const char*)subject_lengths.data(),
+                                   sizeof(uint32_t) * subject_lengths.size());
+    }
+
+    std::vector<std::string> properties;
+
+    {
+        MmapFile property_file(source_path / "meds_reader.properties");
+
+        const char* data = property_file.bytes().data();
+        const char* end =
+            property_file.bytes().data() + property_file.bytes().size();
+
+        while (data != end) {
+            size_t name_length = *(const size_t*)data;
+            data += sizeof(size_t);
+            properties.push_back(std::string(data, name_length));
+            data += name_length;
+            data += sizeof(int64_t);
+
+            std::cout << "Got property " << properties.back() << std::endl;
+        }
+    }
+
+    properties.push_back("meds_reader.null_map");
+
+    for (const auto& property : properties) {
+        filter_database_property(source_path, destination_path, offsets,
+                                 property);
+    }
+
+    std::filesystem::copy(source_path / "meds_reader.properties",
+                          destination_path / "meds_reader.properties");
+
+    std::filesystem::copy(source_path / "meds_reader.version",
+                          destination_path / "meds_reader.version");
+}

native/filter_database.hh

@@ -0,0 +1,4 @@
+#pragma once
+
+void filter_database(const char* source, const char* destination,
+                     const char* subject_ids_file, int num_threads);

native/meds_reader_filter.cc

@@ -0,0 +1,38 @@
+#include <CLI/CLI.hpp>
+#include <iostream>
+
+#include "filter_database.hh"
+
+int main(int argc, char** argv) {
+    CLI::App app{
+        "meds_reader_filter is a program for converting a MEDS dataset to a "
+        "meds_reader SubjectDatabase.",
+        "meds_reader_convert"};
+    argv = app.ensure_utf8(argv);
+
+    std::string source_dataset, destination_database, subject_ids_file;
+    app.add_option("source_dataset", source_dataset,
+                   "A path to the source MEDS dataset")
+        ->required();
+    app.add_option(
+           "destination_database", destination_database,
+           "A path of where to write the resulting meds_reader database.")
+        ->required();
+
+    app.add_option(
+           "subject_ids_file", subject_ids_file,
+           "A path of where to write the resulting meds_reader database.")
+        ->required();
+
+    int num_threads = 1;
+    app.add_option("--num_threads", num_threads,
+                   "The number of threads to use when processing")
+        ->capture_default_str();
+
+    CLI11_PARSE(app, argc, argv);
+
+    filter_database(source_dataset.c_str(), destination_database.c_str(),
+                    subject_ids_file.c_str(), num_threads);
+
+    return 0;
+}

native/mmap_file.hh

@@ -5,6 +5,7 @@
 #include <sys/stat.h>
 #include <unistd.h>
 
+#include <cstring>
 #include <filesystem>
 #include <string_view>
 

pyproject.toml

@@ -26,6 +26,7 @@ Tracker = "https://github.com/som-shahlab/meds_reader/issues"
 
 [project.scripts]
 meds_reader_convert = "meds_reader:meds_reader_convert"
+meds_reader_filter = "meds_reader:meds_reader_filter"
 meds_reader_verify = "meds_reader:meds_reader_verify"
 
 [tool.isort]

setup.py

@@ -119,6 +119,7 @@ def build_extensions(self) -> None:
     ext_modules=[
         BazelExtension("meds_reader._meds_reader", "_meds_reader.so", "native"),
         BazelExtension("meds_reader.meds_reader_convert", "meds_reader_convert", "native"),
+        BazelExtension("meds_reader.meds_reader_filter", "meds_reader_filter", "native"),
     ],
     cmdclass={"build_ext": cmake_build_ext},
     zip_safe=False,

src/meds_reader/__init__.py

@@ -105,6 +105,14 @@ def meds_reader_convert():
                 os.execv(executible, sys.argv)
 
 
+def meds_reader_filter():
+    submodules = importlib.resources.files("meds_reader")
+    for module in submodules.iterdir():
+        if module.name.startswith("meds_reader_filter"):
+            with importlib.resources.as_file(module) as executible:
+                os.execv(executible, sys.argv)
+
+
 def _row_generator(database: _meds_reader.SubjectDatabase, data: pd.DataFrame):
     current_index = None
     current_rows: List[Any] = []