diff --git a/tiledb/consolidation_plan.py b/tiledb/consolidation_plan.py index 98758ca5ba..a6343ae36e 100644 --- a/tiledb/consolidation_plan.py +++ b/tiledb/consolidation_plan.py @@ -1,3 +1,5 @@ +import pprint + import tiledb.cc as lt from .ctx import Ctx, CtxMixin, default_ctx @@ -25,6 +27,59 @@ def __init__(self, ctx: Ctx, array: lt.Array, fragment_size: int): super().__init__(ctx, lt.Array(ctx, array), fragment_size) + def __len__(self): + """Returns the number of nodes in the consolidation plan""" + return self.num_nodes + + def __repr__(self): + attrs = { + "num_nodes": self.num_nodes, + "fragments": { + f"node_{node_idx}": { + "num_fragments": self.num_fragments(node_idx), + "fragment_uris": [ + self.fragment_uri(node_idx, fragment_idx) + for fragment_idx in range(self.num_fragments(node_idx)) + ], + } + for node_idx in range(self.num_nodes) + }, + } + + return pprint.PrettyPrinter().pformat(attrs) + + def _repr_html_(self): + from io import StringIO + + output = StringIO() + output.write("
\n") + + output.write("

Consolidation Plan

\n") + output.write("\n") + output.write( + "\n" + ) + for node_idx in range(self.num_nodes): + output.write( + f"\n" + ) + output.write("
NodeNum FragmentsFragment URIs
{node_idx}{self.num_fragments(node_idx)}{', '.join(self.fragment_uri(node_idx, fragment_idx) for fragment_idx in range(self.num_fragments(node_idx)))}
\n") + + output.write("
\n") + return output.getvalue() + + def __getitem__(self, idx): + if idx < 0 or idx >= self.num_nodes: + raise IndexError("Index out of bounds") + + return { + "num_fragments": self.num_fragments(idx), + "fragment_uris": [ + self.fragment_uri(idx, fragment_idx) + for fragment_idx in range(self.num_fragments(idx)) + ], + } + @property def num_nodes(self) -> int: """ diff --git a/tiledb/tests/test_consolidation_plan.py b/tiledb/tests/test_consolidation_plan.py index aee4f7b16c..b82190024f 100644 --- a/tiledb/tests/test_consolidation_plan.py +++ b/tiledb/tests/test_consolidation_plan.py @@ -1,6 +1,8 @@ import json +import xml import numpy as np +import pytest import tiledb from tiledb.tests.common import DiskTestCase @@ -16,15 +18,32 @@ def test_consolidation_plan(self): with tiledb.open(path, "r") as A: cons_plan = tiledb.ConsolidationPlan(tiledb.default_ctx(), A, 2) assert cons_plan.num_nodes == 1 + assert cons_plan.num_nodes == len(cons_plan) assert cons_plan.num_fragments(0) == 1 # check that it has a nodes key assert "nodes" in json.loads(cons_plan.dump()) # check that it has a list of nodes assert isinstance(json.loads(cons_plan.dump())["nodes"], list) # check that each node has a uri key - assert "uri" in json.loads(cons_plan.dump())["nodes"][0]["uris"][0] - - # write a second fragment to the array + for node in json.loads(cons_plan.dump())["nodes"]: + assert "uri" in node["uris"][0] + # test __repr__ + try: + assert ( + xml.etree.ElementTree.fromstring(cons_plan._repr_html_()) + is not None + ) + except: + pytest.fail( + f"Could not parse cons_plan._repr_html_(). Saw {cons_plan._repr_html_()}" + ) + # test __getitem__ + assert cons_plan[0] == { + "num_fragments": 1, + "fragment_uris": [cons_plan.fragment_uri(0, 0)], + } + + # write a second fragment to the array and check the new consolidation plan with tiledb.open(path, "w") as A: A[:] = np.random.rand(4) @@ -32,9 +51,3 @@ def test_consolidation_plan(self): cons_plan = tiledb.ConsolidationPlan(tiledb.default_ctx(), A, 4) assert cons_plan.num_nodes == 1 assert cons_plan.num_fragments(0) == 2 - # check that it has a nodes key - assert "nodes" in json.loads(cons_plan.dump()) - # check that it has a list of nodes - assert isinstance(json.loads(cons_plan.dump())["nodes"], list) - # check that each node has a uri key - assert "uri" in json.loads(cons_plan.dump())["nodes"][0]["uris"][0]