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/batch_queue.h

@@ -24,6 +24,7 @@ typedef struct BatchQueueFunctions
 	void (*push_batch)(struct BatchQueue *, DecompressContext *, TupleTableSlot *);
 	void (*reset)(struct BatchQueue *);
 	TupleTableSlot *(*top_tuple)(struct BatchQueue *);
+	void (*is_done)(struct BatchQueue *);
 } BatchQueueFunctions;
 
 typedef struct BatchQueue

tsl/src/nodes/columnar_scan/batch_queue_fifo.h

@@ -57,13 +57,20 @@ batch_queue_fifo_top_tuple(BatchQueue *bq)
 	return compressed_batch_current_tuple(batch_array_get_at(&bq->batch_array, 0));
 }
 
+static inline void
+batch_queue_fifo_is_done(BatchQueue *bq)
+{
+	/* no-op */
+}
+
 static const struct BatchQueueFunctions BatchQueueFunctionsFifo = {
 	.free = batch_queue_fifo_free,
 	.needs_next_batch = batch_queue_fifo_needs_next_batch,
 	.pop = batch_queue_fifo_pop,
 	.push_batch = batch_queue_fifo_push_batch,
 	.reset = batch_queue_fifo_reset,
 	.top_tuple = batch_queue_fifo_top_tuple,
+	.is_done = batch_queue_fifo_is_done,
 };
 
 extern BatchQueue *batch_queue_fifo_create(int num_compressed_cols,

tsl/src/nodes/columnar_scan/batch_queue_heap.c

@@ -46,6 +46,10 @@ typedef struct BatchQueueHeap
 	 */
 	TupleTableSlot *last_batch_first_tuple_slot;
 	HeapEntryColumn *last_batch_first_tuple_entry;
+
+	/* When we have no more compressed batches but may still have tuples on the heap,
+	   do not check if we need next batch when there are none. */
+	bool done_fetching_compressed_batches;
 } BatchQueueHeap;
 
 /*
@@ -199,6 +203,12 @@ batch_queue_heap_needs_next_batch(BatchQueue *_queue)
 {
 	BatchQueueHeap *queue = (BatchQueueHeap *) _queue;
 
+	/* We consumed all of the compressed batches, time to empty the heap */
+	if (queue->done_fetching_compressed_batches)
+	{
+		return false;
+	}
+
 	if (binaryheap_empty(queue->merge_heap))
 	{
 		return true;
@@ -319,6 +329,7 @@ batch_queue_heap_reset(BatchQueue *bq)
 {
 	BatchQueueHeap *bqh = (BatchQueueHeap *) bq;
 	binaryheap_reset(bqh->merge_heap);
+	bqh->done_fetching_compressed_batches = false;
 }
 
 /*
@@ -343,13 +354,21 @@ batch_queue_heap_free(BatchQueue *_queue)
 	pfree(queue);
 }
 
+static void
+batch_queue_heap_is_done(BatchQueue *bq)
+{
+	BatchQueueHeap *bqh = (BatchQueueHeap *) bq;
+	bqh->done_fetching_compressed_batches = true;
+}
+
 const struct BatchQueueFunctions BatchQueueFunctionsHeap = {
 	.free = batch_queue_heap_free,
 	.needs_next_batch = batch_queue_heap_needs_next_batch,
 	.pop = batch_queue_heap_pop,
 	.push_batch = batch_queue_heap_push_batch,
 	.reset = batch_queue_heap_reset,
 	.top_tuple = batch_queue_heap_top_tuple,
+	.is_done = batch_queue_heap_is_done,
 };
 
 static SortSupport
@@ -412,5 +431,7 @@ batch_queue_heap_create(int num_compressed_cols, const List *sortinfo,
 	queue->last_batch_first_tuple_entry = palloc(sizeof(HeapEntryColumn) * queue->nkeys);
 	queue->queue.funcs = funcs;
 
+	queue->done_fetching_compressed_batches = false;
+
 	return &queue->queue;
 }

tsl/src/nodes/columnar_scan/columnar_scan.c

@@ -978,27 +978,44 @@ 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)
+	{
+		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
@@ -1477,8 +1494,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 +1507,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)
@@ -1517,7 +1538,7 @@ build_on_single_compressed_path(PlannerInfo *root, const Chunk *chunk, RelOptInf
 	 * will determine whether to put an actual sort between the decompression
 	 * node and the scan during plan creation.
 	 */
-	if (sort_info->use_compressed_sort)
+	if (sort_info->use_compressed_sort && !sort_info->use_batch_sorted_merge)
 	{
 		ColumnarScanPath *columnar_scan_with_compressed_sort = NULL;
 		Path dummy_sort_path; /* dummy for result of cost_sort */
@@ -3229,6 +3250,15 @@ 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 && !sort_info.reverse)
+			{
+				sort_info.needs_sequence_num = true;
+				sort_info.use_compressed_sort = true;
+			}
 		}
 		return sort_info;
 	}

tsl/src/nodes/columnar_scan/exec.c

@@ -457,6 +457,7 @@ columnar_scan_exec_impl(ColumnarScanState *chunk_state, const BatchQueueFunction
 		if (TupIsNull(subslot))
 		{
 			/* Won't have more compressed tuples. */
+			bqfuncs->is_done(bq);
 			break;
 		}
 

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