test: Increased RediSearch aggregate max results config

Author: Julien Ruaux

pom.xml

@@ -26,9 +26,9 @@
 		<github.owner>redis-field-engineering</github.owner>
 		<github.repo>redis-sql-trino</github.repo>
 
-		<lettucemod.version>3.2.0</lettucemod.version>
+		<lettucemod.version>3.2.1</lettucemod.version>
 		<lettuce.version>6.2.2.RELEASE</lettuce.version>
-		<testcontainers-redis.version>1.6.2</testcontainers-redis.version>
+		<testcontainers-redis.version>1.6.3</testcontainers-redis.version>
 		<ulid.version>5.1.0</ulid.version>
 
 		<air.check.skip-license>true</air.check.skip-license>
@@ -364,12 +364,6 @@
 			</exclusions>
 		</dependency>
 
-		<dependency>
-			<groupId>javax.servlet</groupId>
-			<artifactId>javax.servlet-api</artifactId>
-			<scope>test</scope>
-		</dependency>
-
 		<dependency>
 			<groupId>org.assertj</groupId>
 			<artifactId>assertj-core</artifactId>

src/main/java/com/redis/trino/RediSearchMetadata.java

@@ -246,7 +246,7 @@ public Optional<ConnectorOutputMetadata> finishInsert(ConnectorSession session,
 			Collection<ComputedStatistics> computedStatistics) {
 		return Optional.empty();
 	}
-
+	
 	@Override
 	public RediSearchColumnHandle getDeleteRowIdColumnHandle(ConnectorSession session,
 			ConnectorTableHandle tableHandle) {

src/main/java/com/redis/trino/RediSearchSession.java

@@ -51,8 +51,10 @@
 import com.google.common.util.concurrent.UncheckedExecutionException;
 import com.redis.lettucemod.api.StatefulRedisModulesConnection;
 import com.redis.lettucemod.search.AggregateOperation;
+import com.redis.lettucemod.search.AggregateOptions;
 import com.redis.lettucemod.search.AggregateWithCursorResults;
 import com.redis.lettucemod.search.CreateOptions;
+import com.redis.lettucemod.search.CursorOptions;
 import com.redis.lettucemod.search.Document;
 import com.redis.lettucemod.search.Field;
 import com.redis.lettucemod.search.Group;
@@ -309,10 +311,13 @@ public SearchResults<String, String> search(RediSearchTableHandle tableHandle, S
 	public AggregateWithCursorResults<String> aggregate(RediSearchTableHandle table) {
 		Aggregation aggregation = translator.aggregate(table);
 		log.info("Running %s", aggregation);
-		AggregateWithCursorResults<String> results = connection.sync().ftAggregate(aggregation.getIndex(),
-				aggregation.getQuery(), aggregation.getCursorOptions(), aggregation.getOptions());
-		List<AggregateOperation<?, ?>> groupBys = aggregation.getOptions().getOperations().stream()
-				.filter(o -> o.getType() == AggregateOperation.Type.GROUP).collect(Collectors.toList());
+		String index = aggregation.getIndex();
+		String query = aggregation.getQuery();
+		CursorOptions cursor = aggregation.getCursorOptions();
+		AggregateOptions<String, String> options = aggregation.getOptions();
+		AggregateWithCursorResults<String> results = connection.sync().ftAggregate(index, query, cursor, options);
+		List<AggregateOperation<String, String>> groupBys = aggregation.getOptions().getOperations().stream()
+				.filter(this::isGroupOperation).collect(Collectors.toList());
 		if (results.isEmpty() && !groupBys.isEmpty()) {
 			Group groupBy = (Group) groupBys.get(0);
 			Optional<String> as = groupBy.getReducers()[0].getAs();
@@ -325,6 +330,10 @@ public AggregateWithCursorResults<String> aggregate(RediSearchTableHandle table)
 		return results;
 	}
 
+	private boolean isGroupOperation(AggregateOperation<String, String> operation) {
+		return operation.getType() == AggregateOperation.Type.GROUP;
+	}
+
 	public AggregateWithCursorResults<String> cursorRead(RediSearchTableHandle tableHandle, long cursor) {
 		String index = index(tableHandle.getSchemaTableName());
 		if (config.getCursorCount() > 0) {

src/main/java/com/redis/trino/RediSearchTranslator.java

@@ -222,12 +222,12 @@ public Search build() {
 
 	}
 
-	public Search search(RediSearchTableHandle tableHandle, String[] columnNames) {
-		String index = index(tableHandle);
-		String query = queryBuilder.buildQuery(tableHandle.getConstraint(), tableHandle.getWildcards());
+	public Search search(RediSearchTableHandle table, String[] columnNames) {
+		String index = index(table);
+		String query = queryBuilder.buildQuery(table.getConstraint(), table.getWildcards());
 		Builder<String, String> options = SearchOptions.builder();
 		options.withScores(true);
-		options.limit(Limit.offset(0).num(limit(tableHandle)));
+		options.limit(Limit.offset(0).num(limit(table)));
 		options.returnFields(columnNames);
 		return Search.builder().index(index).query(query).options(options.build()).build();
 	}

src/test/java/com/redis/trino/RediSearchServer.java

@@ -12,7 +12,8 @@ public class RediSearchServer implements Closeable {
 
 	public RediSearchServer() {
 		this.dockerContainer = new RedisStackContainer(
-				RedisStackContainer.DEFAULT_IMAGE_NAME.withTag(RedisStackContainer.DEFAULT_TAG));
+				RedisStackContainer.DEFAULT_IMAGE_NAME.withTag(RedisStackContainer.DEFAULT_TAG))
+				.withEnv("REDISEARCH_ARGS", "MAXAGGREGATERESULTS 1000000");
 		this.dockerContainer.start();
 		this.context = new RedisTestContext(dockerContainer);
 	}

src/test/java/com/redis/trino/TestRediSearchConnectorSmokeTest.java

@@ -39,6 +39,12 @@ public class TestRediSearchConnectorSmokeTest extends BaseConnectorSmokeTest {
 
 	private RediSearchServer redisearch;
 
+	@Override
+	protected QueryRunner createQueryRunner() throws Exception {
+		redisearch = new RediSearchServer();
+		return RediSearchQueryRunner.createRediSearchQueryRunner(redisearch, CUSTOMER, NATION, ORDERS, REGION);
+	}
+
 	private void populateBeers() throws IOException, InterruptedException {
 		Beers.populateIndex(redisearch.getTestContext().getConnection());
 	}
@@ -103,12 +109,6 @@ protected void assertQuery(String sql) {
 		super.assertQuery(sql);
 	}
 
-	@Override
-	protected QueryRunner createQueryRunner() throws Exception {
-		redisearch = new RediSearchServer();
-		return RediSearchQueryRunner.createRediSearchQueryRunner(redisearch, CUSTOMER, NATION, ORDERS, REGION);
-	}
-
 	@Test
 	public void testRediSearchFields() throws IOException, InterruptedException {
 		populateBeers();