Extend tableam->relation_size() functionality to use different calculation methods.

This allows to output size metrics for relations and indexes provided by OrioleDB extension using existing PG functions.

If extension function doesn't support the requested method it should output negative value and thus asking to fall
back to using PG internal calculation e.g.

- Orioledb relation_size outputs -1 for bridged indexes to fallback to counting them as PG indexes
- PG table_block_relation_size outputs -1 for any method except the only DEFAULT_SIZE that it supports.

This API relies on table AM extensibility and doesn't use index AM extensibility yet, which could look more logical,
but also more complicated.

Author: pashkinelfe

src/backend/access/table/tableam.c

@@ -635,12 +635,18 @@ table_block_parallelscan_nextpage(Relation rel,
  * is stored, and if it uses them in the expected manner (e.g. the actual data
  * is in the main fork rather than some other), it can use this implementation
  * of the relation_size callback rather than implementing its own.
+ *
+ * Different counting methods is not supported for this function yet. It's expected
+ * DEFAULT_SIZE in all cases.
  */
-uint64
-table_block_relation_size(Relation rel, ForkNumber forkNumber)
+int64
+table_block_relation_size(Relation rel, ForkNumber forkNumber, uint8 method)
 {
 	uint64		nblocks = 0;
 
+	if (method != DEFAULT_SIZE)
+		return -1;
+
 	/* InvalidForkNumber indicates returning the size for all forks */
 	if (forkNumber == InvalidForkNumber)
 	{
@@ -650,7 +656,7 @@ table_block_relation_size(Relation rel, ForkNumber forkNumber)
 	else
 		nblocks = smgrnblocks(RelationGetSmgr(rel), forkNumber);
 
-	return nblocks * BLCKSZ;
+	return (int64) nblocks * BLCKSZ;
 }
 
 /*

src/backend/utils/adt/dbsize.c

@@ -31,6 +31,7 @@
 #include "utils/relfilenumbermap.h"
 #include "utils/relmapper.h"
 #include "utils/syscache.h"
+#include "access/tableam.h"
 
 /* Divide by two and round away from zero */
 #define half_rounded(x)   (((x) + ((x) < 0 ? -1 : 1)) / 2)
@@ -343,6 +344,61 @@ calculate_relation_size(RelFileLocator *rfn, BackendId backend, ForkNumber forkn
 	return totalsize;
 }
 
+/*
+ * Try to get size using proper relation_size method in table am. Fallback to previous method
+ * if rd_tableam->relation_size doesn't support requested counting method or otherwise refuses
+ * to count (with negative output)
+ */
+static
+int64 try_tableam_relation_size(Relation rel, ForkNumber forkNum, bool allforks, uint8 method)
+{
+	int64 size = 0;
+
+	if (rel->rd_rel->relkind == RELKIND_INDEX)
+	{
+		/* For index check rd_tableam for parent relation */
+		Relation        tbl;
+
+		tbl = relation_open(rel->rd_index->indrelid, AccessShareLock);
+		if(tbl->rd_tableam && tbl->rd_tableam->relation_size)
+		{
+			/* We call relation_size method for parent relation but provide index relation as an argument.
+			 * Method for index is always RELATION_SIZE
+			*/
+			if (allforks)
+			{
+				for (ForkNumber i = 0; i <= MAX_FORKNUM; i++)
+				{
+					size += tbl->rd_tableam->relation_size(rel, i, RELATION_SIZE);
+				}
+			}
+			else
+				size = tbl->rd_tableam->relation_size(rel, forkNum, RELATION_SIZE);
+
+			if (size >= 0)
+			{
+				relation_close(tbl, AccessShareLock);
+				return size;
+			}
+		}
+		relation_close(tbl, AccessShareLock);
+	}
+	else if (rel->rd_tableam && rel->rd_tableam->relation_size)
+	{
+		if (allforks)
+		{
+			for (ForkNumber i = 0; i <= MAX_FORKNUM; i++)
+				size += rel->rd_tableam->relation_size(rel, i, method);
+		}
+		else
+			size = rel->rd_tableam->relation_size(rel, forkNum, method);
+
+		if (size >= 0)
+			return size;
+	}
+	return -1;
+}
+
 Datum
 pg_relation_size(PG_FUNCTION_ARGS)
 {
@@ -363,6 +419,18 @@ pg_relation_size(PG_FUNCTION_ARGS)
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
+	/*
+	 * Try to get size using proper relation_size method in table am. Fallback to previous method
+	 * if rd_tableam->relation_size doesn't support requested counting method or otherwise refuses
+	 * to count (with negative output)
+	 */
+	size = try_tableam_relation_size(rel, forkname_to_number(text_to_cstring(forkName)), false, RELATION_SIZE);
+	if (size >= 0)
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_INT64(size);
+	}
+
 	size = calculate_relation_size(&(rel->rd_locator), rel->rd_backend,
 								   forkname_to_number(text_to_cstring(forkName)));
 
@@ -495,6 +563,18 @@ pg_table_size(PG_FUNCTION_ARGS)
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
+	/*
+	 * Try to get size using proper relation_size method in table am. Fallback to previous method
+	 * if rd_tableam->relation_size doesn't support requested counting method or otherwise refuses
+	 * to count (with negative output)
+	 */
+	size = try_tableam_relation_size(rel, InvalidForkNumber, true, TABLE_SIZE);
+	if (size >= 0)
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_INT64(size);
+	}
+
 	size = calculate_table_size(rel);
 
 	relation_close(rel, AccessShareLock);
@@ -507,13 +587,25 @@ pg_indexes_size(PG_FUNCTION_ARGS)
 {
 	Oid			relOid = PG_GETARG_OID(0);
 	Relation	rel;
-	int64		size;
+	int64		size = 0;
 
 	rel = try_relation_open(relOid, AccessShareLock);
 
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
+	if (rel->rd_tableam && rel->rd_tableam->relation_size)
+	{
+		for (ForkNumber forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
+			size += rel->rd_tableam->relation_size(rel, forkNum , INDEXES_SIZE);
+
+		if(size >= 0)
+		{
+			relation_close(rel, AccessShareLock);
+			PG_RETURN_INT64(size);
+		}
+	}
+
 	size = calculate_indexes_size(rel);
 
 	relation_close(rel, AccessShareLock);
@@ -556,6 +648,18 @@ pg_total_relation_size(PG_FUNCTION_ARGS)
 	if (rel == NULL)
 		PG_RETURN_NULL();
 
+	/*
+	 * Try to get size using proper relation_size method in table am. Fallback to previous method
+	 * if rd_tableam->relation_size doesn't support requested counting method or otherwise refuses
+	 * to count (with negative output)
+	 */
+	size = try_tableam_relation_size(rel, InvalidForkNumber, true, TOTAL_SIZE);
+	if (size >= 0)
+	{
+		relation_close(rel, AccessShareLock);
+		PG_RETURN_INT64(size);
+	}
+
 	size = calculate_total_relation_size(rel);
 
 	relation_close(rel, AccessShareLock);

src/include/access/tableam.h

@@ -279,6 +279,15 @@ typedef struct TM_IndexDeleteOp
 #define TABLE_MODIFY_FETCH_OLD_TUPLE 0x0002
 #define TABLE_MODIFY_LOCK_UPDATED	0x0004
 
+/* "method" flag bits for relation_size */
+/* Default behavior implemented for heap AM */
+#define DEFAULT_SIZE (0)
+/* Extended behavior that AM can provide */
+#define RELATION_SIZE (1)
+#define TABLE_SIZE (2)
+#define TOTAL_SIZE (3)
+#define TOAST_TABLE_SIZE (4)
+#define INDEXES_SIZE (5)
 
 /* Typedef for callback function for table_index_build_scan */
 typedef void (*IndexBuildCallback) (Relation index,
@@ -739,7 +748,7 @@ typedef struct TableAmRoutine
 	 * probable that we'll need to revise the interface for those at some
 	 * point.
 	 */
-	uint64		(*relation_size) (Relation rel, ForkNumber forkNumber);
+	int64		(*relation_size) (Relation rel, ForkNumber forkNumber, uint8 method);
 
 
 	/*
@@ -1906,7 +1915,8 @@ table_index_validate_scan(Relation table_rel,
 static inline uint64
 table_relation_size(Relation rel, ForkNumber forkNumber)
 {
-	return rel->rd_tableam->relation_size(rel, forkNumber);
+	int64 res = rel->rd_tableam->relation_size(rel, forkNumber, DEFAULT_SIZE);
+	return res >= 0 ? res : 0;
 }
 
 /*
@@ -2129,7 +2139,7 @@ extern void table_block_parallelscan_startblock_init(Relation rel,
  * ----------------------------------------------------------------------------
  */
 
-extern uint64 table_block_relation_size(Relation rel, ForkNumber forkNumber);
+extern int64 table_block_relation_size(Relation rel, ForkNumber forkNumber, uint8 method);
 extern void table_block_relation_estimate_size(Relation rel,
 											   int32 *attr_widths,
 											   BlockNumber *pages,