Estimate cost of read_intermediate_results()

src/backend/distributed/metadata/metadata_cache.c

@@ -136,6 +136,7 @@ typedef struct MetadataCacheData
 	Oid citusCatalogNamespaceId;
 	Oid copyFormatTypeId;
 	Oid readIntermediateResultFuncId;
+	Oid readIntermediateResultArrayFuncId;
 	Oid extraDataContainerFuncId;
 	Oid workerHashFunctionId;
 	Oid anyValueFunctionId;
@@ -2065,6 +2066,26 @@ CitusReadIntermediateResultFuncId(void)
 }
 
 
+/* return oid of the read_intermediate_results(text[],citus_copy_format) function */
+Oid
+CitusReadIntermediateResultArrayFuncId(void)
+{
+	if (MetadataCache.readIntermediateResultArrayFuncId == InvalidOid)
+	{
+		List *functionNameList = list_make2(makeString("pg_catalog"),
+											makeString("read_intermediate_results"));
+		Oid copyFormatTypeOid = CitusCopyFormatTypeId();
+		Oid paramOids[2] = { TEXTARRAYOID, copyFormatTypeOid };
+		bool missingOK = false;
+
+		MetadataCache.readIntermediateResultArrayFuncId =
+			LookupFuncName(functionNameList, 2, paramOids, missingOK);
+	}
+
+	return MetadataCache.readIntermediateResultArrayFuncId;
+}
+
+
 /* return oid of the citus.copy_format enum type */
 Oid
 CitusCopyFormatTypeId(void)

src/backend/distributed/planner/distributed_planner.c

@@ -99,6 +99,13 @@ static void CheckNodeIsDumpable(Node *node);
 static Node * CheckNodeCopyAndSerialization(Node *node);
 static void AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry,
 											 RelOptInfo *relOptInfo);
+static void AdjustReadIntermediateResultArrayCost(RangeTblEntry *rangeTableEntry,
+												  RelOptInfo *relOptInfo);
+static void AdjustReadIntermediateResultsCostInternal(RelOptInfo *relOptInfo,
+													  List *columnTypes,
+													  int resultIdCount,
+													  Datum *resultIds,
+													  Const *resultFormatConst);
 static List * OuterPlanParamsList(PlannerInfo *root);
 static List * CopyPlanParamList(List *originalPlanParamList);
 static PlannerRestrictionContext * CreateAndPushPlannerRestrictionContext(void);
@@ -1517,6 +1524,7 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 	DistTableCacheEntry *cacheEntry = NULL;
 
 	AdjustReadIntermediateResultCost(rte, relOptInfo);
+	AdjustReadIntermediateResultArrayCost(rte, relOptInfo);
 
 	if (rte->rtekind != RTE_RELATION)
 	{
@@ -1578,30 +1586,6 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo,
 static void
 AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry, RelOptInfo *relOptInfo)
 {
-	PathTarget *reltarget = relOptInfo->reltarget;
-	List *pathList = relOptInfo->pathlist;
-	Path *path = NULL;
-	RangeTblFunction *rangeTableFunction = NULL;
-	FuncExpr *funcExpression = NULL;
-	Const *resultFormatConst = NULL;
-	Datum resultFormatDatum = 0;
-	Oid resultFormatId = InvalidOid;
-	Const *resultIdConst = NULL;
-	Datum resultIdDatum = 0;
-	char *resultId = NULL;
-	int64 resultSize = 0;
-	ListCell *typeCell = NULL;
-	bool binaryFormat = false;
-	double rowCost = 0.;
-	double rowSizeEstimate = 0;
-	double rowCountEstimate = 0.;
-	double ioCost = 0.;
-#if PG_VERSION_NUM >= 120000
-	QualCost funcCost = { 0., 0. };
-#else
-	double funcCost = 0.;
-#endif
-
 	if (rangeTableEntry->rtekind != RTE_FUNCTION ||
 		list_length(rangeTableEntry->functions) != 1)
 	{
@@ -1620,41 +1604,133 @@ AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry, RelOptInfo *rel
 		return;
 	}
 
-	rangeTableFunction = (RangeTblFunction *) linitial(rangeTableEntry->functions);
+	RangeTblFunction *rangeTableFunction = (RangeTblFunction *) linitial(
-	funcExpression = (FuncExpr *) rangeTableFunction->funcexpr;
+		rangeTableEntry->functions);
-	resultIdConst = (Const *) linitial(funcExpression->args);
+	FuncExpr *funcExpression = (FuncExpr *) rangeTableFunction->funcexpr;
+	Const *resultIdConst = (Const *) linitial(funcExpression->args);
 	if (!IsA(resultIdConst, Const))
 	{
 		/* not sure how to interpret non-const */
 		return;
 	}
 
-	resultIdDatum = resultIdConst->constvalue;
+	Datum resultIdDatum = resultIdConst->constvalue;
-	resultId = TextDatumGetCString(resultIdDatum);
 
-	resultSize = IntermediateResultSize(resultId);
+	Const *resultFormatConst = (Const *) lsecond(funcExpression->args);
-	if (resultSize < 0)
-	{
-		/* result does not exist, will probably error out later on */
-		return;
-	}
-
-	resultFormatConst = (Const *) lsecond(funcExpression->args);
 	if (!IsA(resultFormatConst, Const))
 	{
 		/* not sure how to interpret non-const */
 		return;
 	}
 
-	resultFormatDatum = resultFormatConst->constvalue;
+	AdjustReadIntermediateResultsCostInternal(relOptInfo,
-	resultFormatId = DatumGetObjectId(resultFormatDatum);
+											  rangeTableFunction->funccoltypes,
+											  1, &resultIdDatum, resultFormatConst);
+}
 
-	if (resultFormatId == BinaryCopyFormatId())
+
+/*
+ * AdjustReadIntermediateResultArrayCost adjusts the row count and total cost
+ * of a read_intermediate_results(resultIds, format) call based on the file size.
+ */
+static void
+AdjustReadIntermediateResultArrayCost(RangeTblEntry *rangeTableEntry,
+									  RelOptInfo *relOptInfo)
 {
-		binaryFormat = true;
+	Datum *resultIdArray = NULL;
+	int resultIdCount = 0;
 
+	if (rangeTableEntry->rtekind != RTE_FUNCTION ||
+		list_length(rangeTableEntry->functions) != 1)
+	{
+		/* avoid more expensive checks below for non-functions */
+		return;
+	}
+
+	if (!CitusHasBeenLoaded() || !CheckCitusVersion(DEBUG5))
+	{
+		/* read_intermediate_result may not exist */
+		return;
+	}
+
+	if (!ContainsReadIntermediateResultArrayFunction((Node *) rangeTableEntry->functions))
+	{
+		return;
+	}
+
+	RangeTblFunction *rangeTableFunction =
+		(RangeTblFunction *) linitial(rangeTableEntry->functions);
+	FuncExpr *funcExpression = (FuncExpr *) rangeTableFunction->funcexpr;
+	Const *resultIdConst = (Const *) linitial(funcExpression->args);
+	if (!IsA(resultIdConst, Const))
+	{
+		/* not sure how to interpret non-const */
+		return;
+	}
+
+	Datum resultIdArrayDatum = resultIdConst->constvalue;
+	deconstruct_array(DatumGetArrayTypeP(resultIdArrayDatum), TEXTOID, -1, false,
+					  'i', &resultIdArray, NULL, &resultIdCount);
+
+	Const *resultFormatConst = (Const *) lsecond(funcExpression->args);
+	if (!IsA(resultFormatConst, Const))
+	{
+		/* not sure how to interpret non-const */
+		return;
+	}
+
+	AdjustReadIntermediateResultsCostInternal(relOptInfo,
+											  rangeTableFunction->funccoltypes,
+											  resultIdCount, resultIdArray,
+											  resultFormatConst);
+}
+
+
+/*
+ * AdjustReadIntermediateResultsCostInternal adjusts the row count and total cost
+ * of reading intermediate results based on file sizes.
+ */
+static void
+AdjustReadIntermediateResultsCostInternal(RelOptInfo *relOptInfo, List *columnTypes,
+										  int resultIdCount, Datum *resultIds,
+										  Const *resultFormatConst)
+{
+	PathTarget *reltarget = relOptInfo->reltarget;
+	List *pathList = relOptInfo->pathlist;
+	Path *path = NULL;
+	double rowCost = 0.;
+	double rowSizeEstimate = 0;
+	double rowCountEstimate = 0.;
+	double ioCost = 0.;
+#if PG_VERSION_NUM >= 120000
+	QualCost funcCost = { 0., 0. };
+#else
+	double funcCost = 0.;
+#endif
+	int64 totalResultSize = 0;
+	ListCell *typeCell = NULL;
+
+	Datum resultFormatDatum = resultFormatConst->constvalue;
+	Oid resultFormatId = DatumGetObjectId(resultFormatDatum);
+	bool binaryFormat = (resultFormatId == BinaryCopyFormatId());
+
+	for (int index = 0; index < resultIdCount; index++)
+	{
+		char *resultId = TextDatumGetCString(resultIds[index]);
+		int64 resultSize = IntermediateResultSize(resultId);
+		if (resultSize < 0)
+		{
+			/* result does not exist, will probably error out later on */
+			return;
+		}
+
+		if (binaryFormat)
+		{
 			/* subtract 11-byte signature + 8 byte header + 2-byte footer */
-		resultSize -= 21;
+			totalResultSize -= 21;
+		}
+
+		totalResultSize += resultSize;
 	}
 
 	/* start with the cost of evaluating quals */
@@ -1666,7 +1742,7 @@ AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry, RelOptInfo *rel
 	/* add 2 bytes for column count (binary) or line separator (text) */
 	rowSizeEstimate += 2;
 
-	foreach(typeCell, rangeTableFunction->funccoltypes)
+	foreach(typeCell, columnTypes)
 	{
 		Oid columnTypeId = lfirst_oid(typeCell);
 		Oid inputFunctionId = InvalidOid;
@@ -1702,10 +1778,10 @@ AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry, RelOptInfo *rel
 #endif
 
 	/* estimate the number of rows based on the file size and estimated row size */
-	rowCountEstimate = Max(1, (double) resultSize / rowSizeEstimate);
+	rowCountEstimate = Max(1, (double) totalResultSize / rowSizeEstimate);
 
 	/* cost of reading the data */
-	ioCost = seq_page_cost * resultSize / BLCKSZ;
+	ioCost = seq_page_cost * totalResultSize / BLCKSZ;
 
 	Assert(pathList != NIL);
 

src/backend/distributed/planner/multi_logical_planner.c

@@ -78,6 +78,8 @@ static bool ErrorHintRequired(const char *errorHint, Query *queryTree);
 static bool HasTablesample(Query *queryTree);
 static bool HasComplexRangeTableType(Query *queryTree);
 static bool IsReadIntermediateResultFunction(Node *node);
+static bool IsReadIntermediateResultArrayFunction(Node *node);
+static bool IsFunctionWithOid(Node *node, Oid funcOid);
 static bool ExtractFromExpressionWalker(Node *node,
 										QualifierWalkerContext *walkerContext);
 static List * MultiTableNodeList(List *tableEntryList, List *rangeTableList);
@@ -768,18 +770,52 @@ ContainsReadIntermediateResultFunction(Node *node)
 }
 
 
+/*
+ * ContainsReadIntermediateResultArrayFunction determines whether an expresion
+ * tree contains a call to the read_intermediate_results(result_ids, format)
+ * function.
+ */
+bool
+ContainsReadIntermediateResultArrayFunction(Node *node)
+{
+	return FindNodeCheck(node, IsReadIntermediateResultArrayFunction);
+}
+
+
 /*
  * IsReadIntermediateResultFunction determines whether a given node is a function call
  * to the read_intermediate_result function.
  */
 static bool
 IsReadIntermediateResultFunction(Node *node)
+{
+	return IsFunctionWithOid(node, CitusReadIntermediateResultFuncId());
+}
+
+
+/*
+ * IsReadIntermediateResultArrayFunction determines whether a given node is a
+ * function call to the read_intermediate_results(result_ids, format) function.
+ */
+static bool
+IsReadIntermediateResultArrayFunction(Node *node)
+{
+	return IsFunctionWithOid(node, CitusReadIntermediateResultArrayFuncId());
+}
+
+
+/*
+ * IsFunctionWithOid determines whether a given node is a function call
+ * to the read_intermediate_result function.
+ */
+static bool
+IsFunctionWithOid(Node *node, Oid funcOid)
 {
 	if (IsA(node, FuncExpr))
 	{
 		FuncExpr *funcExpr = (FuncExpr *) node;
 
-		if (funcExpr->funcid == CitusReadIntermediateResultFuncId())
+		if (funcExpr->funcid == funcOid)
 		{
 			return true;
 		}

src/include/distributed/metadata_cache.h

@@ -188,6 +188,7 @@ extern Oid CitusCopyFormatTypeId(void);
 
 /* function oids */
 extern Oid CitusReadIntermediateResultFuncId(void);
+Oid CitusReadIntermediateResultArrayFuncId(void);
 extern Oid CitusExtraDataContainerFuncId(void);
 extern Oid CitusWorkerHashFunctionId(void);
 extern Oid CitusAnyValueFunctionId(void);

src/include/distributed/multi_logical_planner.h

@@ -193,6 +193,7 @@ extern bool FindNodeCheckInRangeTableList(List *rtable, bool (*check)(Node *));
 extern bool IsDistributedTableRTE(Node *node);
 extern bool QueryContainsDistributedTableRTE(Query *query);
 extern bool ContainsReadIntermediateResultFunction(Node *node);
+extern bool ContainsReadIntermediateResultArrayFunction(Node *node);
 extern char * FindIntermediateResultIdIfExists(RangeTblEntry *rte);
 extern MultiNode * ParentNode(MultiNode *multiNode);
 extern MultiNode * ChildNode(MultiUnaryNode *multiNode);

src/test/regress/expected/intermediate_results.out

@@ -201,10 +201,10 @@ SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series
                         632
 (1 row)
 
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
                                     QUERY PLAN                                     
------------------------------------------------
+-----------------------------------------------------------------------------------
- Function Scan on read_intermediate_result res
+ Function Scan on read_intermediate_result res  (cost=0.00..4.55 rows=632 width=8)
 (1 row)
 
 -- less accurate results for variable types
@@ -214,10 +214,10 @@ SELECT create_intermediate_result('hellos', $$SELECT s, 'hello-'||s FROM generat
                          63
 (1 row)
 
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('hellos', 'binary') AS res (x int, y text);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('hellos', 'binary') AS res (x int, y text);
                                     QUERY PLAN                                     
------------------------------------------------
+-----------------------------------------------------------------------------------
- Function Scan on read_intermediate_result res
+ Function Scan on read_intermediate_result res  (cost=0.00..0.32 rows=30 width=36)
 (1 row)
 
 -- not very accurate results for text encoding
@@ -227,10 +227,10 @@ SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_s
                           4
 (1 row)
 
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
                                     QUERY PLAN                                    
------------------------------------------------
+----------------------------------------------------------------------------------
- Function Scan on read_intermediate_result res
+ Function Scan on read_intermediate_result res  (cost=0.00..0.01 rows=1 width=32)
 (1 row)
 
 END;
@@ -368,6 +368,50 @@ ON (squares.x::text = interested_in) WHERE user_id = 'jon' ORDER BY 1,2;
  jon     | 5             | 5 | 25 | (2,3)
 (2 rows)
 
+END;
+-- Cost estimation for read_intermediate_results
+BEGIN;
+-- almost accurate row count estimates for primitive types
+SELECT create_intermediate_result('squares_1', 'SELECT s, s*s FROM generate_series(1,632) s'),
+       create_intermediate_result('squares_2', 'SELECT s, s*s FROM generate_series(633,1024) s');
+ create_intermediate_result | create_intermediate_result 
+----------------------------+----------------------------
+                        632 |                        392
+(1 row)
+
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['squares_1', 'squares_2'], 'binary') AS res (x int, x2 int);
+                                     QUERY PLAN                                      
+-------------------------------------------------------------------------------------
+ Function Scan on read_intermediate_results res  (cost=0.00..7.37 rows=1024 width=8)
+(1 row)
+
+-- less accurate results for variable types
+SELECT create_intermediate_result('hellos_1', $$SELECT s, 'hello-'||s FROM generate_series(1,63) s$$),
+       create_intermediate_result('hellos_2', $$SELECT s, 'hello-'||s FROM generate_series(64,129) s$$);
+ create_intermediate_result | create_intermediate_result 
+----------------------------+----------------------------
+                         63 |                         66
+(1 row)
+
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['hellos_1', 'hellos_2'], 'binary') AS res (x int, y text);
+                                     QUERY PLAN                                     
+------------------------------------------------------------------------------------
+ Function Scan on read_intermediate_results res  (cost=0.00..0.66 rows=62 width=36)
+(1 row)
+
+-- not very accurate results for text encoding
+SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
+ create_intermediate_result 
+----------------------------
+                          4
+(1 row)
+
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['stored_squares'], 'text') AS res (s intermediate_results.square_type);
+                                    QUERY PLAN                                     
+-----------------------------------------------------------------------------------
+ Function Scan on read_intermediate_results res  (cost=0.00..0.01 rows=1 width=32)
+(1 row)
+
 END;
 DROP SCHEMA intermediate_results CASCADE;
 NOTICE:  drop cascades to 5 other objects

src/test/regress/sql/intermediate_results.sql

@@ -121,15 +121,15 @@ END;
 BEGIN;
 -- accurate row count estimates for primitive types
 SELECT create_intermediate_result('squares', 'SELECT s, s*s FROM generate_series(1,632) s');
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('squares', 'binary') AS res (x int, x2 int);
 
 -- less accurate results for variable types
 SELECT create_intermediate_result('hellos', $$SELECT s, 'hello-'||s FROM generate_series(1,63) s$$);
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('hellos', 'binary') AS res (x int, y text);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('hellos', 'binary') AS res (x int, y text);
 
 -- not very accurate results for text encoding
 SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
-EXPLAIN (COSTS OFF) SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_result('stored_squares', 'text') AS res (s intermediate_results.square_type);
 END;
 
 -- pipe query output into a result file and create a table to check the result
@@ -201,5 +201,21 @@ ON (squares.x::text = interested_in) WHERE user_id = 'jon' ORDER BY 1,2;
 
 END;
 
+-- Cost estimation for read_intermediate_results
+BEGIN;
+-- almost accurate row count estimates for primitive types
+SELECT create_intermediate_result('squares_1', 'SELECT s, s*s FROM generate_series(1,632) s'),
+       create_intermediate_result('squares_2', 'SELECT s, s*s FROM generate_series(633,1024) s');
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['squares_1', 'squares_2'], 'binary') AS res (x int, x2 int);
+
+-- less accurate results for variable types
+SELECT create_intermediate_result('hellos_1', $$SELECT s, 'hello-'||s FROM generate_series(1,63) s$$),
+       create_intermediate_result('hellos_2', $$SELECT s, 'hello-'||s FROM generate_series(64,129) s$$);
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['hellos_1', 'hellos_2'], 'binary') AS res (x int, y text);
+
+-- not very accurate results for text encoding
+SELECT create_intermediate_result('stored_squares', 'SELECT square FROM stored_squares');
+EXPLAIN (COSTS ON) SELECT * FROM read_intermediate_results(ARRAY['stored_squares'], 'text') AS res (s intermediate_results.square_type);
+END;
 
 DROP SCHEMA intermediate_results CASCADE;