Add FORCED_CONTROLLED and FORCED_FROZEN modes.

Daniil Anisimov · Daniil Anisimov · commit 656ec07fb5f1 · 2023-06-19T13:44:28.000+07:00
diff --git a/aqo.c b/aqo.c
@@ -53,6 +53,8 @@ bool	change_flex_timeout;
 static const struct config_enum_entry format_options[] = {
 	{"intelligent", AQO_MODE_INTELLIGENT, false},
 	{"forced", AQO_MODE_FORCED, false},
+	{"forced_controlled", AQO_MODE_FORCED_CONTROLLED, false},
+	{"forced_frozen", AQO_MODE_FORCED_FROZEN, false},
 	{"controlled", AQO_MODE_CONTROLLED, false},
 	{"learn", AQO_MODE_LEARN, false},
 	{"frozen", AQO_MODE_FROZEN, false},
diff --git a/aqo.h b/aqo.h
@@ -155,6 +155,10 @@ typedef enum
 	AQO_MODE_INTELLIGENT,
 	/* Treats new query types as linked to the common feature space */
 	AQO_MODE_FORCED,
+	/* Like FORCED, but learn only on old feature sub spaces */
+	AQO_MODE_FORCED_CONTROLLED,
+	/* Like FORCED, but use old knowledge without learning */
+	AQO_MODE_FORCED_FROZEN,
 	/* New query types are not linked with any feature space */
 	AQO_MODE_CONTROLLED,
 	/* Creates new feature space for each query type without auto-tuning */
@@ -183,6 +187,8 @@ typedef struct QueryContextData
 	bool		auto_tuning;
 	bool		collect_stat;
 	bool		adding_query;
+	/* Add new fss in learn phase, otherwise learn only on existing fss */
+	bool		adding_fss;
 	bool		explain_only;
 
 	/*
diff --git a/expected/aqo_forced.out b/expected/aqo_forced.out
@@ -6,6 +6,7 @@ SELECT true AS success FROM aqo_reset();
  t
 (1 row)
 
+SET aqo.show_details = on;
 CREATE TABLE aqo_test0(a int, b int, c int, d int);
 WITH RECURSIVE t(a, b, c, d)
 AS (
@@ -70,22 +71,117 @@ WHERE a < 3 AND b < 3 AND c < 3 AND d < 3;
                   QUERY PLAN                   
 -----------------------------------------------
  Index Scan using aqo_test0_idx_a on aqo_test0
+   AQO: rows=3
    Index Cond: (a < 3)
    Filter: ((b < 3) AND (c < 3) AND (d < 3))
-(3 rows)
+ Using aqo: true
+ AQO mode: FORCED
+ JOINS: 0
+(7 rows)
 
 EXPLAIN (COSTS FALSE)
 SELECT * FROM aqo_test0
 WHERE a < 5 AND b < 5 AND c < 5 AND d < 5;
                   QUERY PLAN                   
 -----------------------------------------------
  Index Scan using aqo_test0_idx_a on aqo_test0
+   AQO: rows=5
    Index Cond: (a < 5)
    Filter: ((b < 5) AND (c < 5) AND (d < 5))
-(3 rows)
+ Using aqo: true
+ AQO mode: FORCED
+ JOINS: 0
+(7 rows)
+
+SET aqo.mode = 'forced_controlled';
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+SELECT count(*) FROM tmp1;
+ count 
+-------
+    10
+(1 row)
+
+DROP TABLE tmp1;
+-- Not predict
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Index Scan using aqo_test0_idx_a on aqo_test0
+   AQO not used
+   Index Cond: (a < 10)
+   Filter: ((b < 10) AND (c < 10))
+ Using aqo: true
+ AQO mode: FORCED_CONTROLLED
+ JOINS: 0
+(7 rows)
+
+SET aqo.mode = 'forced';
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 0 AND b < 0 AND c < 0;
+SELECT count(*) FROM tmp1;
+ count 
+-------
+     0
+(1 row)
+
+DROP TABLE tmp1;
+SET aqo.mode = 'forced_controlled';
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+SELECT count(*) FROM tmp1;
+ count 
+-------
+    10
+(1 row)
+
+DROP TABLE tmp1;
+-- Predict
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Index Scan using aqo_test0_idx_a on aqo_test0
+   AQO: rows=10
+   Index Cond: (a < 10)
+   Filter: ((b < 10) AND (c < 10))
+ Using aqo: true
+ AQO mode: FORCED_CONTROLLED
+ JOINS: 0
+(7 rows)
+
+SET aqo.mode = 'forced_frozen';
+-- Not learn
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 20 AND b < 20 AND c < 20;
+SELECT count(*) FROM tmp1;
+ count 
+-------
+    20
+(1 row)
+
+DROP TABLE tmp1;
+-- Predict approximately 9 (interpolation between 1 and 10)
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 20 AND b < 20 AND c < 20;
+                  QUERY PLAN                   
+-----------------------------------------------
+ Index Scan using aqo_test0_idx_a on aqo_test0
+   AQO: rows=9
+   Index Cond: (a < 20)
+   Filter: ((b < 20) AND (c < 20))
+ Using aqo: true
+ AQO mode: FORCED_FROZEN
+ JOINS: 0
+(7 rows)
 
 DROP INDEX aqo_test0_idx_a;
 DROP TABLE aqo_test0;
 DROP INDEX aqo_test1_idx_a;
 DROP TABLE aqo_test1;
+RESET aqo.show_details;
 DROP EXTENSION aqo;
diff --git a/postprocessing.c b/postprocessing.c
@@ -465,6 +465,14 @@ learnOnPlanState(PlanState *p, void *context)
 	else if (!ctx->learn)
 		return true;
 
+	if (!query_context.adding_fss &&
+		!aqo_data_exist(query_context.fspace_hash, aqo_node->fss))
+		/*
+		 * Skip the node if it does not exist in aqo_data
+		 * and do not need to be added.
+		 */
+		goto end;
+
 	/*
 	 * Need learn.
 	 */
@@ -999,6 +1007,12 @@ print_into_explain(PlannedStmt *plannedstmt, IntoClause *into,
 	case AQO_MODE_FORCED:
 		ExplainPropertyText("AQO mode", "FORCED", es);
 		break;
+	case AQO_MODE_FORCED_CONTROLLED:
+		ExplainPropertyText("AQO mode", "FORCED_CONTROLLED", es);
+		break;
+	case AQO_MODE_FORCED_FROZEN:
+		ExplainPropertyText("AQO mode", "FORCED_FROZEN", es);
+		break;
 	case AQO_MODE_CONTROLLED:
 		ExplainPropertyText("AQO mode", "CONTROLLED", es);
 		break;
diff --git a/preprocessing.c b/preprocessing.c
@@ -176,33 +176,55 @@ aqo_planner(Query *parse, const char *query_string, int cursorOptions,
 		{
 			case AQO_MODE_INTELLIGENT:
 				query_context.adding_query = true;
+				query_context.adding_fss = true;
 				query_context.learn_aqo = true;
 				query_context.use_aqo = false;
 				query_context.auto_tuning = true;
 				query_context.collect_stat = true;
 				break;
 			case AQO_MODE_FORCED:
 				query_context.adding_query = false;
+				query_context.adding_fss = true;
 				query_context.learn_aqo = true;
 				query_context.use_aqo = true;
 				query_context.auto_tuning = false;
 				query_context.fspace_hash = 0; /* Use common feature space */
 				query_context.collect_stat = false;
 				break;
+			case AQO_MODE_FORCED_CONTROLLED:
+				query_context.adding_query = false;
+				query_context.adding_fss = false;
+				query_context.learn_aqo = true;
+				query_context.use_aqo = true;
+				query_context.auto_tuning = false;
+				query_context.fspace_hash = 0; /* Use common feature space */
+				query_context.collect_stat = false;
+				break;
+			case AQO_MODE_FORCED_FROZEN:
+				query_context.adding_query = false;
+				query_context.adding_fss = false;
+				query_context.learn_aqo = false;
+				query_context.use_aqo = true;
+				query_context.auto_tuning = false;
+				query_context.fspace_hash = 0; /* Use common feature space */
+				query_context.collect_stat = false;
+				break;
 			case AQO_MODE_CONTROLLED:
 			case AQO_MODE_FROZEN:
 				/*
 				 * if query is not in the AQO knowledge base than disable AQO
 				 * for this query.
 				 */
 				query_context.adding_query = false;
+				query_context.adding_fss = true;
 				query_context.learn_aqo = false;
 				query_context.use_aqo = false;
 				query_context.auto_tuning = false;
 				query_context.collect_stat = false;
 				break;
 			case AQO_MODE_LEARN:
 				query_context.adding_query = true;
+				query_context.adding_fss = true;
 				query_context.learn_aqo = true;
 				query_context.use_aqo = true;
 				query_context.auto_tuning = false;
@@ -222,6 +244,7 @@ aqo_planner(Query *parse, const char *query_string, int cursorOptions,
 	else /* Query class exists in a ML knowledge base. */
 	{
 		query_context.adding_query = false;
+		query_context.adding_fss = true;
 
 		/* Other query_context fields filled in the find_query() routine. */
 
@@ -243,6 +266,7 @@ aqo_planner(Query *parse, const char *query_string, int cursorOptions,
 			 * In this mode we will suppress all writings to the knowledge base.
 			 * AQO will be used for all known queries, if it is not suppressed.
 			 */
+			query_context.adding_fss = false;
 			query_context.learn_aqo = false;
 			query_context.auto_tuning = false;
 			query_context.collect_stat = false;
@@ -256,6 +280,22 @@ aqo_planner(Query *parse, const char *query_string, int cursorOptions,
 			query_context.collect_stat = true;
 			break;
 
+		case AQO_MODE_FORCED_CONTROLLED:
+			/*
+			 * AQO will be learned for all known feature sub spaces.
+			 */
+			query_context.adding_fss = false;
+			break;
+
+		case AQO_MODE_FORCED_FROZEN:
+			/*
+			 * In this mode we will suppress all writings to the knowledge base.
+			 * AQO will be used for all known feature sub spaces.
+			 */
+			query_context.adding_fss = false;
+			query_context.learn_aqo = false;
+			break;
+
 		case AQO_MODE_INTELLIGENT:
 		case AQO_MODE_FORCED:
 		case AQO_MODE_CONTROLLED:
diff --git a/sql/aqo_forced.sql b/sql/aqo_forced.sql
@@ -1,6 +1,7 @@
 -- Preliminaries
 CREATE EXTENSION IF NOT EXISTS aqo;
 SELECT true AS success FROM aqo_reset();
+SET aqo.show_details = on;
 
 CREATE TABLE aqo_test0(a int, b int, c int, d int);
 WITH RECURSIVE t(a, b, c, d)
@@ -52,9 +53,53 @@ EXPLAIN (COSTS FALSE)
 SELECT * FROM aqo_test0
 WHERE a < 5 AND b < 5 AND c < 5 AND d < 5;
 
+SET aqo.mode = 'forced_controlled';
+
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+SELECT count(*) FROM tmp1;
+DROP TABLE tmp1;
+
+-- Not predict
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+
+SET aqo.mode = 'forced';
+
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 0 AND b < 0 AND c < 0;
+SELECT count(*) FROM tmp1;
+DROP TABLE tmp1;
+
+SET aqo.mode = 'forced_controlled';
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+SELECT count(*) FROM tmp1;
+DROP TABLE tmp1;
+
+-- Predict
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 10 AND b < 10 AND c < 10;
+
+SET aqo.mode = 'forced_frozen';
+
+-- Not learn
+CREATE TABLE tmp1 AS SELECT * FROM aqo_test0
+WHERE a < 20 AND b < 20 AND c < 20;
+SELECT count(*) FROM tmp1;
+DROP TABLE tmp1;
+
+-- Predict approximately 9 (interpolation between 1 and 10)
+EXPLAIN (COSTS FALSE)
+SELECT * FROM aqo_test0
+WHERE a < 20 AND b < 20 AND c < 20;
+
 DROP INDEX aqo_test0_idx_a;
 DROP TABLE aqo_test0;
 DROP INDEX aqo_test1_idx_a;
 DROP TABLE aqo_test1;
 
+RESET aqo.show_details;
 DROP EXTENSION aqo;
diff --git a/storage.c b/storage.c
@@ -1630,6 +1630,25 @@ _fill_knn_data(const DataEntry *entry, List **reloids)
 	return data;
 }
 
+/*
+ * Return TRUE when aqo_data contains a row
+ * for the given feature space and subspace.
+ */
+bool
+aqo_data_exist(uint64 fs, int fss)
+{
+	bool		found;
+	data_key	key = {.fs = fs, .fss = fss};
+
+	dsa_init();
+
+	LWLockAcquire(&aqo_state->data_lock, LW_SHARED);
+	hash_search(data_htab, &key, HASH_FIND, &found);
+	LWLockRelease(&aqo_state->data_lock);
+
+	return found;
+}
+
 /*
  * By given feature space and subspace, build kNN data structure.
  *
diff --git a/storage.h b/storage.h
@@ -142,6 +142,7 @@ extern bool aqo_qtext_store(uint64 queryid, const char *query_string);
 extern void aqo_qtexts_flush(void);
 extern void aqo_qtexts_load(void);
 
+extern bool aqo_data_exist(uint64 fs, int fss);
 extern bool aqo_data_store(uint64 fs, int fss, AqoDataArgs *data,
 						   List *reloids);
 extern bool load_aqo_data(uint64 fs, int fss, OkNNrdata *data, List **reloids,
