[c++] Reindexer overrides and fast COO/CSR

.github/workflows/python-ci-single.yml

@@ -53,7 +53,7 @@ jobs:
 
     - name: Check C++ Format
       shell: bash
-      run: ./scripts/run-clang-format.sh . clang-format 0 $(find libtiledbsoma -name "*.cc" -or -name "*.h")
+      run: ./scripts/run-clang-format.sh . clang-format 0 $(find libtiledbsoma -name "*.cc" -or -name "*.h" | grep -v external)
 
   build:
     runs-on: ${{ inputs.os }}

apis/python/src/tiledbsoma/_experiment.py

@@ -5,16 +5,18 @@
 
 """Implementation of a SOMA Experiment.
 """
+import functools
+from typing import Any, Optional
 
-from typing import Any
-
-from somacore import experiment
+from somacore import experiment, query
+from typing_extensions import Self
 
 from ._collection import Collection, CollectionBase
 from ._dataframe import DataFrame
 from ._measurement import Measurement
 from ._tdb_handles import Wrapper
 from ._tiledb_object import AnyTileDBObject
+from .utils import build_index
 
 
 class Experiment(  # type: ignore[misc]  # __eq__ false positive
@@ -74,3 +76,23 @@ def _set_create_metadata(cls, handle: Wrapper[Any]) -> None:
         # TileDB Cloud UI to detect that they are SOMA datasets.
         handle.metadata["dataset_type"] = "soma"
         return super()._set_create_metadata(handle)
+
+    def axis_query(  # type: ignore
+        self,
+        measurement_name: str,
+        *,
+        obs_query: Optional[query.AxisQuery] = None,
+        var_query: Optional[query.AxisQuery] = None,
+    ) -> query.ExperimentAxisQuery[Self]:  # type: ignore
+        """Creates an axis query over this experiment.
+        Lifecycle: maturing
+        """
+        # mypy doesn't quite understand descriptors so it issues a spurious
+        # error here.
+        return query.ExperimentAxisQuery(  # type: ignore
+            self,
+            measurement_name,
+            obs_query=obs_query or query.AxisQuery(),
+            var_query=var_query or query.AxisQuery(),
+            index_factory=functools.partial(build_index, context=self.context),
+        )

apis/python/src/tiledbsoma/pytiledbsoma.cc

@@ -37,6 +37,7 @@
 #include <pybind11/stl_bind.h>
 
 #include <tiledbsoma/tiledbsoma>
+#include <tiledbsoma/reindexer/reindexer.h>
 
 #include "query_condition.cc"
 
@@ -693,5 +694,63 @@ PYBIND11_MODULE(pytiledbsoma, m) {
         .def("get_enum_is_ordered", get_enum_is_ordered)
 
         .def("get_enum_label_on_attr", &SOMAArray::get_enum_label_on_attr);
+    // Efficient C++ re-indexing (aka hashing unique key values to an index
+    // between 0 and number of keys - 1) based on khash
+    py::class_<IntIndexer>(m, "IntIndexer")
+        .def(py::init<>())
+        .def(py::init<std::vector<int64_t>&, int>())
+        .def(
+            "map_locations",
+            [](IntIndexer& indexer,
+               py::array_t<int64_t> keys,
+               int num_threads) {
+                auto buffer = keys.request();
+                int64_t* data = static_cast<int64_t*>(buffer.ptr);
+                size_t length = buffer.shape[0];
+                indexer.map_locations(keys.data(), keys.size(), num_threads);
+            })
+        .def(
+            "map_locations",
+            [](IntIndexer& indexer,
+               std::vector<int64_t> keys,
+               int num_threads) {
+                indexer.map_locations(keys.data(), keys.size(), num_threads);
+            })
+        // Perform lookup for a large input array of keys and return the looked
+        // up value array (passing ownership from C++ to python)
+        .def(
+            "get_indexer",
+            [](IntIndexer& indexer, py::array_t<int64_t> lookups) {
+                auto input_buffer = lookups.request();
+                int64_t* input_ptr = static_cast<int64_t*>(input_buffer.ptr);
+                size_t size = input_buffer.shape[0];
+                auto results = py::array_t<int64_t>(size);
+                auto results_buffer = results.request();
+                size_t results_size = results_buffer.shape[0];
+
+                int64_t* results_ptr = static_cast<int64_t*>(
+                    results_buffer.ptr);
+
+                indexer.lookup(input_ptr, results_ptr, size);
+                return results;
+            })
+        // Perform lookup for a large input array of keys and writes the looked
+        // up values into previously allocated array (works for the cases in
+        // which python and R pre-allocate the array)
+        .def(
+            "get_indexer",
+            [](IntIndexer& indexer,
+               py::array_t<int64_t> lookups,
+               py::array_t<int64_t>& results) {
+                auto input_buffer = lookups.request();
+                int64_t* input_ptr = static_cast<int64_t*>(input_buffer.ptr);
+                size_t size = input_buffer.shape[0];
+
+                auto results_buffer = results.request();
+                int64_t* results_ptr = static_cast<int64_t*>(
+                    results_buffer.ptr);
+                size_t results_size = input_buffer.shape[0];
+                indexer.lookup(input_ptr, input_ptr, size);
+            });
 }
 }  // namespace tiledbsoma

apis/python/src/tiledbsoma/utils.py

@@ -0,0 +1,24 @@
+import numpy as np
+import pandas as pd
+
+from tiledbsoma import pytiledbsoma as clib
+
+from .options import SOMATileDBContext
+
+"""Builds an indexer object compatible with :meth:`pd.Index.get_indexer`."""
+
+
+def build_index(
+    keys: np.typing.NDArray[np.int64], context: SOMATileDBContext
+) -> clib.IntIndexer:
+    if len(np.unique(keys)) != len(keys):
+        raise pd.errors.InvalidIndexError(
+            "Reindexing only valid with uniquely valued Index objects"
+        )
+    tdb_concurrency = int(
+        context.tiledb_ctx.config().get("sm.compute_concurrency_level", 10)
+    )
+    thread_count = tdb_concurrency // 2
+    reindexer = clib.IntIndexer()
+    reindexer.map_locations(keys, thread_count)
+    return reindexer

apis/r/inst/include/tiledbsoma_types.h

@@ -19,6 +19,7 @@
 #include <tiledb/tiledb>					// for QueryCondition etc
 #define ARROW_SCHEMA_AND_ARRAY_DEFINED 1
 #include <tiledbsoma/tiledbsoma>
+#include <tiledbsoma/reindexer/reindexer.h>
 #include "rutilities.h"
 
 namespace tdbs = tiledbsoma;

libtiledbsoma/src/CMakeLists.txt

@@ -50,6 +50,7 @@ message(STATUS "Building with commit hash ${BUILD_COMMIT_HASH}")
 # Common object library
 # ###########################################################
 add_library(TILEDB_SOMA_OBJECTS OBJECT
+  ${CMAKE_CURRENT_SOURCE_DIR}/reindexer/reindexer.cc
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/managed_query.cc
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_array.cc
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_group.cc
@@ -67,9 +68,8 @@ add_library(TILEDB_SOMA_OBJECTS OBJECT
   ${CMAKE_CURRENT_SOURCE_DIR}/utils/util.cc
   ${CMAKE_CURRENT_SOURCE_DIR}/utils/version.cc
 
-  # TODO: uncomment when thread_pool is enabled
-  # ${CMAKE_CURRENT_SOURCE_DIR}/external/src/thread_pool/thread_pool.cc
-  # ${CMAKE_CURRENT_SOURCE_DIR}/external/src/thread_pool/status.cc
+  ${CMAKE_CURRENT_SOURCE_DIR}/external/src/thread_pool/thread_pool.cc
+  ${CMAKE_CURRENT_SOURCE_DIR}/external/src/thread_pool/status.cc
 )
 
 message(STATUS "Building TileDB without deprecation warnings")
@@ -87,6 +87,9 @@ set_property(TARGET TILEDB_SOMA_OBJECTS PROPERTY POSITION_INDEPENDENT_CODE ON)
 target_include_directories(TILEDB_SOMA_OBJECTS
   PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}
+  ${CMAKE_CURRENT_SOURCE_DIR}/vendor
+  ${CMAKE_CURRENT_SOURCE_DIR}/soma
+  ${CMAKE_CURRENT_SOURCE_DIR}/external/khash
   ${CMAKE_CURRENT_SOURCE_DIR}/external/include
   $<TARGET_PROPERTY:TileDB::tiledb_shared,INTERFACE_INCLUDE_DIRECTORIES>
   $<TARGET_PROPERTY:spdlog::spdlog,INTERFACE_INCLUDE_DIRECTORIES>
@@ -197,11 +200,16 @@ install(FILES
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_sparse_ndarray.h 
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_experiment.h 
   ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_measurement.h 
-  ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_object.h 
+  ${CMAKE_CURRENT_SOURCE_DIR}/soma/soma_object.h
   DESTINATION "include/tiledbsoma/soma"
 )
 
-install(FILES 
+install(FILES
+  ${CMAKE_CURRENT_SOURCE_DIR}/reindexer/reindexer.h
+  DESTINATION "include/tiledbsoma/reindexer/"
+)
+
+install(FILES
   ${CMAKE_CURRENT_SOURCE_DIR}/tiledbsoma/tiledbsoma
   DESTINATION "include/tiledbsoma"
 )
@@ -213,10 +221,11 @@ install(FILES
   ${CMAKE_CURRENT_SOURCE_DIR}/utils/stats.h  
   ${CMAKE_CURRENT_SOURCE_DIR}/utils/util.h  
   ${CMAKE_CURRENT_SOURCE_DIR}/utils/version.h
+
   DESTINATION "include/tiledbsoma/utils"
 )
 
-install(FILES 
+install(FILES
   ${CMAKE_CURRENT_SOURCE_DIR}/external/include/span/span.hpp
   DESTINATION "include/tiledbsoma/soma/span"
 )
@@ -294,6 +303,7 @@ if(TILEDBSOMA_BUILD_CLI)
     PRIVATE
     ${CMAKE_CURRENT_SOURCE_DIR}
     ${CMAKE_CURRENT_SOURCE_DIR}/external/include
+    ${CMAKE_CURRENT_SOURCE_DIR}/external/khash
     ${TILEDB_SOMA_EXPORT_HEADER_DIR}
     ${pybind11_INCLUDE_DIRS}
   )

libtiledbsoma/src/external/include/thread_pool/producer_consumer_queue.h

@@ -42,7 +42,7 @@
 #include <deque>
 #include <queue>
 
-#include "soma/logger_public.h"
+#include <logger_public.h>
 
 namespace tiledbsoma {