Hi there- an API question for you. I’m working on a pipeline to generate input data for Apple’s new

embedding-atlas

library. At the end, I need to run UMAP against the full set of embeddings. I currently resort to the clunky approach below. Am I missing a better way to add a column from a numpy array of same length as the dataframe?

vecs = np.stack([r["embedding"] for r in df.select("embedding").to_pylist()], axis=0)

# little helper wrapping umap
xy, knn_indices, knn_distances = umap_with_neighbors(vecs).values()

df = df.with_column("_row_index": (daft.sql_expr("row_number()") - 1))

umap_cols = daft.from_pydict({
    "_row_index": list(range(xy.shape[0])),
    "umap_x": daft.Series.from_arrow(pa.array(xy[:,0])),
    "umap_y": daft.Series.from_arrow(pa.array(xy[:,1])),
})
 
df = df.join(umap_cols, on="_row_index")

I suggest using the

.with_column

method on a

daft.DataFrame

to add a new column. If I assume you have

df: daft.DataFrame

loaded up from somewhere, and assuming

umap_with_neighbors: np.ndarray -> tuple[np.ndarray, np.ndarray,np.ndarray]

then you'd want to do something like:

from typing import TypedDict

import daft
import numpy as np
from numpy.typing import NDArray

class UmapReturn(TypedDict):
  xy: NDArray[float]
  knn_indices: NDArray[int]
  knn_distances: NDArray[float]


@daft.func.batch(return_dtype=UmapReturn)
def umap_udf(vecs: daft.Series):
    xy, knn_indices, knn_distances = umap_with_neighbors(vecs).values()
    return [UmapReturn(xy, knn_indices, knn_distances)]

And use it like:

df: daft.DataFrame = ...

df = df.with_column("umap_results", umap_udf(daft.col("NAME OF COLUMN WITH BATCHES OF VECTORS HERE"))(

df = df.with_column("xy", daft.col("umap_results").get("xy"))

df = df.with_column("knn_indices", daft.col("umap_results").get("knn_indices"))

df = df.with_column("knn_distances", daft.col("umap_results").get("knn_distances"))

df = df.select(<OTHER COLUMNS HERE>, "xy", "knn_indices", "knn_distances")

Thanks @Malcolm Greaves!

umap_with_neighbors

is intended to be run against all rows at once. (It will give non-useful results if run in batches.) For this, should I just pass a

batch_size

greater than the number of rows? (I suppose

None

will disable batching i.e. same behavior as

batch_size = 1

?)

Hmmmm, this may be hearsay 😬 but this may not be a good use case for Daft. If your umap implementation assumes everything needs to exist in memory and needs to be able to view all embeddings at once for it to work, then it's kind of outside of what Daft is good for / the assumptions Daft makes. Daft assumes that you're doing things at the record level, or batch of records. For dataset-sized things, Daft has operations like groupby, sort, and joins. But those are still about individual records. To do umap on the whole dataset, you'd have to implement the actual umap algorithm in terms of Daft's operators.

Sorry to wake the dead on this thread, but this makes sense as a non-goal / less-needful-use-case for daft, and thanks for confirming!