feat: Added built-in fields (id and score) to resultset columns

Author: Julien Ruaux

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

@@ -0,0 +1,61 @@
+package com.redis.trino;
+
+import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static io.trino.spi.type.RealType.REAL;
+import static io.trino.spi.type.VarcharType.VARCHAR;
+import static java.util.Arrays.stream;
+import static java.util.function.Function.identity;
+
+import java.util.Map;
+import java.util.Optional;
+
+import com.redis.lettucemod.search.Field;
+
+import io.trino.spi.connector.ColumnMetadata;
+import io.trino.spi.type.Type;
+
+enum RediSearchBuiltinField {
+
+	ID("_id", VARCHAR, Field.Type.TAG), SCORE("_score", REAL, Field.Type.NUMERIC);
+
+	private static final Map<String, RediSearchBuiltinField> COLUMNS_BY_NAME = stream(values())
+			.collect(toImmutableMap(RediSearchBuiltinField::getName, identity()));
+
+	private final String name;
+	private final Type type;
+	private final Field.Type fieldType;
+
+	RediSearchBuiltinField(String name, Type type, Field.Type fieldType) {
+		this.name = name;
+		this.type = type;
+		this.fieldType = fieldType;
+	}
+
+	public static Optional<RediSearchBuiltinField> of(String name) {
+		return Optional.ofNullable(COLUMNS_BY_NAME.get(name));
+	}
+
+	public static boolean isBuiltinColumn(String name) {
+		return COLUMNS_BY_NAME.containsKey(name);
+	}
+
+	public String getName() {
+		return name;
+	}
+
+	public Type getType() {
+		return type;
+	}
+
+	public Field.Type getFieldType() {
+		return fieldType;
+	}
+
+	public ColumnMetadata getMetadata() {
+		return ColumnMetadata.builder().setName(name).setType(type).setHidden(true).build();
+	}
+
+	public RediSearchColumnHandle getColumnHandle() {
+		return new RediSearchColumnHandle(name, type, fieldType, true, false);
+	}
+}

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

@@ -58,8 +58,8 @@ public class RediSearchPageSource implements ConnectorPageSource {
 
 	public RediSearchPageSource(RediSearchSession rediSearchSession, RediSearchTableHandle tableHandle,
 			List<RediSearchColumnHandle> columns) {
-		this.columnNames = columns.stream().map(RediSearchColumnHandle::getName).collect(Collectors.toList());
-		this.columnTypes = columns.stream().map(RediSearchColumnHandle::getType).collect(Collectors.toList());
+		this.columnNames = columns.stream().map(RediSearchColumnHandle::getName).collect(Collectors.toUnmodifiableList());
+		this.columnTypes = columns.stream().map(RediSearchColumnHandle::getType).collect(Collectors.toUnmodifiableList());
 		this.cursor = rediSearchSession.search(tableHandle, columns).iterator();
 		this.currentDoc = null;
 		this.pageBuilder = new PageBuilder(columnTypes);
@@ -100,7 +100,8 @@ public Page getNextPage() {
 			pageBuilder.declarePosition();
 			for (int column = 0; column < columnTypes.size(); column++) {
 				BlockBuilder output = pageBuilder.getBlockBuilder(column);
-				String value = currentDoc.get(columnNames.get(column));
+				String columnName = columnNames.get(column);
+				String value = currentValue(columnName);
 				if (value == null) {
 					output.appendNull();
 				} else {
@@ -114,6 +115,18 @@ public Page getNextPage() {
 		return page;
 	}
 
+	private String currentValue(String columnName) {
+		if (RediSearchBuiltinField.isBuiltinColumn(columnName)) {
+			if (RediSearchBuiltinField.ID.getName().equals(columnName)) {
+				return currentDoc.getId();
+			}
+			if (RediSearchBuiltinField.SCORE.getName().equals(columnName)) {
+				return String.valueOf(currentDoc.getScore());
+			}
+		}
+		return currentDoc.get(columnName);
+	}
+
 	public static JsonGenerator createJsonGenerator(JsonFactory factory, SliceOutput output) throws IOException {
 		return factory.createGenerator((OutputStream) output);
 	}

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

@@ -217,25 +217,28 @@ private RediSearchTable loadTableSchema(SchemaTableName schemaTableName) {
 			throw new TableNotFoundException(schemaTableName, format("Index '%s' not found", index), null);
 		}
 		Set<String> fields = new HashSet<>();
-		ImmutableList.Builder<RediSearchColumnHandle> columnHandles = ImmutableList.builder();
-		for (Field<String> field : indexInfo.get().getFields()) {
-			RediSearchColumnHandle column = buildColumnHandle(field);
+		ImmutableList.Builder<RediSearchColumnHandle> columns = ImmutableList.builder();
+		for (RediSearchBuiltinField builtinfield : RediSearchBuiltinField.values()) {
+			fields.add(builtinfield.getName());
+			columns.add(builtinfield.getColumnHandle());
+		}
+		for (Field<String> indexedField : indexInfo.get().getFields()) {
+			RediSearchColumnHandle column = buildColumnHandle(indexedField);
 			fields.add(column.getName());
-			columnHandles.add(column);
+			columns.add(column);
 		}
 		SearchResults<String, String> results = connection.sync().ftSearch(index, "*");
 		for (Document<String, String> doc : results) {
-			for (String field : doc.keySet()) {
-				if (fields.contains(field)) {
+			for (String docField : doc.keySet()) {
+				if (fields.contains(docField)) {
 					continue;
 				}
-				columnHandles
-						.add(new RediSearchColumnHandle(field, VarcharType.VARCHAR, Field.Type.TEXT, false, false));
-				fields.add(field);
+				columns.add(new RediSearchColumnHandle(docField, VarcharType.VARCHAR, Field.Type.TEXT, false, false));
+				fields.add(docField);
 			}
 		}
 		return new RediSearchTable(new RediSearchTableHandle(RediSearchTableHandle.Type.SEARCH, schemaTableName),
-				columnHandles.build());
+				columns.build());
 	}
 
 	private Optional<IndexInfo> indexInfo(String index) {

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

@@ -226,6 +226,7 @@ public Search search(RediSearchTableHandle tableHandle, List<RediSearchColumnHan
 		String index = index(tableHandle);
 		String query = queryBuilder.buildQuery(tableHandle.getConstraint(), tableHandle.getWildcards());
 		Builder<String, String> options = SearchOptions.builder();
+		options.withScores(true);
 		options.limit(Limit.offset(0).num(limit(tableHandle)));
 		options.returnFields(columns.stream().map(RediSearchColumnHandle::getName).toArray(String[]::new));
 		return Search.builder().index(index).query(query).options(options.build()).build();

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

@@ -109,10 +109,26 @@ protected QueryRunner createQueryRunner() throws Exception {
 
 	@Test
 	public void testNonIndexedFields() throws IOException {
+		try {
+			redisearch.getTestContext().sync().ftDropindexDeleteDocs(Beers.INDEX);
+		} catch (Exception e) {
+			// ignore
+		}
 		Beers.populateIndex(redisearch.getTestContext().getConnection());
 		getQueryRunner().execute("select id, last_mod from beers");
 	}
 
+	@Test
+	public void testBuiltinFields() throws IOException {
+		try {
+			redisearch.getTestContext().sync().ftDropindexDeleteDocs(Beers.INDEX);
+		} catch (Exception e) {
+			// ignore
+		}
+		Beers.populateIndex(redisearch.getTestContext().getConnection());
+		getQueryRunner().execute("select _id, _score from beers");
+	}
+
 	@SuppressWarnings("unchecked")
 	@Test
 	public void testJsonSearch() throws IOException {