Fix covering plan for a Lucene index over a Synthetic Record Type

docs/ReleaseNotes.md

@@ -15,7 +15,7 @@ The Guava dependency version has been updated to 31.1. Projects may need to chec
 // begin next release
 ### NEXT_RELEASE
 
-* **Bug fix** Fix 1 [(Issue #NNN)](https://github.com/FoundationDB/fdb-record-layer/issues/NNN)
+* **Bug fix** Fix covering plan for a Lucene index over a Synthetic Record Type [(Issue #1927)](https://github.com/FoundationDB/fdb-record-layer/issues/1927)
 * **Bug fix** Fix 2 [(Issue #NNN)](https://github.com/FoundationDB/fdb-record-layer/issues/NNN)
 * **Bug fix** Fix 3 [(Issue #NNN)](https://github.com/FoundationDB/fdb-record-layer/issues/NNN)
 * **Bug fix** Fix 4 [(Issue #NNN)](https://github.com/FoundationDB/fdb-record-layer/issues/NNN)

fdb-record-layer-core/src/main/java/com/apple/foundationdb/record/query/plan/visitor/RecordQueryPlannerSubstitutionVisitor.java

@@ -24,6 +24,9 @@
 import com.apple.foundationdb.record.metadata.Index;
 import com.apple.foundationdb.record.metadata.RecordType;
 import com.apple.foundationdb.record.metadata.expressions.KeyExpression;
+import com.apple.foundationdb.record.metadata.expressions.KeyExpressionWithChildren;
+import com.apple.foundationdb.record.metadata.expressions.ListKeyExpression;
+import com.apple.foundationdb.record.metadata.expressions.ThenKeyExpression;
 import com.apple.foundationdb.record.query.plan.AvailableFields;
 import com.apple.foundationdb.record.query.plan.IndexKeyValueToPartialRecord;
 import com.apple.foundationdb.record.query.plan.PlannableIndexTypes;
@@ -97,10 +100,10 @@ public static RecordQueryPlan removeIndexFetch(@Nonnull RecordMetaData recordMet
             final RecordType recordType = Iterables.getOnlyElement(recordTypes);
             AvailableFields fieldsFromIndex = AvailableFields.fromIndex(recordType, index, indexTypes, commonPrimaryKey, indexPlan);
 
-            Set<KeyExpression> fields = new HashSet<>(requiredFields);
+            final Set<KeyExpression> fields = new HashSet<>(requiredFields);
             if (commonPrimaryKey != null) {
                 // Need the primary key, even if it wasn't one of the explicit result fields.
-                fields.addAll(commonPrimaryKey.normalizeKeyForPositions());
+                flattenKeys(commonPrimaryKey, fields);
             }
             fields.removeIf(keyExpression -> !keyExpression.needsCopyingToPartialRecord());
 
@@ -119,6 +122,17 @@ public static RecordQueryPlan removeIndexFetch(@Nonnull RecordMetaData recordMet
         return null;
     }
 
+    private static void flattenKeys(@Nonnull KeyExpression commonPrimaryKey, @Nonnull Set<KeyExpression> fields) {
+        // Not just normalizeKeyForPositions, because while List doesn't flatten _positions_, that doesn't matter.
+        if (commonPrimaryKey instanceof ThenKeyExpression || commonPrimaryKey instanceof ListKeyExpression) {
+            for (KeyExpression child : ((KeyExpressionWithChildren)commonPrimaryKey).getChildren()) {
+                flattenKeys(child, fields);
+            }
+        } else {
+            fields.addAll(commonPrimaryKey.normalizeKeyForPositions());
+        }
+    }
+
     @Nullable
     public static RecordQueryFetchFromPartialRecordPlan.FetchIndexRecords resolveFetchIndexRecordsFromPlan(@Nonnull final RecordQueryPlan plan) {
         if (plan instanceof RecordQueryPlanWithIndex) {

fdb-record-layer-core/src/test/proto/test_records_join_index.proto

@@ -123,6 +123,7 @@ message OrderWithHeader {
   optional int32 order_no = 3 [(field).primary_key = true];
   optional int32 quantity = 4;
   repeated Ref cc = 5;
+  optional string order_desc = 6;
 }
 
 message RecordTypeUnion {

fdb-record-layer-lucene/src/main/java/com/apple/foundationdb/record/lucene/LuceneIndexQueryPlan.java

@@ -32,6 +32,7 @@
 import com.apple.foundationdb.record.metadata.RecordType;
 import com.apple.foundationdb.record.metadata.expressions.GroupingKeyExpression;
 import com.apple.foundationdb.record.metadata.expressions.KeyExpression;
+import com.apple.foundationdb.record.metadata.expressions.NestingKeyExpression;
 import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreBase;
 import com.apple.foundationdb.record.provider.foundationdb.IndexScanParameters;
 import com.apple.foundationdb.record.query.plan.AvailableFields;
@@ -161,18 +162,8 @@ public void getStoredFields(@Nonnull List<KeyExpression> keyFields, @Nonnull Lis
                                 @Nonnull List<KeyExpression> otherFields) {
         int i = 0;
         while (i < nonStoredFields.size()) {
-            KeyExpression field = nonStoredFields.get(i);
-            KeyExpression origField = field;
-            // Unwrap functions that are just annotations.
-            while (field instanceof LuceneFunctionKeyExpression) {
-                if (field instanceof LuceneFunctionKeyExpression.LuceneSorted) {
-                    field = ((LuceneFunctionKeyExpression.LuceneSorted)field).getSortedExpression();
-                } else if (field instanceof LuceneFunctionKeyExpression.LuceneStored) {
-                    field = ((LuceneFunctionKeyExpression.LuceneStored)field).getStoredExpression();
-                } else {
-                    break;
-                }
-            }
+            KeyExpression origField = nonStoredFields.get(i);
+            KeyExpression field = removeLuceneAnnotations(origField);
             if (field != origField) {
                 nonStoredFields.set(i, field);
                 int j = keyFields.indexOf(origField);
@@ -196,6 +187,29 @@ public void getStoredFields(@Nonnull List<KeyExpression> keyFields, @Nonnull Lis
         }
     }
 
+    // Unwrap functions that are just annotations.
+    @SuppressWarnings("PMD.CompareObjectsWithEquals")
+    private static KeyExpression removeLuceneAnnotations(@Nonnull KeyExpression field) {
+        if (field instanceof NestingKeyExpression) {
+            KeyExpression origChild = ((NestingKeyExpression)field).getChild();
+            KeyExpression child = removeLuceneAnnotations(origChild);
+            if (child == origChild) {
+                return field;
+            }
+            return new NestingKeyExpression(((NestingKeyExpression)field).getParent(), child);
+        }
+        while (field instanceof LuceneFunctionKeyExpression) {
+            if (field instanceof LuceneFunctionKeyExpression.LuceneSorted) {
+                field = ((LuceneFunctionKeyExpression.LuceneSorted)field).getSortedExpression();
+            } else if (field instanceof LuceneFunctionKeyExpression.LuceneStored) {
+                field = ((LuceneFunctionKeyExpression.LuceneStored)field).getStoredExpression();
+            } else {
+                break;
+            }
+        }
+        return field;
+    }
+
     @Override
     public boolean equals(final Object o) {
         if (this == o) {

fdb-record-layer-lucene/src/test/java/com/apple/foundationdb/record/lucene/LuceneSyntheticPlannerTest.java

@@ -0,0 +1,125 @@
+/*
+ * LuceneSyntheticPlannerTest.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2022 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.apple.foundationdb.record.lucene;
+
+import com.apple.foundationdb.record.RecordMetaData;
+import com.apple.foundationdb.record.RecordMetaDataBuilder;
+import com.apple.foundationdb.record.TestRecordsJoinIndexProto;
+import com.apple.foundationdb.record.metadata.Index;
+import com.apple.foundationdb.record.metadata.IndexTypes;
+import com.apple.foundationdb.record.metadata.JoinedRecordTypeBuilder;
+import com.apple.foundationdb.record.metadata.Key;
+import com.apple.foundationdb.record.provider.foundationdb.FDBRecordContext;
+import com.apple.foundationdb.record.provider.foundationdb.FDBRecordStoreTestBase;
+import com.apple.foundationdb.record.query.RecordQuery;
+import com.apple.foundationdb.record.query.expressions.QueryComponent;
+import com.apple.foundationdb.record.query.plan.PlannableIndexTypes;
+import com.apple.foundationdb.record.query.plan.plans.RecordQueryPlan;
+import com.google.common.collect.Sets;
+import org.hamcrest.Matcher;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static com.apple.foundationdb.record.lucene.LucenePlanMatchers.query;
+import static com.apple.foundationdb.record.lucene.LucenePlanMatchers.scanParams;
+import static com.apple.foundationdb.record.metadata.Key.Expressions.concat;
+import static com.apple.foundationdb.record.metadata.Key.Expressions.field;
+import static com.apple.foundationdb.record.metadata.Key.Expressions.function;
+import static com.apple.foundationdb.record.query.plan.match.PlanMatchers.coveringIndexScan;
+import static com.apple.foundationdb.record.query.plan.match.PlanMatchers.indexName;
+import static com.apple.foundationdb.record.query.plan.match.PlanMatchers.indexScan;
+import static com.apple.foundationdb.record.query.plan.match.PlanMatchers.indexScanType;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.allOf;
+import static org.hamcrest.Matchers.hasToString;
+
+/**
+ * Tests about applying the Lucene index when used as part of a SyntheticRecord join.
+ */
+public class LuceneSyntheticPlannerTest extends FDBRecordStoreTestBase {
+
+    protected void openRecordStore(FDBRecordContext context, FDBRecordStoreTestBase.RecordMetaDataHook hook, boolean attemptWholeFilter) {
+        RecordMetaDataBuilder metaDataBuilder = RecordMetaData.newBuilder().setRecords(TestRecordsJoinIndexProto.getDescriptor());
+        hook.apply(metaDataBuilder);
+        recordStore = getStoreBuilder(context, metaDataBuilder.getRecordMetaData()).createOrOpen();
+
+        PlannableIndexTypes indexTypes = new PlannableIndexTypes(
+                    Sets.newHashSet(IndexTypes.VALUE, IndexTypes.VERSION),
+                    Sets.newHashSet(IndexTypes.RANK, IndexTypes.TIME_WINDOW_LEADERBOARD),
+                    Sets.newHashSet(IndexTypes.TEXT),
+                    Sets.newHashSet(LuceneIndexTypes.LUCENE)
+            );
+        planner = new LucenePlanner(recordStore.getRecordMetaData(), recordStore.getRecordStoreState(), indexTypes, recordStore.getTimer());
+        planner.setConfiguration(planner.getConfiguration()
+                .asBuilder()
+                .setPlanOtherAttemptWholeFilter(attemptWholeFilter)
+                .build());
+    }
+
+    /**
+     * Verify that Lucene index on Joined record gives covering scan.
+     */
+    @Test
+    void canPlanQueryAgainstSyntheticLuceneType() {
+        try (FDBRecordContext context = openContext()) {
+            openRecordStore(context, metaDataBuilder -> {
+                metaDataBuilder.getRecordType("CustomerWithHeader")
+                        .setPrimaryKey(Key.Expressions.concat(field("___header").nest("z_key"), field("___header").nest("rec_id")));
+                metaDataBuilder.getRecordType("OrderWithHeader")
+                        .setPrimaryKey(Key.Expressions.concat(field("___header").nest("z_key"), field("___header").nest("rec_id")));
+
+                //set up the joined index
+                final JoinedRecordTypeBuilder joined = metaDataBuilder.addJoinedRecordType("luceneJoinedIdx");
+                joined.addConstituent("order", "OrderWithHeader");
+                joined.addConstituent("cust", "CustomerWithHeader");
+                joined.addJoin("order", field("___header").nest("z_key"),
+                        "cust", field("___header").nest("z_key"));
+                joined.addJoin("order", field("custRef").nest("string_value"),
+                        "cust", field("___header").nest("rec_id"));
+
+                metaDataBuilder.addIndex(joined, new Index("joinNestedConcat", concat(
+                        field("cust").nest(function(LuceneFunctionNames.LUCENE_STORED, field("name"))),
+                        field("order").nest(concat(function(LuceneFunctionNames.LUCENE_STORED, field("order_no")),
+                                function(LuceneFunctionNames.LUCENE_TEXT, field("order_desc"))
+                        ))
+                ), LuceneIndexTypes.LUCENE));
+                metaDataBuilder.addIndex("OrderWithHeader", "order$custRef", concat(field("___header").nest("z_key"), field("custRef").nest("string_value")));
+            }, false);
+
+            String luceneSearch = "order_order_desc: \"twelve pineapple\" and cust_name: \"steve\"";
+            QueryComponent filter = new LuceneQueryComponent(luceneSearch, List.of("order", "cust"));
+            RecordQuery query = RecordQuery.newBuilder()
+                    .setRecordType("luceneJoinedIdx")
+                    .setFilter(filter)
+                    .setRequiredResults(List.of(Key.Expressions.field("order").nest("order_no")))
+                    .build();
+            final RecordQueryPlan plan = planner.plan(query);
+            Matcher<RecordQueryPlan> matcher = coveringIndexScan(indexScan(allOf(
+                    indexScanType(LuceneScanTypes.BY_LUCENE),
+                    indexName("joinNestedConcat"),
+                    scanParams(query(hasToString(luceneSearch)))
+            )));
+            assertThat(plan, matcher);
+        }
+
+    }
+}