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.
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.
Gets the value of inputCol or its default value.
Gets the value of inputCols or its default value.
getK
()Get the value of k.
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)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
Returns the dictionary of parameters intended for the underlying cuML class.
Number of cuML workers, where each cuML worker corresponds to one Spark task running on one GPU.
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.
- 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.
- 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.
- 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.
- 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:
- dataset
pyspark.sql.DataFrame
input dataset
- paramsdict, optional
an optional param map that overrides embedded params.
- dataset
- 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 typeParam
.