Repository: Chia-Network/chia-blockchain
File: chia/data_layer/data_store.py (Python)

Code before:

                async with reader.execute(query, (store_id, *chunk)) as cursor:
                    rows = await cursor.fetchall()
                    result.update({row["kv_id"]: row["blob"] for row in rows})

        if len(result) != len(kv_ids):
            raise Exception("Cannot retrieve all the requested kv_ids")

        return result

    async def write_tree_to_file(
        self,
        root: Root,
        node_hash: bytes32,
        store_id: bytes32,
        deltas_only: bool,
        writer: BinaryIO,
    ) -> None:
        if node_hash == bytes32.zeros:
            return

        delta_file_cache = DeltaFileCache()
        with log_exceptions(log=log, message="Error while getting merkle blob"):
            root_path = self.get_merkle_path(store_id=store_id, root_hash=root.node_hash)
        delta_file_cache.load_hash_to_index(root_path)
        merkle_blob = await self.get_merkle_blob(store_id=store_id, root_hash=root.node_hash)

        if root.generation > 0:
            previous_root = await self.get_tree_root(store_id=store_id, generation=root.generation - 1)
            with log_exceptions(log=log, message="Error while getting previous merkle blob"):
                previous_root_path = self.get_merkle_path(store_id=store_id, root_hash=previous_root.node_hash)
            delta_file_cache.load_previous_hashes(previous_root_path)

        tree_nodes: list[SerializedNode] = []

        await self.get_nodes_for_file(root, node_hash, store_id, deltas_only, merkle_blob, delta_file_cache, tree_nodes)
        kv_ids = (
            KeyOrValueId.from_bytes(raw_id)
            for node in tree_nodes
            if node.is_terminal
            for raw_id in (node.value1, node.value2)
        )
        table_blobs = await self.get_table_blobs(kv_ids, store_id)

        for node in tree_nodes:
            if node.is_terminal:
                blobs = []
                for raw_id in (node.value1, node.value2):
                    id = KeyOrValueId.from_bytes(raw_id)
                    id_table_blob = table_blobs[id]
                    blobs.append(self.get_blob_from_table_blob(id_table_blob, store_id))

Review comment (performance):
i think this is loading the blob from disk twice?  i would have to look back, but i think i figured out how to pass a pyo3 object back into rust without overhead which would allow these to be re-ordered with the already loaded merkle blob being passed to the delta file cache.  this would still avoid asking the blob to create a python dict or whatever and then passing that back into a conversion back to rust.

Code after:
            await self.get_nodes_for_file(root, left_hash, store_id, deltas_only, delta_file_cache, tree_nodes)
            await self.get_nodes_for_file(root, right_hash, store_id, deltas_only, delta_file_cache, tree_nodes)
            tree_nodes.append(SerializedNode(False, bytes(left_hash), bytes(right_hash)))
        elif isinstance(raw_node, chia_rs.datalayer.LeafNode):
            tree_nodes.append(
                SerializedNode(
                    True,
                    raw_node.key.to_bytes(),
                    raw_node.value.to_bytes(),
                )
            )
        else:
            raise Exception(f"Node is neither InternalNode nor TerminalNode: {raw_node}")

    async def get_table_blobs(
        self, kv_ids_iter: Iterable[KeyOrValueId], store_id: bytes32
    ) -> dict[KeyOrValueId, tuple[bytes32, Optional[bytes]]]:
        result: dict[KeyOrValueId, tuple[bytes32, Optional[bytes]]] = {}
        batch_size = min(500, SQLITE_MAX_VARIABLE_NUMBER - 10)
        kv_ids = list(dict.fromkeys(kv_ids_iter))

        async with self.db_wrapper.reader() as reader:
            for i in range(0, len(kv_ids), batch_size):
                chunk = kv_ids[i : i + batch_size]
                placeholders = ",".join(["?"] * len(chunk))
                query = f"""
                    SELECT hash, blob, kv_id
                    FROM ids
                    WHERE store_id = ? AND kv_id IN ({placeholders})
                    LIMIT {len(chunk)}
                """

                async with reader.execute(query, (store_id, *chunk)) as cursor:
                    rows = await cursor.fetchall()
                    result.update({row["kv_id"]: (row["hash"], row["blob"]) for row in rows})

        if len(result) != len(kv_ids):
            raise Exception("Cannot retrieve all the requested kv_ids")

        return result

    async def write_tree_to_file(
        self,
        root: Root,
        node_hash: bytes32,
        store_id: bytes32,
        deltas_only: bool,
        writer: BinaryIO,
    ) -> None:
        if node_hash == bytes32.zeros:
            return