[User question] speed of a random access query

[sjfleming]

I was interested in testing out TileDB for some in-house single-cell data I've been working with. I created a TileDB datastore in a google cloud bucket using tiledbsoma.io.from_anndata (uploading lots of h5ad files after running tiledbsoma.io.register_anndatas.

I then tried querying data from this TileDB datastore backed by a google cloud bucket (3 million cells in there).

A query like this (10k cells)

    logger.info("starting quick check")
    inds = range(10000)
    with tiledbsoma.Experiment.open(tiledb_bucket_path) as exp:
        with exp.axis_query(
            measurement_name,
            obs_query=tiledbsoma.AxisQuery(coords=(inds,)),
        ) as query:
            adata = query.to_anndata(
                X_name=x_layer_name,
                column_names={"obs": ["soma_joinid"], "var": ["soma_joinid"]},
            )
            logger.info("quick check done")

ran in 30 seconds, and I was thrilled!

But as soon as I tried to query 10k random cell indices, I ran into a long delay:

    logger.info("starting quick shuffled check")
    inds = np.arange(3_000_000)
    inds_shuffled = np.random.permutation(inds)
    inds = [i for i in inds_shuffled[:10000]]
    with tiledbsoma.Experiment.open(tiledb_bucket_path) as exp:
        with exp.axis_query(
            measurement_name,
            obs_query=tiledbsoma.AxisQuery(coords=(inds,)),
        ) as query:
            adata = query.to_anndata(
                X_name=x_layer_name,
                column_names={"obs": ["soma_joinid"], "var": ["soma_joinid"]},
            )
            logger.info("quick check done")

The above took an hour to run.

Questions:

1. Am I doing something wrong / suboptimal above?
2. Is this kind of much longer query time for random access expected? Is it just part of TileDB, where a truly random query forces TileDB to open a ton of tiles, and so it's just gonna take a really long time?

Thanks!!

[johnkerl]

@sjfleming I ran the above for a couple experiments I have handy, one 2M cells and one 4M cells, both stored on S3, and the execution times were no more than three minutes for either. So we'll need to analyze a bit more -- we may have some questions for you about your schemas etc.

[spencerseale]

@sjfleming

What version are you using?

tiledbsoma.show_package_versions()

For your appending (loading your AnnDatas into the tiledbsoma.Experiment, was this done in parallel or in a sequential manner?

Have you assured that all of your obs_id_name are unique? These are the values we use when creating the registration mapping to tell SOMA where to fit in all of your new data being added into the dest tiledbsoma.Experiment.

[sjfleming]

Thanks so much for your responses!

@johnkerl sounds like if you ran the same test and they both ran in 3 mins, then I'm probably doing something bad. My data is stored in GCS rather than AWS... and I should say I'm just using my laptop to run that query, it's not like I'm running an enterprise version of TileDB in Kubernetes or something. This is home-brewed.

@spencerseale the version returns

tiledbsoma.__version__              1.15.0rc3
TileDB core version (libtiledbsoma) 2.26.2
python version                      3.10.16.final.0
OS version                          Darwin 23.6.0

I did the anndata uploads in a sequential manner (simply because I was having trouble coming up with a working parallelization strategy on my setup)... but I did run the tiledbsoma.io.register_anndatas function first.

The question about obs_id_name, I'm afraid I'll need to ask more: do you mean like adata.obs_names in the anndata files? Whether those are unique across all anndata files? Or does obs_id_name refer to something TileDB-specific (and how would I check to answer your question)?

[sjfleming]

In case you have the time and the inclination, this is the python code I used to create my TileDB datastore

"""Script to upload data from a bucket to a TileDB datastore in google cloud."""

import anndata
import click
from google.cloud import storage
import tiledbsoma.io
import os
import pandas as pd
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Generator

tiledbsoma.logging.info()


client = storage.Client()


def bucket_and_blob_from_path(cloud_path: str) -> tuple[str, str]:
    """Get the bucket and blob from a google cloud path."""
    assert cloud_path.startswith("gs://")
    bucket_name = cloud_path.split("/")[2]
    bucket = client.bucket(bucket_name)
    blob_name = "/".join(cloud_path.split("/")[3:])
    blob = bucket.blob(blob_name)
    return bucket, blob, bucket_name


def list_bucket_files(bucket_path: str, suffix: str = ".h5ad") -> list[str]:
    """Get a list of files in a bucket ending with a specific suffix."""
    bucket, _, bucket_name = bucket_and_blob_from_path(bucket_path)
    blobs = bucket.list_blobs()
    return [
        f"gs://{bucket_name}/{blob.name}"
        for blob in blobs
        if blob.name.endswith(suffix)
        and f"gs://{bucket_name}/{blob.name}".startswith(bucket_path)
    ]


def convert_to_boolean(df: pd.DataFrame) -> pd.DataFrame:
    """
    Convert columns with boolean-like values to categorical strings: 'TRUE', 'FALSE', or 'NA'.

    Args:
        df (pd.DataFrame): Input DataFrame.

    Returns:
        pd.DataFrame: Updated DataFrame with standardized string boolean values.

    By Jim Nemesh
    """
    for col in df.columns:
        col_non_na = df[col].dropna().astype(str)
        # Accept 'NA' as a missing value too (in string form)
        col_non_na_filtered = col_non_na[
            col_non_na.map(lambda x: str(x).strip().lower() != "na")
        ]
        is_boolean = col_non_na_filtered.map(
            lambda x: str(x).strip().lower() in ["true", "false", "1", "0"]
        ).all()
        if is_boolean:
            df[col] = df[col].astype(str)
            df.loc[:, col] = df[col].map(
                lambda x: (
                    "TRUE" if str(x).strip().lower() in ["true", "1"]
                    else "FALSE" if str(x).strip().lower() in ["false", "0"]
                    else "NA"
                )
        )
    return df.astype("category")


def float_nan_to_string_na(s: pd.Series) -> pd.Series:
    """
    Convert NaN values in a Series to 'NA' string.

    Args:
        s (pd.Series): Input Series.

    Returns:
        pd.Series: Updated Series with NaN values replaced by 'NA'.
    """
    return s.where(s.notna(), "NA").astype(str)


def anndata_from_h5ad(h5ad_local_file: str, obs_schema: pd.Series | None) -> anndata.AnnData:
    """Load an AnnData object from a h5ad file and remove uns, varm, and obsm,
    and convert boolean columns to a categorical [TRUE, FALSE, NA] of strings,
    and convert float NaN columns."""
    adata = anndata.read_h5ad(h5ad_local_file)
    if hasattr(adata, "uns"):
        del adata.uns
    if hasattr(adata, "varm"):
        del adata.varm
    if hasattr(adata, "obsm"):
        del adata.obsm
    adata.obs = convert_to_boolean(adata.obs)

    if obs_schema is not None:
    # for obs dtypes that do not adhere to schema (differences in categories are okay)
        for col, dtype in obs_schema[
            (obs_schema != adata.obs.dtypes) 
            & ((obs_schema != 'category') | (adata.obs.dtypes != 'category'))
        ].items():
            # try to convert to the schema type
            adata.obs[col] = float_nan_to_string_na(adata.obs[col])

    return adata


def ingest_initial_h5ad(
    h5ad_file: str,
    experiment_cloud_uri: str,
    measurement_name: str,
    x_layer_name: str,
    ingest_mode: str,
) -> pd.Series:
    """Ingest one h5ad file into a TileDB datastore."""
    _, blob, _ = bucket_and_blob_from_path(h5ad_file)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".h5ad") as temp_file:
        temp_file_path = temp_file.name
        blob.download_to_filename(temp_file_path)
        adata = anndata_from_h5ad(temp_file_path, obs_schema=None)

        tiledbsoma.io.from_anndata(
            experiment_cloud_uri,
            adata,
            measurement_name=measurement_name,
            X_layer_name=x_layer_name,
            ingest_mode=ingest_mode,
        )
    return adata.obs.dtypes


def anndata_generator(cloud_paths: list[str], obs_schema: pd.Series) -> Generator[anndata.AnnData, None, None]:
    for cloud_path in cloud_paths:
        print(cloud_path, flush=True)
        _, blob, _ = bucket_and_blob_from_path(cloud_path)

        with tempfile.NamedTemporaryFile(delete=False, suffix=".h5ad") as temp_file:
            temp_file_path = temp_file.name
            blob.download_to_filename(temp_file_path)
            adata = anndata_from_h5ad(temp_file_path, obs_schema=obs_schema)
            yield adata

        os.remove(temp_file_path)


def ingest_file_list(
    h5ad_bucket_path: str,
    tiledb_bucket_path: str,
    measurement_name: str = "RNA",
    x_layer_name: str = "data",
    ingest_mode: str = "write",
    max_workers: int = 2,
    var_index_if_named: str = "var_id",
) -> bool:
    """
    Ingest all h5ad files from a bucket path into a TileDB datastore.

    Args:
        h5ad_bucket_path: Bucket path containing all h5ad files for ingestion.
        tiledb_bucket_path: (Empty) bucket path for the new TileDB datastore.
        measurement_name: Name of the measurement.
        x_layer_name: Name of the count data internal to TileDB.
        ingest_mode: Ingestion mode for tiledb-soma, usually 'write'. Can use
            'resume' for a retry which will skip over existing data.
        max_workers: Number of parallel workers to use for ingestion.
        var_index_if_named: Name of the var index in the h5ads. "var_id" if unnamed.

    Returns:
        True if successful, False otherwise.
    """
    cloud_h5ads = list_bucket_files(h5ad_bucket_path)
    experiment_cloud_uri = tiledb_bucket_path

    if len(cloud_h5ads) == 0:
        raise ValueError(f"No h5ad files found in {h5ad_bucket_path}")

    # ingest the first file to establish the database
    print(f"Ingesting first file {cloud_h5ads[0]}", flush=True)
    obs_schema = ingest_initial_h5ad(
        cloud_h5ads[0],
        experiment_cloud_uri=experiment_cloud_uri,
        measurement_name=measurement_name,
        x_layer_name=x_layer_name,
        ingest_mode=ingest_mode,
    )

    if len(cloud_h5ads) == 1:
        return

    # register all of the h5ad files together
    # takes about 10 minutes per 100 h5ad files with 10k cells each, ~100MB files
    print(f"Registering {len(cloud_h5ads)} h5ad files", flush=True)
    rd = tiledbsoma.io.register_anndatas(
        experiment_cloud_uri,
        anndata_generator(cloud_h5ads, obs_schema=obs_schema),
        measurement_name=measurement_name,
        obs_field_name="obs_id",
        var_field_name=var_index_if_named,
    )

    def _ingest_h5ad(cloud_path: str):
        """Download h5ad to a temporary file and ingest it into the TileDB datastore."""
        # stagger the threads
        if cloud_path == cloud_h5ads[1]:
            time.sleep(150)

        print(f'{time.ctime()} -- {cloud_path}', flush=True)
        _, blob, _ = bucket_and_blob_from_path(cloud_path)

        with tempfile.NamedTemporaryFile(delete=True, suffix=".h5ad") as temp_file:
            temp_file_path = temp_file.name
            blob.download_to_filename(temp_file_path)
            adata = anndata_from_h5ad(temp_file_path, obs_schema=obs_schema)

            # ingest the h5ad file
            try:
                tiledbsoma.io.from_anndata(
                    experiment_cloud_uri,
                    adata,
                    measurement_name=measurement_name,
                    X_layer_name=x_layer_name,
                    ingest_mode=ingest_mode,
                    registration_mapping=rd,
                )
            except Exception as e:  # tiledbsoma._exception.SOMAError as e:
                print(f"Error ingesting {cloud_path}: {e}", flush=True)
                print(f"Retrying download and ingestion of {cloud_path}...", flush=True)
                blob.download_to_filename(temp_file_path)
                adata = anndata_from_h5ad(temp_file_path, obs_schema=obs_schema)
                tiledbsoma.io.from_anndata(
                    experiment_cloud_uri,
                    adata,
                    measurement_name=measurement_name,
                    X_layer_name=x_layer_name,
                    ingest_mode=ingest_mode,
                    registration_mapping=rd,
                )

    # run the remaining ingestions
    if max_workers == 1:
        print(f"Ingesting {len(cloud_h5ads) - 1} h5ad files sequentially", flush=True)
        for cloud_path in cloud_h5ads[1:]:
            _ingest_h5ad(cloud_path)
    else:
        print(f"Ingesting {len(cloud_h5ads) - 1} h5ad files in parallel", flush=True)
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            executor.map(_ingest_h5ad, cloud_h5ads[1:])

    return True


@click.command()
@click.argument("h5ad_bucket_path", type=str)
@click.argument("tiledb_bucket_path", type=str)
@click.option(
    "--measurement_name",
    type=str,
    default="RNA",
    help="Name of the measurement (default: 'RNA')",
)
@click.option(
    "--x_layer_name",
    type=str,
    default="data",
    help="Name of the count data internal to TileDB (default: 'data')",
)
@click.option(
    "--max_workers",
    type=int,
    default=2,
    help="Number of parallel worker threads to use for ingestion (default: 2)",
)
@click.option(
    "--var_index_if_named",
    type=str,
    default="var_id",
    help="If the var index is a named column, input it here. Default for an unnamed index is 'var_id'.",
)
def main(
    h5ad_bucket_path: str,
    tiledb_bucket_path: str,
    measurement_name: str,
    x_layer_name: str,
    max_workers: int,
    var_index_if_named: str,
):
    """
    This script ingests all h5ad files from a specified google bucket path H5AD_BUCKET_PATH
    into a TileDB datastore in google cloud at TILEDB_BUCKET_PATH.

    \b
    Arguments:
        h5ad_bucket_path (str): Bucket path containing all h5ad files for ingestion.
        tiledb_bucket_path (str): (Empty) bucket path for the new TileDB datastore.
        max_workers (int): Number of parallel workers to use for ingestion.
        var_index_if_named (str): If the var index in the h5ads is a named column, input it here.
            Default for an unnamed index is 'var_id'.
    \b
    Example:
        python tiledb_upload_script.py gs://my-bucket/h5ad_extracts/curriculum gs://my-other-bucket/tiledb-datastore
    """
    if max_workers < 1:
        raise ValueError("max_workers must be >= 1. 1 will run a sequential for-loop.")

    succeeded = ingest_file_list(
        h5ad_bucket_path=h5ad_bucket_path,
        tiledb_bucket_path=tiledb_bucket_path,
        measurement_name=measurement_name,
        x_layer_name=x_layer_name,
        ingest_mode="write",
        max_workers=max_workers,
        var_index_if_named=var_index_if_named,
    )
    if succeeded:
        print("Ingestion completed successfully.")
    else:
        print("Ingestion incomplete.")


if __name__ == "__main__":
    main()

I ran this using --max_workers 1 (which runs upload in series) due to (if I remember right) memory overflows on my machine when trying to have several workers run in parallel.

[spencerseale]

@sjfleming:

If you're able to bump you version to the latest stable release, we're on 1.16.1.

Yes similar to AnnData.obs_names. obs_id_name is the value we use to identify unique values in the input data and the existing tiledbsoma.Experiment source. SOMA basically says, what obs_id_name values currently exist in the tiledbsoma.Experiment and what obs_id_name values are needing to be added. It then generates the soma_joinid dim values based on those values. soma_joinids are a proxy for the obs and var IDs that are more performant over cloud storage.

So when you tiledbsoma.io.from_anndata or any other of the related operations, make sure you're properly wiring up those values by including them as kwargs just as you did when you ran tiledbsoma.io.register_anndatas. This may not be impacting you, but just wanted to clarify.

Additionally, you can do a test to make sure all of your obs IDs are truly unique, by first adding a uuid.uuid4() suffix to each obs ID and that will make sure all of those IDs (your cell barcode or whatever is the cell identifier) is truly unique!

Also in ingest_initial_h5ad, we have a "schema_only" mode you could use to instantiate your tiledbsoma.Experiment prior to loading any data.

You'll need to add a resize step after you create the registration mapping. We have a tutorial here you can checkout that shows the order of operations: https://cloud.tiledb.com/academy/structure/life-sciences/single-cell/tutorials/append-data/

[sjfleming]

Thanks for your message! I did check the obs IDs to make sure they are unique, and they are.

I see your other thoughts, but I guess I'm not sure what I really need to try and what should work as is. Would you expect the script above to work (it does run) as is, if the obs IDs in the h5ad files are all globally unique?

[spencerseale]

@sjfleming without trying it directly, it looks good to me. You can transform the soma into an AnnData and compare to your original to validate it.