Protect against silent error in root IDs with select_columns and timestamp

caveclient/materializationengine.py

@@ -3,17 +3,17 @@
 import logging
 import re
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from typing import Iterable, Optional, Union
-from requests import HTTPError
 
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pytz
-from concurrent.futures import ThreadPoolExecutor
 from cachetools import TTLCache, cached
 from IPython.display import HTML
+from requests import HTTPError
 
 from .auth import AuthClient
 from .base import (
@@ -136,6 +136,35 @@ def string_format_timestamp(ts):
         return ts
 
 
+def _check_select_columns_and_timestamp(
+    select_columns: Optional[Union[list, dict]], timestamp: Optional[datetime]
+) -> None:
+    if timestamp is not None and select_columns is not None:
+        if isinstance(select_columns, list):
+            select_values = select_columns
+        elif isinstance(select_columns, dict):
+            select_values = list(select_columns.values())
+        else:
+            return
+
+        has_root = False
+        has_supervoxel = False
+        for select_value in select_values:
+            if "root_id" in select_value:
+                has_root = True
+            if "supervoxel_id" in select_value:
+                has_supervoxel = True
+        if has_root and not has_supervoxel:
+            msg = (
+                "It looks like you are selecting a root ID column without "
+                "selecting the corresponding supervoxel ID column. This is likely "
+                "to cause issues with incorrect root IDs being returned for some "
+                "versions of the materialization engine. Please select the "
+                "corresponding supervoxel ID column as well."
+            )
+            raise ValueError(msg)
+
+
 def MaterializationClient(
     server_address,
     datastack_name=None,
@@ -258,7 +287,7 @@ def __init__(
     @property
     def datastack_name(self):
         return self._datastack_name
-    
+
     @property
     def cg_client(self):
         if self._cg_client is None:
@@ -296,7 +325,7 @@ def tables(self) -> TableManager:
             else:
                 raise ValueError("No full CAVEclient specified")
         return self._tables
-    
+
     @property
     def views(self) -> ViewManager:
         if self._views is None:
@@ -1513,6 +1542,8 @@ def live_query(
         if self.cg_client is None:
             raise ValueError("You must have a cg_client to run live_query")
 
+        _check_select_columns_and_timestamp(select_columns, timestamp)
+
         if datastack_name is None:
             datastack_name = self.datastack_name
         if desired_resolution is None:
@@ -1894,25 +1925,15 @@ def __init__(self, *args, **kwargs):
                     self.get_tables_metadata,
                 )
             )
-            metadata.append(
-                executor.submit(
-                    self.fc.schema.schema_definition_all
-                )
-            )
-            metadata.append(
-                executor.submit(
-                    self.get_views
-                )
-            )
-            metadata.append(
-                executor.submit(
-                    self.get_view_schemas
-                )
-            )
+            metadata.append(executor.submit(self.fc.schema.schema_definition_all))
+            metadata.append(executor.submit(self.get_views))
+            metadata.append(executor.submit(self.get_view_schemas))
         tables = None
         if self.fc is not None:
             if metadata[0].result() is not None and metadata[1].result() is not None:
-                tables = TableManager(self.fc, metadata[0].result(), metadata[1].result())
+                tables = TableManager(
+                    self.fc, metadata[0].result(), metadata[1].result()
+                )
         self._tables = tables
         if self.fc is not None:
             views = ViewManager(self.fc, metadata[2].result(), metadata[3].result())