Modify check for ctx between versions

TileDB-Inc · Feb 19, 2025 · 6c69c48 · 6c69c48
1 parent 9c8cdcf
commit 6c69c48
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/tiledb/ml/readers/_tensor_schema/base_sparse.py b/tiledb/ml/readers/_tensor_schema/base_sparse.py
@@ -38,6 +38,12 @@ def max_partition_weight(self) -> int:
         try:
             # getattr for compatibility with TileDB-Py <0.33
             ctx = getattr(self._array, "ctx", self._array._ctx_())
+            if hasattr(self._array, "ctx"):  # Check if 'ctx' exists (TileDB-Py >= 0.33)
+                ctx = self._array.ctx
+            elif hasattr(self._array, "_ctx_"):  # Check if '_ctx_' exists (TileDB-Py < 0.33)
+                ctx = self._array._ctx_()  # Note the parentheses - it was a method
+            else:
+                raise AttributeError("TileDB context not found. Incompatible TileDB-Py version?")
             memory_budget = int(ctx.config()["py.init_buffer_bytes"])
         except KeyError:
             memory_budget = 10 * 1024**2