NearestNeighborsModel#

class spark_rapids_ml.knn.NearestNeighborsModel(item_df_withid: DataFrame, processed_item_df: DataFrame, label_isdata: int, label_isquery: int)#

Methods

`clear`(param)	Reset a Spark ML Param to its default value, setting matching cuML parameter, if exists.
`copy`([extra])
`cpu`()	Return the equivalent PySpark CPU model
`exactNearestNeighborsJoin`(query_df[, distCol])	This function returns the k exact nearest neighbors (knn) in item_df of each query vector in query_df.
`explainParam`(param)	Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.
`explainParams`()	Returns the documentation of all params with their optionally default values and user-supplied values.
`extractParamMap`([extra])	Extracts the embedded default param values and user-supplied values, and then merges them with extra values from input into a flat param map, where the latter value is used if there exist conflicts, i.e., with ordering: default param values < user-supplied values < extra.
`getIdCol`()	Gets the value of idCol.
`getInputCol`()	Gets the value of inputCol or its default value.
`getInputCols`()	Gets the value of inputCols or its default value.
`getK`()	Get the value of k.
`getLabelCol`()	Gets the value of labelCol or its default value.
`getOrDefault`(param)	Gets the value of a param in the user-supplied param map or its default value.
`getParam`(paramName)	Gets a param by its name.
`hasDefault`(param)	Checks whether a param has a default value.
`hasParam`(paramName)	Tests whether this instance contains a param with a given (string) name.
`isDefined`(param)	Checks whether a param is explicitly set by user or has a default value.
`isSet`(param)	Checks whether a param is explicitly set by user.
`kneighbors`(query_df[, sort_knn_df_by_query_id])	Return the exact nearest neighbors for each query in query_df.
`set`(param, value)	Sets a parameter in the embedded param map.
`setIdCol`(value)	Sets the value of idCol.
`setInputCol`(value)	Sets the value of `inputCol` or `inputCols`.
`setInputCols`(value)	Sets the value of `inputCols`.
`setK`(value)	Sets the value of k.
`transform`(dataset[, params])	Transforms the input dataset with optional parameters.

Attributes

`cuml_params`	Returns the dictionary of parameters intended for the underlying cuML class.
`idCol`
`inputCol`
`inputCols`
`k`
`labelCol`
`num_workers`	Number of cuML workers, where each cuML worker corresponds to one Spark task running on one GPU.
`params`	Returns all params ordered by name.

Methods Documentation

clear(param: Param) → None#: Reset a Spark ML Param to its default value, setting matching cuML parameter, if exists.

copy(extra: Optional[ParamMap] = None) → P#

cpu() → Model#: Return the equivalent PySpark CPU model

exactNearestNeighborsJoin(query_df: DataFrame, distCol: str = 'distCol') → DataFrame#

This function returns the k exact nearest neighbors (knn) in item_df of each query vector in query_df. item_df is the dataframe passed to the fit function of the NearestNeighbors estimator. Note that the knn relationship is asymmetric with respect to the input datasets (e.g., if x is a knn of y , y is not necessarily a knn of x).

Parameters:

query_df: pyspark.sql.DataFrame: the query_df dataframe. Each row represents a query vector.
distCol: str: the name of the output distance column

Returns:

knnjoin_df: pyspark.sql.DataFrame: the result dataframe that has three columns (item_df, query_df, distCol). item_df column is of struct type that includes as fields all the columns of input item dataframe. Similarly, query_df column is of struct type that includes as fields all the columns of input query dataframe. distCol is the distance column. A row in knnjoin_df is in the format (v1, v2, dist(v1, v2)), where item_vector v1 is one of the k nearest neighbors of query_vector v2 and their distance is dist(v1, v2).

explainParam(param: Union[str, Param]) → str#: Explains a single param and returns its name, doc, and optional default value and user-supplied value in a string.

explainParams() → str#: Returns the documentation of all params with their optionally default values and user-supplied values.

extractParamMap(extra: Optional[ParamMap] = None) → ParamMap#

Extracts the embedded default param values and user-supplied values, and then merges them with extra values from input into a flat param map, where the latter value is used if there exist conflicts, i.e., with ordering: default param values < user-supplied values < extra.

Parameters:

extradict, optional: extra param values

Returns:

dict: merged param map

getIdCol() → str#: Gets the value of idCol.

getInputCol() → str#: Gets the value of inputCol or its default value.

getInputCols() → List[str]#: Gets the value of inputCols or its default value.

getK() → int#: Get the value of k.

getLabelCol() → str#: Gets the value of labelCol or its default value.

getOrDefault(param: Union[str, Param[T]]) → Union[Any, T]#: Gets the value of a param in the user-supplied param map or its default value. Raises an error if neither is set.

getParam(paramName: str) → Param#: Gets a param by its name.

hasDefault(param: Union[str, Param[Any]]) → bool#: Checks whether a param has a default value.

hasParam(paramName: str) → bool#: Tests whether this instance contains a param with a given (string) name.

isDefined(param: Union[str, Param[Any]]) → bool#: Checks whether a param is explicitly set by user or has a default value.

isSet(param: Union[str, Param[Any]]) → bool#: Checks whether a param is explicitly set by user.

kneighbors(query_df: DataFrame, sort_knn_df_by_query_id: bool = True) → Tuple[DataFrame, DataFrame, DataFrame]#

Return the exact nearest neighbors for each query in query_df. The data vectors (or equivalently item vectors) should be provided through the fit function (see Examples in the spark_rapids_ml.knn.NearestNeighbors). The distance measure here is euclidean distance and the number of target exact nearest neighbors can be set through setK(). The function currently only supports float32 type and will convert other data types into float32.

Parameters:

query_df: pyspark.sql.DataFrame: query vectors where each row corresponds to one query. The query_df can be in the format of a single array column, a single vector column, or multiple float columns.
sort_knn_df_by_query_id: bool (default=True): whether to sort the returned dataframe knn_df by query_id

Returns:

query_df: pyspark.sql.DataFrame: the query_df itself if it has an id column set through setIdCol(). If not, a monotonically increasing id column will be added.
item_df: pyspark.sql.DataFrame: the item_df (or equivalently data_df) itself if it has an id column set through setIdCol(). If not, a monotonically increasing id column will be added.
knn_df: pyspark.sql.DataFrame: the result k nearest neighbors (knn) dataframe that has three columns (id, indices, distances). Each row of knn_df corresponds to the knn result of a query vector, identified by the id column. The indices/distances column stores the ids/distances of knn item_df vectors.

set(param: Param, value: Any) → None#: Sets a parameter in the embedded param map.

setIdCol(value: str) → P#: Sets the value of idCol. If not set, an id column will be added with column name unique_id. The id column is used to specify nearest neighbor vectors by associated id value.

setInputCol(value: Union[str, List[str]]) → P#: Sets the value of inputCol or inputCols.

setInputCols(value: List[str]) → P#: Sets the value of inputCols. Used when input vectors are stored as multiple feature columns.

setK(value: int) → P#: Sets the value of k.

transform(dataset: DataFrame, params: Optional[ParamMap] = None) → DataFrame#

Transforms the input dataset with optional parameters.

New in version 1.3.0.

Parameters:

datasetpyspark.sql.DataFrame: input dataset
paramsdict, optional: an optional param map that overrides embedded params.

Returns:

pyspark.sql.DataFrame: transformed dataset

Attributes Documentation

cuml_params#: Returns the dictionary of parameters intended for the underlying cuML class.

idCol = Param(parent='undefined', name='idCol', doc='id column name.')#

inputCol: Param[str] = Param(parent='undefined', name='inputCol', doc='input column name.')#

inputCols: Param[List[str]] = Param(parent='undefined', name='inputCols', doc='input column names.')#

k = Param(parent='undefined', name='k', doc='The number nearest neighbors to retrieve. Must be >= 1.')#

labelCol: Param[str] = Param(parent='undefined', name='labelCol', doc='label column name.')#

num_workers#: Number of cuML workers, where each cuML worker corresponds to one Spark task running on one GPU.

params#: Returns all params ordered by name. The default implementation uses dir() to get all attributes of type Param.