Add support for TileDB-Py 0.33

Code was using internal function which no longer exists.
TileDB-Inc · Feb 19, 2025 · 9c8cdcf · 9c8cdcf
1 parent e9eb15f
commit 9c8cdcf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/tiledb/ml/readers/_tensor_schema/base_sparse.py b/tiledb/ml/readers/_tensor_schema/base_sparse.py
@@ -36,7 +36,9 @@ def key_range(self) -> WeightedRange[Any, int]:
     @property
     def max_partition_weight(self) -> int:
         try:
-            memory_budget = int(self._array._ctx_().config()["py.init_buffer_bytes"])
+            # getattr for compatibility with TileDB-Py <0.33
+            ctx = getattr(self._array, "ctx", self._array._ctx_())
+            memory_budget = int(ctx.config()["py.init_buffer_bytes"])
         except KeyError:
             memory_budget = 10 * 1024**2
 

diff --git a/tiledb/ml/readers/_tensor_schema/dense.py b/tiledb/ml/readers/_tensor_schema/dense.py
@@ -66,7 +66,9 @@ def iter_tensors(
 
     @property
     def max_partition_weight(self) -> int:
-        memory_budget = int(self._array._ctx_().config()["sm.mem.total_budget"])
+        # getattr for compatibility with TileDB-Py <0.33
+        ctx = getattr(self._array, "ctx", self._array._ctx_())
+        memory_budget = int(ctx.config()["sm.mem.total_budget"])
 
         # The memory budget should be large enough to read the cells of the largest field
         bytes_per_cell = max(dtype.itemsize for dtype in self.field_dtypes)