Avoid local dir creation, ensure dense array ordering during UMAP save() #823
Conversation
Signed-off-by: Rishi Chandra <[email protected]>
|
build |
rishic3
changed the title
python/src/spark_rapids_ml/umap.py
Outdated
| pd.DataFrame( | ||
| { | ||
| "row_id": range(array.shape[0]), | ||
| "data": array.tolist(), |
There was a problem hiding this comment.
Is this better than list(array) ?
Pretty sure we create pandas dfs array columns from np arrays and vice versa elsewhere in our code and would be good to be consistent and/or use best way through out.
There was a problem hiding this comment.
While list(array) is more efficient since tolist() does a deep conversion of each row to python lists, Spark will throw an error with list(array) if spark.sql.execution.arrow.pyspark.enabled=false, since pyarrow would no longer handle the numpy -> arrow array conversion.
There was a problem hiding this comment.
We pretty much require that to be enabled to get good data transfer from jvm to python workers.
| } | ||
| ), | ||
| schema=schema, | ||
| ) |
There was a problem hiding this comment.
Below and elsewhere in this class is it correct to use overwrite when writing? This might be counter to the overwrite MLWriter api. If that is not invoked, a user would not expect overwrite to be allowed.
python/src/spark_rapids_ml/umap.py
Outdated
| data_df = spark.read.parquet(df_path) | ||
| return np.array(data_df.collect(), dtype=np.float32) | ||
| data_df = spark.read.parquet(df_path).orderBy("row_id") | ||
| return np.array([row.data for row in data_df.collect()], dtype=np.float32) |
There was a problem hiding this comment.
toPandas might be better here followed by np.array(list(data_pandas_df.data))
| @@ -1495,8 +1504,6 @@ def write_dense_array(array: np.ndarray, df_path: str) -> None: | |||
| assert model_attributes is not None | |||
|
|
|||
| data_path = os.path.join(path, "data") | |||
| if not os.path.exists(data_path): | |||
There was a problem hiding this comment.
Would be good to have a test that checks for expected files and directories?
| @@ -1547,8 +1554,8 @@ def read_sparse_array( | |||
| return scipy.sparse.csr_matrix((data, indices, indptr), shape=csr_shape) | |||
|
|
|||
| def read_dense_array(df_path: str) -> np.ndarray: | |||
| data_df = spark.read.parquet(df_path) | |||
| return np.array(data_df.collect(), dtype=np.float32) | |||
| data_df = spark.read.parquet(df_path).orderBy("row_id") | |||
There was a problem hiding this comment.
I wonder if there is test for the order, one that would fail if orderby was omitted.
There was a problem hiding this comment.
A multi-gpu env (e.g., DGX) where Spark's default parallelism is >1 would have caught it and I should have tested there with the last PR.
Forcing >1 parallelism would require changing CleanSparkSession to allow a new conf to override the default conf - not sure if that's worth it
|
build |
No description provided.