Enable UMAP to properly handle sparse data

[rishic3]

Addressing bug found by QA, where UMAP did not properly handle sparse inputs for 'jaccard' metric.

Changes:

Handling sparse data:

• Convert UMAP input data to CSR matrix in case of SparseVectors by inheriting enable_sparse_data_optim code.
• Return and store "raw_data" attribute in dict of lists containing the CSR attributes after fit.
• Chunk and broadcast each CSR attribute independently, store result within dict.
• Fixed bug where UMAP was inheriting from the wrong CumlModel and was returning a new dataframe, rather than adding columns to the input dataframe.

Reading/writing:

• Added logic to save/load CSR array components.
• Refactored to save model attributes as single compressed .npz file for all chunks (rather than separate .npy files) - requires compression time but less disk space.

Testing:

• Added test for sparse inputs.
• Added test for model persistence in sparse case, including when data is saved/loaded in chunks.

Todos:

These features can be targeted for 24.12:

• Extend SparseDataGen to generate binary sparse dataset for testing jaccard metric (can be reused for exact/approx KNN)
• Refactor UMAP model persistence to enable saving/loading arrays to cloud storage (e.g., put arrays into a spark dataframe and use Spark dataframe writer)
• Chunking by bytes/row currently rests on the assumption of relatively uniform memory distribution - add better handling of skewed data.

[rishic3]

build

[rishic3]

build

[rishic3]

build

[rishic3]

build

[rishic3]

Removing since it looks resolved: NVIDIA/spark-rapids#10770

[eordentlich]

Good find.

[lijinf2]

What is the range of trustworthiness score? Would the tolerance 0.15 be too large, given the ci actually runs single-GPU umap.

[rishic3]

Re: discussion on how to avoid chunking twice - chunking when returning from fit is by # rows, and is meant to help conform with maxRecordsPerBatch so that Spark only loads a batch at a time into memory, whereas chunking on the driver side is based on # bytes to stay below the 8GB broadcast limit.

During fit we can calculate the tighter of the two restrictions prior to yielding to the driver and then reuse these chunks during broadcast. Though if the maxRecordsPerBatch chunk is significantly smaller in bytes, we'd have the overhead of unnecessarily broadcasting lots of small chunks. Not sure if there's a fix that makes sense here.

[rishic3]

Re: discussion on how to avoid chunking twice - chunking when returning from fit is by # rows, and is meant to help conform with maxRecordsPerBatch so that Spark only loads a batch at a time into memory, whereas chunking on the driver side is based on # bytes to stay below the 8GB broadcast limit.

During fit we can calculate the tighter of the two restrictions prior to yielding to the driver and then reuse these chunks during broadcast. Though if the maxRecordsPerBatch chunk is significantly smaller in bytes, we'd have the overhead of unnecessarily broadcasting lots of small chunks. Not sure if there's a fix that makes sense here.

Also, chunking the CSR matrix by # rows when returning from _fit poses separate challenges since the CSR components and the embedding matrix will have different # rows. It's unclear the number of chunks we should return here.
If we could chunk the CSR matrix directly and reconcatenate using scipy.sparse.vstack that would be ideal, but there's no easy way to get the CSR matrix into a serializable format (i.e., python lists) unless we extract the components. Will think about this more.

[eordentlich]

Good find.

[rishic3]

build

[lijinf2]

Thanks for addressing the comments. I am good with merging the PR. Some corner cases can be studied in the future.

[rishic3]

build

[lijinf2]

build