Is there a nice way to marshal a sparse tensor into a numpy array or a coo array in a UDF? It seems like

to_pylist

and

to_arrow

give you a format which is a list of maps.

cc @Sagi @MichaelV who contributed sparse tensor support! I think we just haven’t really defined the expected behavior yet. Feel free to propose something 🙂

Do you have a preferred place for me to write that up? Also happy to write down my thoughts here.

An issue could be a good medium for this!

Cool I'll create one when I get a chance today.

Hi, for native sparse matrix implementation, we’d need to use something like scipy or pytorch, but I didn’t want to add those dependencies to daft just for casting. I also thought about creating a custom numpy sparse dtype, but keeping the coo components in a simple python dict felt cleaner and simpler to understand. Each component (data, indices of non zero elements) is just stored as a numpy array. We can always use any coo constructor to reconstruct the flattened list of values and indices, then reshape it to the original size. If you have any suggestion we would love to hear! :)

@Sagi that makes sense. Will think about it a bit more, I understand the concern of not wanting to bloat the deps. > btw I’d recommend using daft to convert sparse tensor columns to dense tensors when possible since it’s already supported! This is where I landed for converting from the raw data and it works well for serializing to a smaller footprint. I did have a question about going back to dense though for the purposes of inference. If I do that in a UDF call like this:

return experiment_table.with_column(
        "predictions",
        VITBFullNeighborhood.with_concurrency(concurrency)(
            daft.col("positions"),
            daft.col("expressions").cast(DT.tensor(DT.float32(), shape=(1000, 1))),
            daft.col("masks"),
            daft.col("input_expression"),
            daft.col("input_mask"),
        ),
    )

will that lazily cast partitions or batches as they get fed into the UDF?

We “could” do a check like

if _IS_PYTORCH_IMPORTABLE: …

in our code to dynamically produce certain types. The question is whether or not we’d want to, because it is arguably a bit unexpected in terms of behavior. Perhaps this is more of an API problem where

.to_pylist()

should be able to take in certain arguments. Or maybe users should only use an explicit

.to_pytorch()

and

.to_numpy()

function.

The behavior or

to_numpy()

was the most surprising for me. I can understand

to_pylist

returning a list of maps but I wasn't expecting

to_numpy

to do the same.