Add aqo_k as custom parameter to define few number of features

Alena Rybakina · Alena Rybakina · commit 50ee33c1eacc · 2022-07-09T18:30:20.000+03:00
for prediction. Its default value is 3. Queries can contain
a more number of features than 3 especially generic queries.

Also add predict_a_few_neibours parameter for switch avalable
to predict a few neibors than 3. It is done for not to change
the previous logic of the code
diff --git a/aqo.c b/aqo.c
@@ -48,6 +48,7 @@ bool	force_collect_stat;
  */
 bool	aqo_show_hash;
 bool	aqo_show_details;
+bool	aqo_predict_with_few_neighbors;
 
 /* GUC variables */
 static const struct config_enum_entry format_options[] = {
@@ -72,7 +73,7 @@ int			auto_tuning_infinite_loop = 8;
 /* Machine learning parameters */
 
 /* The number of nearest neighbors which will be chosen for ML-operations */
-int			aqo_k = 3;
+int			aqo_k;
 double		log_selectivity_lower_bound = -30;
 
 /*
@@ -237,6 +238,29 @@ _PG_init(void)
 							NULL
 	);
 
+		DefineCustomIntVariable("aqo.k_neighbors_threshold",
+							"Set the threshold of number of neighbors for predicting.",
+							NULL,
+							&aqo_k,
+							3,
+							1, INT_MAX / 1000,
+							PGC_USERSET,
+							0,
+							NULL,
+							NULL,
+							NULL);
+
+	DefineCustomBoolVariable("aqo.predict_with_few_neighbors",
+							"Make prediction with less neighbors than we should have.",
+							 NULL,
+							 &aqo_predict_with_few_neighbors,
+							 true,
+							 PGC_USERSET,
+							 0,
+							 NULL,
+							 lc_assign_hook,
+							 NULL);
+
 	prev_shmem_startup_hook						= shmem_startup_hook;
 	shmem_startup_hook							= aqo_init_shmem;
 	prev_planner_hook							= planner_hook;
diff --git a/aqo.h b/aqo.h
@@ -211,6 +211,7 @@ extern double auto_tuning_exploration;
 extern int	auto_tuning_max_iterations;
 extern int	auto_tuning_infinite_loop;
 extern double auto_tuning_convergence_error;
+extern bool aqo_predict_with_few_neighbors;
 
 /* Machine learning parameters */
 
diff --git a/machine_learning.c b/machine_learning.c
@@ -158,6 +158,9 @@ OkNNr_predict(OkNNrdata *data, double *features)
 
 	Assert(data != NULL);
 
+	if (!aqo_predict_with_few_neighbors && data->rows < aqo_k)
+		return -1.;
+
 	for (i = 0; i < data->rows; ++i)
 		distances[i] = fs_distance(data->matrix[i], features, data->cols);
 
