Do not sort batches for Batch sorted merge over unordered chunks if query sort matches compressed sort

Author: natalya-aksman

tsl/src/nodes/columnar_scan/columnar_scan.c

@@ -978,27 +978,45 @@ cost_batch_sorted_merge(PlannerInfo *root, const CompressionInfo *compression_in
 {
 	Path sort_path; /* dummy for result of cost_sort */
 
-	/*
-	 * Don't disable the compressed batch sorted merge plan with the enable_sort
-	 * GUC. We have a separate GUC for it, and this way you can try to force the
-	 * batch sorted merge plan by disabling sort.
-	 */
-	const bool old_enable_sort = enable_sort;
-	enable_sort = true;
-	cost_sort(&sort_path,
-			  root,
-			  dcpath->required_compressed_pathkeys,
+	/* We are utilizing compressed sort for batch sorted merge: do not need extra sort */
+	if (dcpath->required_compressed_pathkeys &&
+		pathkeys_contained_in(dcpath->required_compressed_pathkeys, compressed_path->pathkeys))
+	{
+		sort_path.rows = compressed_path->rows;
+		sort_path.startup_cost = compressed_path->startup_cost;
+		sort_path.total_cost = compressed_path->total_cost;
 #if PG18_GE
-			  compressed_path->disabled_nodes,
+		/* PG18 changes the way we handle disabled nodes so we
+		 * need to take those into account as well.
+		 *
+		 * https://github.com/postgres/postgres/commit/e2225346
+		 */
+		sort_path.disabled_nodes = compressed_path->disabled_nodes;
 #endif
-			  compressed_path->total_cost,
-			  compressed_path->rows,
-			  compressed_path->pathtarget->width,
-			  0.0,
-			  work_mem,
-			  -1);
-	enable_sort = old_enable_sort;
-
+	}
+	else
+	{
+		/*
+		 * Don't disable the compressed batch sorted merge plan with the enable_sort
+		 * GUC. We have a separate GUC for it, and this way you can try to force the
+		 * batch sorted merge plan by disabling sort.
+		 */
+		const bool old_enable_sort = enable_sort;
+		enable_sort = true;
+		cost_sort(&sort_path,
+				  root,
+				  dcpath->required_compressed_pathkeys,
+#if PG18_GE
+				  compressed_path->disabled_nodes,
+#endif
+				  compressed_path->total_cost,
+				  compressed_path->rows,
+				  compressed_path->pathtarget->width,
+				  0.0,
+				  work_mem,
+				  -1);
+		enable_sort = old_enable_sort;
+	}
 	/*
 	 * In compressed batch sorted merge, for each distinct segmentby value we
 	 * have to keep the corresponding latest batch open. Estimate the number of
@@ -1432,6 +1450,18 @@ build_on_single_compressed_path(PlannerInfo *root, const Chunk *chunk, RelOptInf
 		return NIL;
 	}
 
+	if (ts_guc_debug_require_batch_sorted_merge == DRO_Require ||
+		ts_guc_debug_require_batch_sorted_merge == DRO_Force)
+	{
+		if (!(sort_info->use_batch_sorted_merge && ts_guc_enable_decompression_sorted_merge))
+		{
+			ereport(ERROR,
+					(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
+					 errmsg("debug: batch sorted merge is required but not possible at planning "
+							"time")));
+		}
+	}
+
 	/*
 	 * Filter out all paths that try to JOIN the compressed chunk on the
 	 * hypertable or the uncompressed chunk
@@ -1477,8 +1507,6 @@ build_on_single_compressed_path(PlannerInfo *root, const Chunk *chunk, RelOptInf
 	 */
 	if (sort_info->use_batch_sorted_merge && ts_guc_enable_decompression_sorted_merge)
 	{
-		Assert(!sort_info->use_compressed_sort);
-
 		ColumnarScanPath *path_copy =
 			copy_columnar_scan_path((ColumnarScanPath *) chunk_path_no_sort);
 
@@ -1492,6 +1520,12 @@ build_on_single_compressed_path(PlannerInfo *root, const Chunk *chunk, RelOptInf
 		 * query here.
 		 */
 		path_copy->custom_path.path.pathkeys = sort_info->decompressed_sort_pathkeys;
+
+		/* Batch sorted merge over unordered chunk can utilize compressed sort, copy the relevant
+		 * fields */
+		path_copy->needs_sequence_num = sort_info->needs_sequence_num;
+		path_copy->required_compressed_pathkeys = sort_info->required_compressed_pathkeys;
+
 		cost_batch_sorted_merge(root, compression_info, path_copy, compressed_path);
 
 		if (ts_guc_debug_require_batch_sorted_merge == DRO_Force)
@@ -1502,22 +1536,13 @@ build_on_single_compressed_path(PlannerInfo *root, const Chunk *chunk, RelOptInf
 
 		decompressed_paths = lappend(decompressed_paths, path_copy);
 	}
-	else if (ts_guc_debug_require_batch_sorted_merge == DRO_Require ||
-			 ts_guc_debug_require_batch_sorted_merge == DRO_Force)
-	{
-		ereport(ERROR,
-				(errcode(ERRCODE_OBJECT_NOT_IN_PREREQUISITE_STATE),
-				 errmsg("debug: batch sorted merge is required but not possible at planning "
-						"time")));
-	}
-
 	/*
 	 * If we can push down the sort below the ColumnarScan node, we set the
 	 * pathkeys of the decompress node to the decompressed_sort_pathkeys. We
 	 * will determine whether to put an actual sort between the decompression
 	 * node and the scan during plan creation.
 	 */
-	if (sort_info->use_compressed_sort)
+	else if (sort_info->use_compressed_sort)
 	{
 		ColumnarScanPath *columnar_scan_with_compressed_sort = NULL;
 		Path dummy_sort_path; /* dummy for result of cost_sort */
@@ -3229,6 +3254,18 @@ build_sortinfo(PlannerInfo *root, const Chunk *chunk, RelOptInfo *chunk_rel,
 													  compression_info,
 													  /* for_batch_sorted_merge = */ true,
 													  &sort_info.reverse);
+
+			/* Pathkeys are matching leading orderby metadata column:
+			 * can use already sorted compressed data for batch sorted merge.
+			 */
+			if (sort_info.use_batch_sorted_merge && ts_guc_enable_decompression_sorted_merge &&
+				!sort_info.reverse && list_length(pathkeys) == 1 &&
+				/* Needs to be sorted on metadata, not sequence number */
+				!compression_info->has_seq_num)
+			{
+				sort_info.needs_sequence_num = true;
+				sort_info.use_compressed_sort = true;
+			}
 		}
 		return sort_info;
 	}

tsl/src/nodes/columnar_scan/exec.c

@@ -98,6 +98,7 @@ columnar_scan_state_create(CustomScan *cscan)
 	Assert(list_length(chunk_state->decompression_map) ==
 		   list_length(chunk_state->is_segmentby_column));
 
+	chunk_state->done_fetching_batches = false;
 	return (Node *) chunk_state;
 }
 
@@ -201,6 +202,7 @@ columnar_scan_begin(CustomScanState *node, EState *estate, int eflags)
 	Plan *compressed_scan = linitial(cscan->custom_plans);
 	Assert(list_length(cscan->custom_plans) == 1);
 
+	chunk_state->done_fetching_batches = false;
 	ts_stats_compression_acc_init(&dcontext->observ_acc);
 
 	PlanState *ps = &node->ss.ps;
@@ -451,12 +453,13 @@ columnar_scan_exec_impl(ColumnarScanState *chunk_state, const BatchQueueFunction
 
 	bqfuncs->pop(bq, dcontext);
 
-	while (bqfuncs->needs_next_batch(bq))
+	while (!chunk_state->done_fetching_batches && bqfuncs->needs_next_batch(bq))
 	{
 		TupleTableSlot *subslot = ExecProcNode(linitial(chunk_state->csstate.custom_ps));
 		if (TupIsNull(subslot))
 		{
 			/* Won't have more compressed tuples. */
+			chunk_state->done_fetching_batches = true;
 			break;
 		}
 
@@ -491,6 +494,7 @@ static void
 columnar_scan_rescan(CustomScanState *node)
 {
 	ColumnarScanState *chunk_state = (ColumnarScanState *) node;
+	chunk_state->done_fetching_batches = false;
 	BatchQueue *bq = chunk_state->batch_queue;
 
 	bq->funcs->reset(bq);

tsl/src/nodes/columnar_scan/exec.h

@@ -42,6 +42,10 @@ typedef struct ColumnarScanState
 	 * evaluate to constant false, hence the flag.
 	 */
 	List *vectorized_quals_original;
+
+	/* We may be done with compressed batches but not done with the heap,
+	 * let ColumnarScanState know to stop fetching batches in this case. */
+	bool done_fetching_batches;
 } ColumnarScanState;
 
 extern Node *columnar_scan_state_create(CustomScan *cscan);

tsl/src/nodes/columnar_scan/planner.c

@@ -1298,82 +1298,95 @@ columnar_scan_plan_create(PlannerInfo *root, RelOptInfo *rel, CustomPath *path,
 
 		sort_options = list_make4(sort_col_idx, sort_ops, sort_collations, sort_nulls);
 
-		/*
-		 * Build a sort node for the compressed batches. The sort function is
-		 * derived from the sort function of the pathkeys, except that it refers
-		 * to the min and max metadata columns of the batches. We have already
-		 * verified that the pathkeys match the compression order_by, so this
-		 * mapping is possible.
-		 */
-		AttrNumber *sortColIdx = palloc(sizeof(AttrNumber) * numsortkeys);
-		Oid *sortOperators = palloc(sizeof(Oid) * numsortkeys);
-		Oid *collations = palloc(sizeof(Oid) * numsortkeys);
-		bool *nullsFirst = palloc(sizeof(bool) * numsortkeys);
-		for (int i = 0; i < numsortkeys; i++)
+		/* We can utilize compressed sort for batch sorted merge over unordered chunks: do not need
+		 * to add extra Sort node */
+		if (dcpath->required_compressed_pathkeys &&
+			pathkeys_contained_in(dcpath->required_compressed_pathkeys, compressed_path->pathkeys))
+		{
+			decompress_plan->custom_plans = custom_plans;
+		}
+		else
 		{
-			Oid sortop = list_nth_oid(sort_ops, i);
-
-			/* Find the operator in pg_amop --- failure shouldn't happen */
-			Oid opfamily, opcintype;
-			CompareType strategy;
-			if (!get_ordering_op_properties(list_nth_oid(sort_ops, i),
-											&opfamily,
-											&opcintype,
-											&strategy))
-			{
-				elog(ERROR, "operator %u is not a valid ordering operator", sortOperators[i]);
-			}
-
 			/*
-			 * This way to determine the matching metadata column works, because
-			 * we have already verified that the pathkeys match the compression
-			 * orderby.
+			 * Build a sort node for the compressed batches. The sort function is
+			 * derived from the sort function of the pathkeys, except that it refers
+			 * to the min and max metadata columns of the batches. We have already
+			 * verified that the pathkeys match the compression order_by, so this
+			 * mapping is possible.
 			 */
-			Assert(strategy == BTLessStrategyNumber || strategy == BTGreaterStrategyNumber);
-			char *lower_name;
-			char *upper_name;
-			orderby_sparse_metadata_names(dcpath->info->settings, i + 1, &lower_name, &upper_name);
-			char *meta_col_name = strategy == BTLessStrategyNumber ? lower_name : upper_name;
+			AttrNumber *sortColIdx = palloc(sizeof(AttrNumber) * numsortkeys);
+			Oid *sortOperators = palloc(sizeof(Oid) * numsortkeys);
+			Oid *collations = palloc(sizeof(Oid) * numsortkeys);
+			bool *nullsFirst = palloc(sizeof(bool) * numsortkeys);
+			for (int i = 0; i < numsortkeys; i++)
+			{
+				Oid sortop = list_nth_oid(sort_ops, i);
+
+				/* Find the operator in pg_amop --- failure shouldn't happen */
+				Oid opfamily, opcintype;
+				CompareType strategy;
+				if (!get_ordering_op_properties(list_nth_oid(sort_ops, i),
+												&opfamily,
+												&opcintype,
+												&strategy))
+				{
+					elog(ERROR, "operator %u is not a valid ordering operator", sortOperators[i]);
+				}
 
-			AttrNumber attr_position =
-				get_attnum(dcpath->info->compressed_rte->relid, meta_col_name);
+				/*
+				 * This way to determine the matching metadata column works, because
+				 * we have already verified that the pathkeys match the compression
+				 * orderby.
+				 */
+				Assert(strategy == BTLessStrategyNumber || strategy == BTGreaterStrategyNumber);
+				char *lower_name;
+				char *upper_name;
+				orderby_sparse_metadata_names(dcpath->info->settings,
+											  i + 1,
+											  &lower_name,
+											  &upper_name);
+				char *meta_col_name = strategy == BTLessStrategyNumber ? lower_name : upper_name;
+
+				AttrNumber attr_position =
+					get_attnum(dcpath->info->compressed_rte->relid, meta_col_name);
+
+				if (attr_position == InvalidAttrNumber)
+				{
+					elog(ERROR, "couldn't find metadata column \"%s\"", meta_col_name);
+				}
 
-			if (attr_position == InvalidAttrNumber)
-			{
-				elog(ERROR, "couldn't find metadata column \"%s\"", meta_col_name);
-			}
+				/*
+				 * If the compressed target list is not based on the layout of
+				 * the uncompressed chunk (see comment for physical_tlist above),
+				 * adjust the position of the attribute.
+				 */
+				if (target_list_compressed_is_physical)
+				{
+					sortColIdx[i] = attr_position;
+				}
+				else
+				{
+					sortColIdx[i] =
+						find_attr_pos_in_tlist(compressed_scan->plan.targetlist, attr_position);
+				}
 
-			/*
-			 * If the compressed target list is not based on the layout of
-			 * the uncompressed chunk (see comment for physical_tlist above),
-			 * adjust the position of the attribute.
-			 */
-			if (target_list_compressed_is_physical)
-			{
-				sortColIdx[i] = attr_position;
+				sortOperators[i] = sortop;
+				collations[i] = list_nth_oid(sort_collations, i);
+				nullsFirst[i] = list_nth_oid(sort_nulls, i);
 			}
-			else
-			{
-				sortColIdx[i] =
-					find_attr_pos_in_tlist(compressed_scan->plan.targetlist, attr_position);
-			}
-
-			sortOperators[i] = sortop;
-			collations[i] = list_nth_oid(sort_collations, i);
-			nullsFirst[i] = list_nth_oid(sort_nulls, i);
-		}
 
-		/* Now build the compressed batches sort node */
-		Sort *sort = ts_make_sort((Plan *) compressed_scan,
-								  numsortkeys,
-								  sortColIdx,
-								  sortOperators,
-								  collations,
-								  nullsFirst);
+			/* Now build the compressed batches sort node */
+			Sort *sort = ts_make_sort((Plan *) compressed_scan,
+									  numsortkeys,
+									  sortColIdx,
+									  sortOperators,
+									  collations,
+									  nullsFirst);
 
-		ts_label_sort_with_costsize(root, sort, /* limit_tuples = */ -1.0);
+			ts_label_sort_with_costsize(root, sort, /* limit_tuples = */ -1.0);
 
-		decompress_plan->custom_plans = list_make1(sort);
+			decompress_plan->custom_plans = list_make1(sort);
+		}
 	}
 	else
 	{

tsl/test/expected/compress_unordered_sort.out

@@ -369,10 +369,8 @@ SET timescaledb.debug_require_batch_sorted_merge = 'force';
  GroupAggregate
    Group Key: _hyper_1_1_chunk."time"
    ->  Custom Scan (ColumnarScan) on _hyper_1_1_chunk
-         ->  Sort
-               Sort Key: _hyper_1_1_chunk_compressed._ts_meta_v2_first_time DESC
-               ->  Index Scan using _hyper_1_1_chunk_compressed_device_sensor__ts_meta_v2_first_idx on _hyper_1_1_chunk_compressed
-                     Index Cond: ((device = 'd1'::text) AND (sensor = 'A'::text))
+         ->  Index Scan using _hyper_1_1_chunk_compressed_device_sensor__ts_meta_v2_first_idx on _hyper_1_1_chunk_compressed
+               Index Cond: ((device = 'd1'::text) AND (sensor = 'A'::text))
 
 select time, avg(value) from metrics  where device = 'd1' and sensor='A' group by time order by time DESC;
              time             |  avg