diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/query/FDBInQueryTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/query/FDBInQueryTest.java index 4dccdf09c5..b53bea51ee 100644 --- a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/query/FDBInQueryTest.java +++ b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/provider/foundationdb/query/FDBInQueryTest.java @@ -28,6 +28,7 @@ import com.apple.foundationdb.record.RecordCursorIterator; import com.apple.foundationdb.record.RecordCursorResult; import com.apple.foundationdb.record.TestHelpers; +import com.apple.foundationdb.record.TestRecords1Proto; import com.apple.foundationdb.record.TestRecordsEnumProto; import com.apple.foundationdb.record.TestRecordsWithHeaderProto; import com.apple.foundationdb.record.metadata.Index; @@ -76,6 +77,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -142,6 +144,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.anyOf; import static org.hamcrest.Matchers.containsString; +import static org.hamcrest.Matchers.everyItem; import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.greaterThanOrEqualTo; import static org.hamcrest.Matchers.hasSize; @@ -1015,6 +1018,104 @@ void testConstantObjectInQuerySortedWithExtraUnorderedColumns(boolean reverse) t } } + @DualPlannerTest(planner = DualPlannerTest.Planner.CASCADES) + @ParameterizedTest + @BooleanSource + void testColumnsNeededForUnionNotInProjection(boolean reverse) { + Assumptions.assumeTrue(useCascadesPlanner); + final Index index = new Index("indexForInUnion", concatenateFields("str_value_indexed", "num_value_3_indexed", "num_value_2", "num_value_unique")); + final RecordMetaDataHook hook = complexQuerySetupHook() + .andThen(metaDataBuilder -> metaDataBuilder.addIndex("MySimpleRecord", index)); + final Map> primaryKeysByStrValue = new HashMap<>(); + try (FDBRecordContext context = openContext()) { + openSimpleRecordStore(context, hook); + + // For each str_value and numValue3, save different values for numValue2. That way, the union + // plan will need to contain something other than the sort key to avoid marking records as dupes + long recNo = 0; + for (String strValue : List.of("a", "b", "c")) { + for (int numValue3 = 0; numValue3 < 10; numValue3++) { + for (int numValue2 = 0; numValue2 < 5; numValue2++) { + long thisRecNo = recNo++; + final TestRecords1Proto.MySimpleRecord rec = TestRecords1Proto.MySimpleRecord.newBuilder() + .setNumValue3Indexed(numValue3) + .setNumValue2(((int)thisRecNo * 7 + numValue2) % 5) + .setStrValueIndexed(strValue) + .setRecNo(thisRecNo) + .setNumValueUnique(1000 + (int)thisRecNo) + .build(); + primaryKeysByStrValue.computeIfAbsent(strValue, ignore -> new HashSet<>()) + .add(thisRecNo); + recordStore.saveRecord(rec); + } + } + } + + + commit(context); + } + + final CorrelationIdentifier constantAlias = CorrelationIdentifier.uniqueID(); + final ConstantObjectValue strValueListConstant = ConstantObjectValue.of(constantAlias, "0", new Type.Array(false, Type.primitiveType(Type.TypeCode.STRING, false))); + planner.setConfiguration(planner.getConfiguration().asBuilder().setAttemptFailedInJoinAsUnionMaxSize(10).build()); + + // Equivalent to query like: + // SELECT num_value_3_indexed, num_value_unique, rec_no FROM MySimpleRecord WHERE str_value_indexed IN ?strValueList ORDER BY num_value_3_indexed + // Note that it does *not* project num_value_2. However, num_value_2 is a part of the order of the underlying index, so it is necessary to make sure the merge happens correctly + final RecordQueryPlan plan = planGraph(() -> { + final Quantifier base = FDBSimpleQueryGraphTest.fullTypeScan(recordStore.getRecordMetaData(), "MySimpleRecord"); + final Quantifier select = Quantifier.forEach(Reference.of(GraphExpansion.builder() + .addQuantifier(base) + .addPredicate(FieldValue.ofFieldName(base.getFlowedObjectValue(), "str_value_indexed") + .withComparison(new Comparisons.ValueComparison(Comparisons.Type.IN, strValueListConstant))) + .addResultColumn(FDBSimpleQueryGraphTest.projectColumn(base, "num_value_3_indexed")) + .addResultColumn(FDBSimpleQueryGraphTest.projectColumn(base, "num_value_unique")) + .addResultColumn(FDBSimpleQueryGraphTest.projectColumn(base, "rec_no")) + .build() + .buildSelect())); + + final AliasMap aliasMap = AliasMap.ofAliases(select.getAlias(), Quantifier.current()); + return Reference.of(new LogicalSortExpression(ImmutableList.of(FieldValue.ofFieldName(select.getFlowedObjectValue(), "num_value_3_indexed").rebase(aliasMap)), reverse, select)); + }); + + // This doesn't quite plan correctly, due to: https://github.com/FoundationDB/fdb-record-layer/issues/2728 + assertMatchesExactly(plan, + inUnionOnValuesPlan( + mapPlan( + coveringIndexPlan().where(indexPlanOf(indexPlan() + .where(indexName(index.getName())) + .and(reverse ? isReverse() : isNotReverse()) + .and(scanComparisons(equalities(exactly(anyValueComparison())))) + )) + ) + ).where(inUnionComparisonValues(exactly(fieldValueWithFieldNames("num_value_3_indexed")))) // This comparison key is insufficient! + ); + + assertEquals(reverse ? 398379949 : 398350158, plan.planHash(CURRENT_LEGACY)); + assertEquals(reverse ? 1341745787 : 1347286913, plan.planHash(CURRENT_FOR_CONTINUATION)); + + try (FDBRecordContext context = openContext()) { + openSimpleRecordStore(context, hook); + + List> results = queryAsMaps(plan, constantBindings(strValueListConstant, List.of("a", "c"))); + final Set foundPrimaryKeys = new HashSet<>(); + int lastNumValue3 = reverse ? Integer.MAX_VALUE : Integer.MIN_VALUE; + for (Map res : results) { + int numValue3 = (Integer) res.get("num_value_3_indexed"); + assertThat(numValue3, reverse ? lessThanOrEqualTo(lastNumValue3) : greaterThanOrEqualTo(lastNumValue3)); + lastNumValue3 = numValue3; + foundPrimaryKeys.add((Long) res.get("rec_no")); + } + + // Once planner is fixed, should be able to replace everyItem(in( ... )) with containsInAnyOrder(...), + // and assertDiscardedAtMost(50, ... ) should be replaced with assertDiscardedNone() + assertThat(foundPrimaryKeys, everyItem(in(Stream.concat(primaryKeysByStrValue.get("a").stream(), primaryKeysByStrValue.get("c").stream()).toArray()))); + TestHelpers.assertDiscardedAtMost(50, context); + + commit(context); + } + } + enum InAsOrUnionMode { NONE, AS_OR,