postgrespro
diff --git a/‎aqo.c
Lines changed: 25 additions & 1 deletion b/‎aqo.c
Lines changed: 25 additions & 1 deletion
diff --git a/‎aqo.h
Lines changed: 1 addition & 0 deletions b/‎aqo.h
Lines changed: 1 addition & 0 deletions
diff --git a/‎cardinality_estimation.c
Lines changed: 1 addition & 1 deletion b/‎cardinality_estimation.c
Lines changed: 1 addition & 1 deletion
@@ -34,6 +34,7 @@ void _PG_init(void);
 /* Strategy of determining feature space for new queries. */
 int		aqo_mode = AQO_MODE_CONTROLLED;
 bool	force_collect_stat;
+bool	aqo_predict_with_few_neighbors;
 
 /*
  * Show special info in EXPLAIN mode.
@@ -71,7 +72,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;
 
 /*
@@ -306,6 +307,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;
 
@@ -217,6 +217,7 @@ extern double auto_tuning_convergence_error;
 /* Machine learning parameters */
 
 extern int	aqo_k;
+extern bool aqo_predict_with_few_neighbors;
 extern double log_selectivity_lower_bound;
 
 /* Parameters for current query */
 
@@ -93,7 +93,7 @@ predict_for_relation(List *clauses, List *selectivities, List *relsigns,
 		 */
 
 		/* Try to search in surrounding feature spaces for the same node */
-		if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, use_wide_search))
+		if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, use_wide_search, features))
 			result = -1;
 		else
 		{
Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,7 @@ predict_for_relation(List clauses, List selectivities, List *relsigns,`
`93`	`93`	`*/`
`94`	`94`
`95`	`95`	`/* Try to search in surrounding feature spaces for the same node */`
`96`		`- if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, use_wide_search))`
	`96`	`+ if (!load_aqo_data(query_context.fspace_hash, *fss, data, NULL, use_wide_search, features))`
`97`	`97`	`result = -1;`
`98`	`98`	`else`
`99`	`99`	`{`