Full Text: Queries Need to Sometimes Add Type Filters
If a full text query plan uses a multi-type or universal index and the query itself should only return records from a subset of the types (e.g., a single-type query), then a type filter needs to be added.
diff --git a/fdb-record-layer-core/src/com/apple/foundationdb/record/query/plan/RecordQueryPlanner.java b/fdb-record-layer-core/src/com/apple/foundationdb/record/query/plan/RecordQueryPlanner.java
index 080fa450..9bd6a878 100644
--- a/fdb-record-layer-core/src/com/apple/foundationdb/record/query/plan/RecordQueryPlanner.java
+++ b/fdb-record-layer-core/src/com/apple/foundationdb/record/query/plan/RecordQueryPlanner.java
@@ -940,6 +940,28 @@ public class RecordQueryPlanner implements QueryPlanner {
         }
     }
 
+    @Nonnull
+    private Set<String> getPossibleTypes(@Nonnull Index index) {
+        final Collection<RecordType> recordTypes = metaData.recordTypesForIndex(index);
+        if (recordTypes.size() == 1) {
+            final RecordType singleRecordType = recordTypes.iterator().next();
+            return Collections.singleton(singleRecordType.getName());
+        } else {
+            return recordTypes.stream().map(RecordType::getName).collect(Collectors.toSet());
+        }
+    }
+
+    @Nonnull
+    private RecordQueryPlan addTypeFilterIfNeeded(@Nonnull CandidateScan candidateScan, @Nonnull RecordQueryPlan plan,
+                                                  @Nonnull Set<String> possibleTypes) {
+        Collection<String> allowedTypes = candidateScan.planContext.query.getRecordTypes();
+        if (!allowedTypes.isEmpty() && !allowedTypes.containsAll(possibleTypes)) {
+            return new RecordQueryTypeFilterPlan(plan, allowedTypes);
+        } else {
+            return plan;
+        }
+    }
+
     @Nullable
     private ScoredPlan planVersion(@Nonnull CandidateScan candidateScan,
                                    @Nonnull KeyExpression index,
@@ -1038,7 +1060,10 @@ public class RecordQueryPlanner implements QueryPlanner {
             return null;
         }
         // TODO: Check the rest of the fields of the text index expression to see if the sort and unsatisfied filters can be helped.
-        final RecordQueryTextIndexPlan plan = new RecordQueryTextIndexPlan(index.getName(), scan, candidateScan.reverse);
+        RecordQueryPlan plan = new RecordQueryTextIndexPlan(index.getName(), scan, candidateScan.reverse);
+        // Add a type filter if the index is over more types than those the query specifies
+        Set<String> possibleTypes = getPossibleTypes(index);
+        plan = addTypeFilterIfNeeded(candidateScan, plan, possibleTypes);
         // This weight is fairly arbitrary, but it is supposed to be higher than for most indexes because
         // most of the time, the full text scan is believed to be more selective (and expensive to run as a post-filter)
         // than other indexes.
@@ -1063,7 +1088,6 @@ public class RecordQueryPlanner implements QueryPlanner {
         if (scanComparisons == null) {
             scanComparisons = ScanComparisons.EMPTY;
         }
-
         RecordQueryPlan plan;
         Set<String> possibleTypes;
         if (candidateScan.index == null) {
@@ -1078,14 +1102,10 @@ public class RecordQueryPlanner implements QueryPlanner {
                 scanType = IndexScanType.BY_VALUE;
             }
             plan = new RecordQueryIndexPlan(candidateScan.index.getName(), scanType, scanComparisons, candidateScan.reverse);
-            possibleTypes = metaData.recordTypesForIndex(candidateScan.index).stream().map(RecordType::getName).collect(Collectors.toSet());
+            possibleTypes = getPossibleTypes(candidateScan.index);
         }
-
-        Collection<String> allowedTypes = candidateScan.planContext.query.getRecordTypes();
-        if (allowedTypes != null && allowedTypes.size() > 0 && !allowedTypes.containsAll(possibleTypes)) {
-            plan = new RecordQueryTypeFilterPlan(plan, allowedTypes);
-        }
-
+        // Add a type filter if the query plan might return records of more types than the query specified
+        plan = addTypeFilterIfNeeded(candidateScan, plan, possibleTypes);
         return plan;
     }
 
diff --git a/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/provider/foundationdb/indexes/TextIndexTest.java b/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/provider/foundationdb/indexes/TextIndexTest.java
index 474bad75..21b90813 100644
--- a/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/provider/foundationdb/indexes/TextIndexTest.java
+++ b/fdb-record-layer-core/test/unit/com/apple/foundationdb/record/provider/foundationdb/indexes/TextIndexTest.java
@@ -189,6 +189,7 @@ public class TextIndexTest extends FDBRecordStoreTestBase {
     private static final Index COMBINED_TEXT_BY_GROUP = new Index("Combined$text_by_group", field("text").groupBy(field("group")), IndexTypes.TEXT);
     private static final Index COMPLEX_MULTI_TAG_INDEX = new Index("Complex$multi_tag", field("text").groupBy(field("tag", FanType.FanOut)), IndexTypes.TEXT);
     private static final Index COMPLEX_THEN_TAG_INDEX = new Index("Complex$text_tag", concat(field("text"), field("tag", FanType.FanOut)), IndexTypes.TEXT);
+    private static final Index MULTI_TYPE_INDEX = new Index("Simple&Complex$text", field("text"), IndexTypes.TEXT);
     private static final Index MAP_ON_VALUE_INDEX = new Index("Map$entry-value", new GroupingKeyExpression(field("entry", FanType.FanOut).nest(concatenateFields("key", "value")), 1), IndexTypes.TEXT);
     private static final Index MAP_ON_VALUE_PREFIX_LEGACY = new Index("Map$entry-value_prefix", new GroupingKeyExpression(field("entry", FanType.FanOut).nest(concatenateFields("key", "value")), 1), IndexTypes.TEXT,
             ImmutableMap.of(Index.TEXT_TOKENIZER_NAME_OPTION, PrefixTextTokenizer.NAME, Index.TEXT_TOKENIZER_VERSION_OPTION, "0"));
@@ -1662,6 +1663,86 @@ public class TextIndexTest extends FDBRecordStoreTestBase {
         }
     }
 
+    @Nonnull
+    private List<Long> queryMultiTypeDocuments(@Nonnull QueryComponent textFilter, @Nonnull List<String> recordTypes, int planHash) throws InterruptedException, ExecutionException {
+        if (!(textFilter instanceof ComponentWithComparison)) {
+            throw new RecordCoreArgumentException("filter without comparison provided as text filter");
+        }
+        Matcher<RecordQueryPlan> planMatcher = descendant(textIndexScan(allOf(
+                indexName(MULTI_TYPE_INDEX.getName()),
+                textComparison(equalTo(((ComponentWithComparison)textFilter).getComparison())
+        ))));
+        if (recordTypes.size() != 2) {
+            planMatcher = descendant(typeFilter(contains(recordTypes.get(0)), planMatcher));
+        }
+        return queryDocuments(recordTypes, Collections.singletonList(field("doc_id")), textFilter, planHash, planMatcher)
+                .map(t -> t.getLong(0))
+                .asList()
+                .join();
+    }
+
+    @Test
+    public void queryMultiTypeDocuments() throws Exception {
+        final List<String> bothTypes = Arrays.asList(SIMPLE_DOC, COMPLEX_DOC);
+        final List<String> simpleTypes = Collections.singletonList(SIMPLE_DOC);
+        final List<String> complexTypes = Collections.singletonList(COMPLEX_DOC);
+
+        final List<String> textSamples = Arrays.asList(
+                TextSamples.ROMEO_AND_JULIET_PROLOGUE,
+                TextSamples.ROMEO_AND_JULIET_PROLOGUE,
+                TextSamples.ANGSTROM,
+                TextSamples.AETHELRED,
+                TextSamples.FRENCH,
+                TextSamples.GERMAN
+        );
+        final List<Message> documents = IntStream.range(0, textSamples.size())
+                .mapToObj(i -> {
+                    final String text = textSamples.get(i);
+                    if (i % 2 == 0) {
+                        return SimpleDocument.newBuilder()
+                                .setDocId(i)
+                                .setText(text)
+                                .setGroup(i % 4)
+                                .build();
+                    } else {
+                        return ComplexDocument.newBuilder()
+                                .setDocId(i)
+                                .setText(text)
+                                .setGroup(i % 4)
+                                .build();
+                    }
+                })
+                .collect(Collectors.toList());
+
+        try (FDBRecordContext context = openContext()) {
+            openRecordStore(context, metaDataBuilder -> {
+                metaDataBuilder.getRecordType(COMPLEX_DOC).setPrimaryKey(field("doc_id"));
+                metaDataBuilder.removeIndex(SIMPLE_DEFAULT_NAME);
+                metaDataBuilder.addMultiTypeIndex(Arrays.asList(
+                        metaDataBuilder.getRecordType(SIMPLE_DOC),
+                        metaDataBuilder.getRecordType(COMPLEX_DOC)
+                ), MULTI_TYPE_INDEX);
+            });
+            documents.forEach(recordStore::saveRecord);
+
+            assertEquals(Arrays.asList(0L, 1L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPhrase("where we lay our scene"), bothTypes, 1755757799));
+            assertEquals(Collections.singletonList(0L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPhrase("where we lay our scene"), simpleTypes, -1489953292));
+            assertEquals(Collections.singletonList(1L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPhrase("where we lay our scene"), complexTypes, -1333764430));
+
+            assertEquals(Arrays.asList(2L, 4L, 5L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPrefix("na"), bothTypes, -714642562));
+            assertEquals(Arrays.asList(2L, 4L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPrefix("na"), simpleTypes, 334613643));
+            assertEquals(Collections.singletonList(5L),
+                    queryMultiTypeDocuments(Query.field("text").text().containsPrefix("na"), complexTypes, 490802505));
+
+            commit(context);
+        }
+    }
+
     @Nonnull
     private List<Long> queryMapDocumentsWithScan(@Nonnull QueryComponent filter, int planHash) throws InterruptedException, ExecutionException {
         return queryDocuments(Collections.singletonList(MAP_DOC), Collections.singletonList(field("doc_id")), filter, planHash,
['fdb-record-layer-core/src/com/apple/foundationdb/record/query/plan/RecordQueryPlanner.java', 'fdb-record-layer-core/test/unit/com/apple/foundationdb/record/provider/foundationdb/indexes/TextIndexTest.java']