Try reducing the memory overhead.

Free some allocated memory right after use.
Reset AQOPredictMemCtx as soon as posible.
Remove learning attempts on SubPlan nodes.
Bugfix. Free allocated memory on save/load data.
Add memory context for storage.
Change copyright to 2016-2023.

Author: Daniil Anisimov

aqo.c

@@ -2,7 +2,7 @@
  * aqo.c
  *		Adaptive query optimization extension
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/aqo.c
@@ -92,6 +92,9 @@ MemoryContext 		AQOPredictMemCtx = NULL;
 /* Is released at the end of learning */
 MemoryContext 		AQOLearnMemCtx = NULL;
 
+/* Is released at the end of load/store routines */
+MemoryContext 		AQOStorageMemCtx = NULL;
+
 /* Additional plan info */
 int njoins;
 
@@ -348,6 +351,12 @@ _PG_init(void)
 	AQOLearnMemCtx = AllocSetContextCreate(AQOTopMemCtx,
 											 "AQOLearnMemoryContext",
 											 ALLOCSET_DEFAULT_SIZES);
+	/*
+	 * AQOStorageMemoryContext containe data for load/store routines.
+	 */
+	AQOStorageMemCtx = AllocSetContextCreate(AQOTopMemCtx,
+											 "AQOStorageMemoryContext",
+											 ALLOCSET_DEFAULT_SIZES);
 	RegisterResourceReleaseCallback(aqo_free_callback, NULL);
 	RegisterAQOPlanNodeMethods();
 

aqo.h

@@ -105,7 +105,7 @@
  * Module storage.c is responsible for storage query settings and models
  * (i. e. all information which is used in extension).
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/aqo.h
@@ -232,6 +232,7 @@ extern MemoryContext AQOTopMemCtx;
 extern MemoryContext AQOCacheMemCtx;
 extern MemoryContext AQOPredictMemCtx;
 extern MemoryContext AQOLearnMemCtx;
+extern MemoryContext AQOStorageMemCtx;
 
 extern int aqo_statement_timeout;
 

auto_tuning.c

@@ -8,7 +8,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/auto_tuning.c

cardinality_estimation.c

@@ -8,7 +8,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/cardinality_estimation.c

cardinality_hooks.c

@@ -18,7 +18,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/cardinality_hooks.c
@@ -82,6 +82,7 @@ aqo_set_baserel_rows_estimate(PlannerInfo *root, RelOptInfo *rel)
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -100,6 +101,7 @@ aqo_set_baserel_rows_estimate(PlannerInfo *root, RelOptInfo *rel)
 
 	/* Return to the caller's memory context. */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	if (predicted < 0)
 		goto default_estimator;
@@ -191,12 +193,15 @@ aqo_get_parameterized_baserel_size(PlannerInfo *root,
 			cache_selectivity(current_hash, rel->relid, rte->relid,
 							  *((double *) lfirst(l2)));
 		}
+
+		pfree(args_hash);
+		pfree(eclass_hash);
 	}
 
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(oldctx);
-
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -211,6 +216,7 @@ aqo_get_parameterized_baserel_size(PlannerInfo *root,
 
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(oldctx);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	predicted_ppi_rows = predicted;
 	fss_ppi_hash = fss;
@@ -265,6 +271,7 @@ aqo_set_joinrel_size_estimates(PlannerInfo *root, RelOptInfo *rel,
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -284,6 +291,7 @@ aqo_set_joinrel_size_estimates(PlannerInfo *root, RelOptInfo *rel,
 
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	rel->fss_hash = fss;
 
@@ -343,6 +351,7 @@ aqo_get_parameterized_joinrel_size(PlannerInfo *root,
 	if (!query_context.use_aqo)
 	{
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		goto default_estimator;
 	}
 
@@ -359,6 +368,7 @@ aqo_get_parameterized_joinrel_size(PlannerInfo *root,
 									 &fss);
 	/* Return to the caller's memory context */
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 	predicted_ppi_rows = predicted;
 	fss_ppi_hash = fss;
@@ -450,6 +460,7 @@ aqo_estimate_num_groups(PlannerInfo *root, List *groupExprs,
 		grouped_rel->rows = predicted;
 		grouped_rel->fss_hash = fss;
 		MemoryContextSwitchTo(old_ctx_m);
+		MemoryContextReset(AQOPredictMemCtx);
 		return predicted;
 	}
 	else
@@ -460,6 +471,7 @@ aqo_estimate_num_groups(PlannerInfo *root, List *groupExprs,
 		grouped_rel->predicted_cardinality = -1;
 
 	MemoryContextSwitchTo(old_ctx_m);
+	MemoryContextReset(AQOPredictMemCtx);
 
 default_estimator:
 	if (aqo_estimate_num_groups_next)

expected/unsupported.out

@@ -311,6 +311,59 @@ EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
  JOINS: 0
 (23 rows)
 
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+                        QUERY PLAN                         
+-----------------------------------------------------------
+ Seq Scan on t (actual rows=1000 loops=1)
+   AQO not used
+   Filter: ((x = (SubPlan 1)) AND (SubPlan 2))
+   SubPlan 1
+     ->  Limit (actual rows=1 loops=1000)
+           AQO not used
+           ->  Seq Scan on t t0 (actual rows=1 loops=1000)
+                 AQO not used
+                 Filter: (x = t.x)
+                 Rows Removed by Filter: 475
+   SubPlan 2
+     ->  Seq Scan on t t0_1 (actual rows=1 loops=1000)
+           AQO not used
+           Filter: (x = t.x)
+           Rows Removed by Filter: 475
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(18 rows)
+
+-- No prediction for top SeqScan, because it fss is changed
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
+	SELECT * FROM t WHERE
+			x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND
+			x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
+                        QUERY PLAN                         
+-----------------------------------------------------------
+ Seq Scan on t (actual rows=1000 loops=1)
+   AQO not used
+   Filter: ((SubPlan 2) AND (x = (SubPlan 1)))
+   SubPlan 2
+     ->  Seq Scan on t t0_1 (actual rows=1 loops=1000)
+           AQO: rows=1, error=0%
+           Filter: (x = t.x)
+           Rows Removed by Filter: 475
+   SubPlan 1
+     ->  Limit (actual rows=1 loops=1000)
+           AQO not used
+           ->  Seq Scan on t t0 (actual rows=1 loops=1000)
+                 AQO: rows=1, error=0%
+                 Filter: (x = t.x)
+                 Rows Removed by Filter: 475
+ Using aqo: true
+ AQO mode: LEARN
+ JOINS: 0
+(18 rows)
+
 -- It's OK to use the knowledge for a query with different constants.
 EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)
 	SELECT count(*) FROM t WHERE
@@ -580,6 +633,10 @@ ORDER BY (md5(query_text),error) DESC;
 -------+------------------------------------------------------------------------------------------------
  0.768 | SELECT count(*) FROM (SELECT count(*) FROM t1 GROUP BY (x,y)) AS q1;
  0.070 | SELECT count(*) FROM (SELECT * FROM t GROUP BY (x) HAVING x > 3) AS q1;
+ 1.554 | EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF)                                         +
+       |         SELECT * FROM t WHERE                                                                 +
+       |                         x = (SELECT x FROM t t0 WHERE t0.x = t.x LIMIT 1) AND                 +
+       |                         x IN (SELECT x FROM t t0 WHERE t0.x = t.x);
  0.000 | SELECT count(*) FROM t WHERE x < 3 AND mod(x,3) = 1;
  0.000 | SELECT * FROM                                                                                 +
        |         (SELECT * FROM t WHERE x < 0) AS t0                                                   +
@@ -612,13 +669,13 @@ ORDER BY (md5(query_text),error) DESC;
        |                 JOIN                                                                          +
        |         (SELECT * FROM t WHERE x % 3 < (SELECT avg(x) FROM t t0 WHERE t0.x <> t.x)) AS q2     +
        |                 ON q1.x = q2.x+1;
-(13 rows)
+(14 rows)
 
 DROP TABLE t,t1 CASCADE; -- delete all tables used in the test
 SELECT count(*) FROM aqo_data; -- Just to detect some changes in the logic. May some false positives really bother us here?
  count 
 -------
-    44
+    48
 (1 row)
 
 SELECT true AS success FROM aqo_cleanup();

hash.c

@@ -12,7 +12,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/hash.c
@@ -157,6 +157,8 @@ get_grouped_exprs_hash(int child_fss, List *group_exprs)
 	final_hashes[0] = child_fss;
 	final_hashes[1] = get_int_array_hash(hashes, i);
 
+	pfree(hashes);
+
 	return get_int_array_hash(final_hashes, 2);
 }
 
@@ -224,6 +226,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		clause_has_consts[i] = (args != NULL && has_consts(*args));
 		i++;
 	}
+	pfree(args_hash);
 
 	idx = argsort(clause_hashes, n, sizeof(*clause_hashes), int_cmp);
 	inverse_idx = inverse_permutation(idx, n);
@@ -234,6 +237,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		sorted_clauses[inverse_idx[i]] = clause_hashes[i];
 		i++;
 	}
+	pfree(clause_hashes);
 
 	i = 0;
 	foreach(lc, selectivities)
@@ -249,6 +253,7 @@ get_fss_for_object(List *relsigns, List *clauselist,
 		}
 		i++;
 	}
+	pfree(inverse_idx);
 
 	for (i = 0; i < n;)
 	{
@@ -272,6 +277,8 @@ get_fss_for_object(List *relsigns, List *clauselist,
 				  sizeof(**features), double_cmp);
 		i = j;
 	}
+	pfree(idx);
+	pfree(clause_has_consts);
 
 	/*
 	 * Generate feature subspace hash.
@@ -281,6 +288,8 @@ get_fss_for_object(List *relsigns, List *clauselist,
 	eclasses_hash = get_int_array_hash(eclass_hash, nargs);
 	relations_hash = get_relations_hash(relsigns);
 	fss_hash = get_fss_hash(clauses_hash, eclasses_hash, relations_hash);
+	pfree(sorted_clauses);
+	pfree(eclass_hash);
 
 	if (nfeatures != NULL)
 	{
@@ -340,11 +349,17 @@ static int
 get_node_hash(Node *node)
 {
 	char	   *str;
+	char	   *no_consts;
+	char	   *no_locations;
 	int			hash;
 
-	str = remove_locations(remove_consts(nodeToString(node)));
-	hash = get_str_hash(str);
+	str = nodeToString(node);
+	no_consts = remove_consts(str);
 	pfree(str);
+	no_locations = remove_locations(no_consts);
+	pfree(no_consts);
+	hash = get_str_hash(no_locations);
+	pfree(no_locations);
 	return hash;
 }
 
@@ -467,6 +482,7 @@ get_relations_hash(List *relsigns)
 
 	result = DatumGetInt32(hash_any((const unsigned char *) hashes,
 									nhashes * sizeof(uint32)));
+	pfree(hashes);
 
 	return result;
 }
@@ -479,9 +495,11 @@ static char *
 remove_consts(const char *str)
 {
 	char *res;
+	char *tmp;
 
-	res = replace_patterns(str, "{CONST", is_brace);
-	res = replace_patterns(res, ":stmt_len", is_brace);
+	tmp = replace_patterns(str, "{CONST", is_brace);
+	res = replace_patterns(tmp, ":stmt_len", is_brace);
+	pfree(tmp);
 	return res;
 }
 
@@ -683,6 +701,8 @@ get_eclasses(List *clauselist, int *nargs, int **args_hash, int **eclass_hash)
 
 	for (i = 0; i < *nargs; ++i)
 		(*eclass_hash)[i] = e_hashes[disjoint_set_get_parent(p, i)];
+
+	pfree(e_hashes);
 }
 
 /*

machine_learning.c

@@ -12,7 +12,7 @@
  *
  *******************************************************************************
  *
- * Copyright (c) 2016-2022, Postgres Professional
+ * Copyright (c) 2016-2023, Postgres Professional
  *
  * IDENTIFICATION
  *	  aqo/machine_learning.c