Repository: activeloopai/Hub
File: deeplake/core/vectorstore/vector_search/dataset/dataset.py (Python)

Code before:

    return processed_tensors, tensors[ids_tensor]


def read_tensor_data_if_needed(tensor_data, dataset, tensor_name):
    # generalize this method for other htypes that need reading.
    if dataset and tensor_name != "id" and dataset[tensor_name].htype == "image":
        tensor_data = [
            deeplake.read(data) if isinstance(data, str) else data
            for data in tensor_data
        ]
    return tensor_data


def convert_tensor_data_to_list(tensor_data, tensors, ids_tensor):
    if tensor_data is None:
        tensor_data = [None] * len(tensors[ids_tensor])
    elif not isinstance(tensor_data, list):
        tensor_data = list(tensor_data)
    return tensor_data


def get_not_none_tensors(tensors, embedding_data):
    not_non_tensors = {k: v for k, v in tensors.items() if v is not None}
    try:
        num_items = len(next(iter(not_non_tensors.values())))
    except StopIteration:
        if embedding_data:
            num_items = len(embedding_data[0])
        else:
            num_items = 0
    return not_non_tensors, num_items


def populate_id_tensor_if_needed(ids_tensor, tensors, not_non_tensors, num_items):
    if "id" not in not_non_tensors and "ids" not in not_non_tensors:
        id = [str(uuid.uuid1()) for _ in range(num_items)]
        tensors[ids_tensor] = id
    else:
        for tensor in not_non_tensors:
            if tensor in ("id", "ids"):
                break

        tensors[ids_tensor] = not_non_tensors[tensor]
    return tensors


def get_id_tensor(dataset):
    return "ids" if "ids" in dataset.tensors else "id"



Review comment (suggestion):
```suggestion
        num_items = len(next(iter(not_none_tensors.values())))
```

Code after:

    return processed_tensors, tensors[ids_tensor]


def read_tensor_data_if_needed(tensor_data, dataset, tensor_name):
    # generalize this method for other htypes that need reading.
    if dataset and tensor_name != "id" and dataset[tensor_name].htype == "image":
        tensor_data = [
            deeplake.read(data) if isinstance(data, str) else data
            for data in tensor_data
        ]
    return tensor_data


def convert_tensor_data_to_list(tensor_data, tensors, ids_tensor):
    if tensor_data is None:
        tensor_data = [None] * len(tensors[ids_tensor])
    elif not isinstance(tensor_data, list):
        tensor_data = list(tensor_data)
    return tensor_data


def get_not_none_tensors(tensors, embedding_data):
    not_none_tensors = {k: v for k, v in tensors.items() if v is not None}
    try:
        num_items = len(next(iter(not_none_tensors.values())))
    except StopIteration:
        if embedding_data:
            num_items = len(embedding_data[0])
        else:
            num_items = 0
    return not_none_tensors, num_items


def populate_id_tensor_if_needed(ids_tensor, tensors, not_none_tensors, num_items):
    if "id" not in not_non_tensors and "ids" not in not_non_tensors:
        id = [str(uuid.uuid1()) for _ in range(num_items)]
        tensors[ids_tensor] = id
    else:
        for tensor in not_non_tensors:
            if tensor in ("id", "ids"):
                break

        tensors[ids_tensor] = not_non_tensors[tensor]
    return tensors


def get_id_tensor(dataset):
    return "ids" if "ids" in dataset.tensors else "id"