[Feature request] docs on chunked reading for large databases?

[nick-youngblut]

Is your feature request related to a problem? Please describe.

For large tiledb-soma databases hosted on GCP or AWS, read of the obs dataframe can be quite slow (minutes) and require a substantial amount of memory.

Describe the solution you'd like

I've implemented the following chunked-read approach:

from typing import Optional, List
import pandas as pd
import tiledbsoma
import tiledb
import pyarrow as pa

def get_obs(
        db_uri: str,
        columns: Optional[List[str]] = None,
        obs_query: Optional[tiledbsoma.AxisQuery] = None,
        axis: str = "RNA",
        group_by: Optional[List[str]] = None,
        agg_name: str = "count_all"
    ) -> pd.DataFrame:
    """
    Retrieve the obs table from a tiledbsoma Experiment.
    
    If group_by is provided, aggregated counts by the specified grouping columns are returned.
    Otherwise, the full obs subset with the specified columns is returned.
    
    Parameters:
      db_uri: URI for the database.
      columns: List of column names to retrieve (required if group_by is None).
      obs_query: Optional AxisQuery to filter the obs data.
      axis: The axis to query (default "RNA").
      group_by: Optional list of columns to group by for aggregation.
      agg_name: Name for the aggregated count column (default "count_all").
    
    Returns:
      A pandas DataFrame containing either the subset of obs data or the aggregated counts.
    """
    if group_by is not None:
        # When grouping, we only need the group_by columns.
        group_cols = group_by
        chunk_results: List[pd.DataFrame] = []
        with tiledbsoma.Experiment.open(db_uri) as exp:
            # Use axis_query if provided.
            if obs_query is not None:
                try:
                    reader = exp.axis_query(axis, obs_query=obs_query).obs(column_names=group_cols)
                except TypeError:
                    reader = exp.axis_query(axis, obs_query=obs_query).obs()
                chunk_iter = reader
            else:
                chunk_iter = exp.obs.read(column_names=group_cols)
            # Process each chunk.
            for chunk in chunk_iter:
                if chunk.num_rows == 0:
                    continue
                df_chunk = chunk.to_pandas()
                # Group and count within this chunk.
                df_group = df_chunk.groupby(group_cols, as_index=False).size().rename(columns={"size": agg_name})
                chunk_results.append(df_group)
        if not chunk_results:
            return pd.DataFrame(columns=group_cols + [agg_name])
        # Combine chunk results and re-aggregate to get overall counts.
        df_all = pd.concat(chunk_results, ignore_index=True)
        df_final = df_all.groupby(group_cols, as_index=False)[agg_name].sum()
        return df_final.sort_values(by=agg_name, ascending=False)
    else:
        # Without grouping, ensure columns is provided.
        if columns is None:
            raise ValueError("When group_by is not specified, the 'columns' parameter must be provided.")
        chunks: List[pd.DataFrame] = []
        with tiledbsoma.Experiment.open(db_uri) as exp:
            if obs_query is not None:
                try:
                    reader = exp.axis_query(axis, obs_query=obs_query).obs(column_names=columns)
                except TypeError:
                    reader = exp.axis_query(axis, obs_query=obs_query).obs()
                chunk_iter = reader
            else:
                chunk_iter = exp.obs.read(column_names=columns)
            for chunk in chunk_iter:
                if chunk.num_rows == 0:
                    continue
                df_chunk = chunk.to_pandas()[columns]
                chunks.append(df_chunk)
        if not chunks:
            return pd.DataFrame(columns=columns)
        return pd.concat(chunks, ignore_index=True)

Which seems to substantially reduce db query times and memory. Maybe something similar can be implemented in the package? Maybe it is, and I've just missed that code?

At least in the docs, it might be helpful to point out alternatives to using .concat as in:

obs = experiment.obs
table = obs.read().concat()
table

...for large databases.

[nick-youngblut]

Similar to the issue above of poor performance for large databases, querying and converting to an anndata is very slow via:

obs_query = tiledbsoma.AxisQuery(value_filter='sample in ["sample1", "sample2"]')
with tiledbsoma.Experiment.open(db_uri) as exp:
    adata = exp.axis_query("RNA", obs_query=obs_query).to_anndata(X_name="data")
adata

I've created the following alternative, but it still can be slow if there are many contiguous blocks to slice:

def group_contiguous(sorted_ids: List[int]) -> List[Tuple[int, int, List[int]]]:
    """
    Group a sorted list of integers into contiguous blocks.
    Each block is returned as a tuple (start, end, group_ids),
    where group_ids is the list of consecutive integers, and the block
    covers indices from start to end-1.
    """
    groups = []
    if not sorted_ids:
        return groups
    current = [sorted_ids[0]]
    for i in sorted_ids[1:]:
        if i == current[-1] + 1:
            current.append(i)
        else:
            groups.append((current[0], current[-1] + 1, current.copy()))
            current = [i]
    groups.append((current[0], current[-1] + 1, current.copy()))
    return groups

def read_block(ms_array, block_range: Tuple[int, int], group_ids: List[int]) -> sp.csr_matrix:
    """
    Read a contiguous slice of the measurement array and extract only the rows
    corresponding to group_ids.
    
    The TileDB array returns a full matrix with data in the original row positions,
    so we use the original group_ids for indexing rather than computing relative indices.
    """
    start, end = block_range
    
    # Read and convert to sparse matrix
    block_tensor = ms_array.read(coords=(slice(start, end), slice(None)))
    block_sparse = block_tensor.coos().concat().to_scipy().tocsr()
    
    # Extract rows using original group_ids
    return block_sparse[group_ids, :]

def get_anndata(
    db_uri: str,
    obs_query: Optional[tiledbsoma.AxisQuery] = None,
    measurement_name: str = "RNA",
    X_name: str = "X",
    obs_columns: Optional[List[str]] = None,
    max_workers: int = 12
) -> anndata.AnnData:
    """
    Retrieve a subset of an experiment as an AnnData object, filtering the measurement
    (X) on-disk by reading only the rows corresponding to the obs query's 'soma_joinid' values.
    
    In addition, the function reads the var DataFrame from the measurement so that
    the returned AnnData object contains full feature metadata.
    
    Steps:
      1. Read obs metadata (optionally in chunks) and filter it via obs_query.
      2. Extract the 'soma_joinid' values from obs and group them into contiguous blocks.
      3. Use a ThreadPoolExecutor to read each contiguous block (via slicing) concurrently.
      4. Vertically stack the resulting blocks and reorder rows to match the original obs order.
      5. Read the var DataFrame from the measurement.
      6. Construct an AnnData object with obs, var, and the sparse X matrix.
    
    Parameters:
      db_uri (str): URI of the TileDB-SOMA experiment.
      obs_query (Optional[tiledbsoma.AxisQuery]): Query to filter observations.
      measurement_name (str): Name of the measurement (e.g. "RNA").
      X_name (str): Name of the measurement matrix (e.g. "X").
      obs_columns (Optional[List[str]]): List of columns to retrieve from obs.
      chunk_obs (bool): Whether to read obs metadata in chunks.
      max_workers (int): Maximum number of threads for parallel block reads.
    
    Returns:
      anndata.AnnData: The filtered AnnData object with a sparse X and complete var metadata.
    """
    import sys
    print("Reading obs metadata...", file=sys.stderr)
    # --- Read obs metadata ---
    if obs_columns is None:
        obs_columns = ["soma_joinid"]
    obs_df = get_obs(db_uri, columns=obs_columns, obs_query=obs_query, axis=measurement_name)
    
    # --- Ensure 'soma_joinid' exists and extract join IDs ---
    if "soma_joinid" not in obs_df.columns:
        raise ValueError("The obs metadata must include the 'soma_joinid' column.")
    required_ids = obs_df["soma_joinid"].tolist()
    sorted_ids = sorted(required_ids)
    groups = group_contiguous(sorted_ids)
    print(f"Found {len(groups)} contiguous blocks.", file=sys.stderr)
    
    print("Reading measurement data by contiguous blocks in parallel...", file=sys.stderr)
    block_list = []
    with tiledbsoma.Experiment.open(db_uri) as exp:
        ms_obj = exp.ms[measurement_name]
        if X_name not in ms_obj:
            available = list(ms_obj.keys())
            raise ValueError(f"Measurement '{measurement_name}' does not contain layer '{X_name}'. "
                             f"Available layers: {available}")
        ms_layer = ms_obj[X_name]
        if isinstance(ms_layer, tiledbsoma.Collection):
            keys = list(ms_layer.keys())
            ms_array = ms_layer["data"] if "data" in keys else ms_layer[keys[0]]
        else:
            ms_array = ms_layer
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_block = {
                executor.submit(read_block, ms_array, (start, end), group_ids): (start, end, group_ids)
                for (start, end, group_ids) in groups
            }
            for future in concurrent.futures.as_completed(future_to_block):
                block_info = future_to_block[future]
                try:
                    block_matrix = future.result()
                    #print(block_matrix); return None
                    block_list.append((block_info[2], block_matrix))
                    if len(block_list) % 10 == 0:
                        print(f"  Processed {len(block_list)} blocks.", file=sys.stderr)
                except Exception as e:
                    raise ValueError(f"Error reading block {block_info[0]} to {block_info[1]}") from e
    
    if not block_list:
        raise ValueError("No measurement data was read.")
    
    print("Stacking blocks...", file=sys.stderr)
    block_list.sort(key=lambda x: x[0][0])  
    blocks = [blk for (_, blk) in block_list]
    # Now blocks are (cells x genes), verify dimensions match
    n_genes = blocks[0].shape[1]  # Check number of genes (columns) instead
    for i, block in enumerate(blocks):
        if block.shape[1] != n_genes:
            raise ValueError(f"Block {i} has {block.shape[1]} genes but expected {n_genes}")
    # Stack vertically since blocks are already (cells x genes)
    X_filtered = sp.vstack(blocks).tocsr() 
    
    print("Reading var metadata...", file=sys.stderr)
    with tiledbsoma.Experiment.open(db_uri) as exp:
        var_reader = exp.ms[measurement_name].var.read()
        var_df = var_reader.concat().to_pandas()
    
    print("Creating AnnData object...", file=sys.stderr)
    adata = anndata.AnnData(X_filtered, obs=obs_df, var=var_df)
    return adata

# Example usage:
obs_query = tiledbsoma.AxisQuery(value_filter='sample in ["sample1", "sample2"]')
adata = get_anndata(db_uri, obs_query=obs_query, measurement_name="RNA", X_name="X")
adata

Any suggestions on improving performance?

[johnkerl]

Thanks @nick-youngblut ! Taking a look ...

[nick-youngblut]

FYI: get_anndata often crashes my Jupyter kernel (~50% of the time), and I'm not sure why.

[johnkerl]

@nick-youngblut do you see Killed, or some other message? If you see no error message at all -- if it's not too much to ask -- would it be possible for you to run the same code you're running in a notebook, but, with code cells copy/pasted into a Jupyter Terminal session? My experience is that a terminal session will often show useful error messages that a notebook session can hide.

[bkmartinjr]

Hi @nick-youngblut. Any chance you could share:

• the dataset you are querying - if it is public (e.g., CZI Census?)? If not public, could you share schema and how you created the SOMA Experiment (I need the low-level TileDB layout info, etc)?
• specific cases that are most problematic (e.g., is it fetching the entire obs, or a specific value_filter? all columns or a subset? etc)

Most of the underlying ExperimentAxisQuery code performs chunked reads/queries (at a layer lower than this class). I'd like to reproduce the performance problem, and and then dig into the root cause. That should either lead to a fix, or some config/param suggestions.

We can also move to slack or email as well if easier to iterate.

[nick-youngblut]

the dataset you are querying

It is a private database. Thanks @bkmartinjr for offering to help! It would be great to figure out the performance issues. For now, we are falling back to h5ad files + parquet for quick access to all of the metadata.

Maybe email is best: nick.youngblut@arcinstitute.org

[nick-youngblut]

would it be possible for you to run the same code you're running in a notebook, but, with code cells copy/pasted into a Jupyter Terminal session? My experience is that a terminal session will often show useful error messages that a notebook session can hide.

@johnkerl here's the error when running the code from the terminal:

Unexpected error 107 on netlink descriptor 234 (address family 10).
Aborted (core dumped)

Thanks for your help with figuring this out!