Add limit option to scanner

This adds an option to set a limit on how many tuples to return in a
relation scan using the scanner implemention. This avoids a common
pattern of manually implementing limits in the tuple handling
function.

src/chunk.c

@@ -534,7 +534,6 @@ chunk_collision_scan(ChunkScanCtx *scanctx, Hypercube *cube)
 	}
 }
 
-
 /*
  * Apply a function to each chunk in the scan context's hash table. If the limit
  * is greater than zero only a limited number of chunks will be processed.

src/chunk_constraint.c

@@ -31,9 +31,6 @@ chunk_constraint_tuple_found(TupleInfo *ti, void *data)
 
 	chunk_constraint_fill(&chunk->constraints[chunk->num_constraints++], ti->tuple);
 
-	if (chunk->capacity == chunk->num_constraints)
-		return false;
-
 	return true;
 }
 
@@ -55,6 +52,7 @@ chunk_constraint_scan_by_chunk_id(Chunk *chunk)
 		.scantype = ScannerTypeIndex,
 		.nkeys = 1,
 		.scankey = scankey,
+		.limit = chunk->num_constraints,
 		.data = chunk,
 		.filter = chunk_constraint_for_dimension_slice,
 		.tuple_found = chunk_constraint_tuple_found,

src/dimension.c

@@ -255,6 +255,7 @@ dimension_scan(int32 hypertable_id, Oid main_table_relid, int16 num_dimensions)
 		.index = catalog->tables[DIMENSION].index_ids[DIMENSION_HYPERTABLE_ID_IDX],
 		.scantype = ScannerTypeIndex,
 		.nkeys = 1,
+		.limit = num_dimensions,
 		.scankey = scankey,
 		.data = space,
 		.tuple_found = dimension_tuple_found,

src/dimension_slice.c

@@ -53,18 +53,15 @@ typedef struct DimensionSliceScanData
 {
 	DimensionVec *slices;
 	int			limit;
-} DimensionSliceScanData;
+}	DimensionSliceScanData;
 
 static bool
 dimension_vec_tuple_found(TupleInfo *ti, void *data)
 {
-	DimensionSliceScanData *scandata = data;
+	DimensionVec *slices = data;
 	DimensionSlice *slice = dimension_slice_from_tuple(ti->tuple);
 
-	dimension_vec_add_slice(&scandata->slices, slice);
-
-	if (scandata->limit == ti->count)
-		return false;
+	dimension_vec_add_slice(&slices, slice);
 
 	return true;
 }
@@ -73,7 +70,8 @@ static int
 dimension_slice_scan_limit_internal(ScanKeyData *scankey,
 									Size num_scankeys,
 									tuple_found_func on_tuple_found,
-									void *scandata)
+									void *scandata,
+									int limit)
 {
 	Catalog    *catalog = catalog_get();
 	ScannerCtx	scanCtx = {
@@ -83,6 +81,7 @@ dimension_slice_scan_limit_internal(ScanKeyData *scankey,
 		.nkeys = num_scankeys,
 		.scankey = scankey,
 		.data = scandata,
+		.limit = limit,
 		.tuple_found = on_tuple_found,
 		.lockmode = AccessShareLock,
 		.scandirection = ForwardScanDirection,
@@ -100,10 +99,7 @@ DimensionVec *
 dimension_slice_scan_limit(int32 dimension_id, int64 coordinate, int limit)
 {
 	ScanKeyData scankey[3];
-	DimensionSliceScanData data = {
-		.slices = dimension_vec_create(limit > 0 ? limit : DIMENSION_VEC_DEFAULT_SIZE),
-		.limit = limit,
-	};
+	DimensionVec *slices = dimension_vec_create(limit > 0 ? limit : DIMENSION_VEC_DEFAULT_SIZE);
 
 	/*
 	 * Perform an index scan for slices matching the dimension's ID and which
@@ -116,9 +112,9 @@ dimension_slice_scan_limit(int32 dimension_id, int64 coordinate, int limit)
 	ScanKeyInit(&scankey[2], Anum_dimension_slice_dimension_id_range_start_range_end_idx_range_end,
 				BTGreaterStrategyNumber, F_INT8GT, Int64GetDatum(coordinate));
 
-	dimension_slice_scan_limit_internal(scankey, 3, dimension_vec_tuple_found, &data);
+	dimension_slice_scan_limit_internal(scankey, 3, dimension_vec_tuple_found, slices, limit);
 
-	return dimension_vec_sort(&data.slices);
+	return dimension_vec_sort(&slices);
 }
 
 /*
@@ -130,10 +126,7 @@ DimensionVec *
 dimension_slice_collision_scan_limit(int32 dimension_id, int64 range_start, int64 range_end, int limit)
 {
 	ScanKeyData scankey[3];
-	DimensionSliceScanData data = {
-		.slices = dimension_vec_create(limit > 0 ? limit : DIMENSION_VEC_DEFAULT_SIZE),
-		.limit = limit,
-	};
+	DimensionVec *slices = dimension_vec_create(limit > 0 ? limit : DIMENSION_VEC_DEFAULT_SIZE);
 
 	ScanKeyInit(&scankey[0], Anum_dimension_slice_dimension_id_range_start_range_end_idx_dimension_id,
 				BTEqualStrategyNumber, F_INT4EQ, Int32GetDatum(dimension_id));
@@ -142,9 +135,9 @@ dimension_slice_collision_scan_limit(int32 dimension_id, int64 range_start, int6
 	ScanKeyInit(&scankey[2], Anum_dimension_slice_dimension_id_range_start_range_end_idx_range_end,
 			  BTGreaterStrategyNumber, F_INT8GT, Int64GetDatum(range_start));
 
-	dimension_slice_scan_limit_internal(scankey, 3, dimension_vec_tuple_found, &data);
+	dimension_slice_scan_limit_internal(scankey, 3, dimension_vec_tuple_found, slices, limit);
 
-	return dimension_vec_sort(&data.slices);
+	return dimension_vec_sort(&slices);
 }
 
 
@@ -173,7 +166,7 @@ dimension_slice_scan_for_existing(DimensionSlice *slice)
 	ScanKeyInit(&scankey[2], Anum_dimension_slice_dimension_id_range_start_range_end_idx_range_end,
 		BTEqualStrategyNumber, F_INT8EQ, Int64GetDatum(slice->fd.range_end));
 
-	dimension_slice_scan_limit_internal(scankey, 3, dimension_slice_fill, &slice);
+	dimension_slice_scan_limit_internal(scankey, 3, dimension_slice_fill, &slice, 1);
 
 	return slice;
 }
@@ -200,6 +193,7 @@ dimension_slice_scan_by_id(int32 dimension_slice_id)
 		.nkeys = 1,
 		.scankey = scankey,
 		.data = &slice,
+		.limit = 1,
 		.tuple_found = dimension_slice_tuple_found,
 		.lockmode = AccessShareLock,
 		.scandirection = ForwardScanDirection,

src/scanner.c

@@ -175,7 +175,6 @@ scanner_scan(ScannerCtx *ctx)
 
 	while (is_valid)
 	{
-
 		if (ctx->filter == NULL || ctx->filter(&ictx.tinfo, ctx->data))
 		{
 			ictx.tinfo.count++;
@@ -203,6 +202,10 @@ scanner_scan(ScannerCtx *ctx)
 				break;
 		}
 
+		/* Check if limit is reached */
+		if (ctx->limit > 0 && ictx.tinfo.count >= ctx->limit)
+			break;
+
 		is_valid = scanner->getnext(&ictx);
 	}
 

src/scanner.h

@@ -40,7 +40,9 @@ typedef struct ScannerCtx
 	ScannerType scantype;
 	ScanKey		scankey;
 	int			nkeys,
-				norderbys;
+				norderbys,
+				limit;			/* Limit on number of tuples to return. 0 or
+								 * less means no limit */
 	bool		want_itup;
 	LOCKMODE	lockmode;
 	struct