diff --git a/plugins/snowflake/superduper_snowflake/__init__.py b/plugins/snowflake/superduper_snowflake/__init__.py
index 117893925..1835c0dd9 100644
--- a/plugins/snowflake/superduper_snowflake/__init__.py
+++ b/plugins/snowflake/superduper_snowflake/__init__.py
@@ -1,7 +1,7 @@
 from .vector_search import SnowflakeVectorSearcher as VectorSearcher
 from .data_backend import SnowflakeDataBackend as DataBackend
 
-__version__ = "0.5.13"
+__version__ = "0.5.16"
 
 __all__ = [
     "VectorSearcher",
diff --git a/pyproject.toml b/pyproject.toml
index e92a37247..84d063f92 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 name = "superduper-framework"
 description = "🔮 Bring AI to your favourite database 🔮"
 readme = "README.md"
-version = '0.5.6'
+version = '0.5.9'
 license = {file = "LICENSE"}
 maintainers = [{name = "superduper.io, Inc.", email = "opensource@superduper.com"}]
 keywords = [
diff --git a/superduper/backends/query_dataset.py b/superduper/backends/query_dataset.py
index d98c6a3af..66342a11a 100644
--- a/superduper/backends/query_dataset.py
+++ b/superduper/backends/query_dataset.py
@@ -64,6 +64,14 @@ def __init__(
                 self._documents = list(
                     self.db.execute(self.select.select_using_ids(ids))
                 )
+                primary_id = self.select.id_field
+                # is it possible that the _documents came out in a different order?
+                # if so, resort them
+                lookup = {
+                    r[primary_id]: r for r in self._documents
+                }
+                self._documents = [lookup[id] for id in ids]
+                # DONE
         else:
             if ids is None:
                 self._ids = [
diff --git a/superduper/components/model.py b/superduper/components/model.py
index 3c08d6070..931356d03 100644
--- a/superduper/components/model.py
+++ b/superduper/components/model.py
@@ -573,7 +573,11 @@ def _prepare_inputs_from_select(
         X_data: t.Any
         mapping = Mapping(X, self.signature)
         if in_memory:
+            logging.info('Loading data into memory...')
             docs = list(self.db.execute(select.select_using_ids(ids)))
+            pid = select.table_or_collection.primary_id
+            lookup = {r[pid]: r for r in docs}
+            docs = [lookup[id_] for id_ in ids]
             X_data = list(map(lambda x: mapping(x), docs))
         else:
             assert isinstance(self.db, Datalayer)
diff --git a/superduper/rest/build.py b/superduper/rest/build.py
index fabb74480..fea407f9d 100644
--- a/superduper/rest/build.py
+++ b/superduper/rest/build.py
@@ -71,6 +71,10 @@ def build_rest_app(app: SuperDuperApp):
     """
     CFG.log_colorize = False
 
+    @app.add("/health", method="get")
+    def health():
+        return {"status": 200}
+
     @app.add("/handshake/config", method="post")
     def handshake(cfg: str):
         from superduper import CFG