Full-text search: a filter on another property is applied after the limit (openCypher) or loads nearly every match when it is restrictive (SQL) #11069

/*
 * Copyright © 2021-present Arcade Data Ltd (info@arcadedata.com)
 *
 * 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.
 *
 * SPDX-FileCopyrightText: 2021-present Arcade Data Ltd (info@arcadedata.com)
 * SPDX-License-Identifier: Apache-2.0
 */
package com.arcadedb.function.sql.fulltext;

import com.arcadedb.TestHelper;
import com.arcadedb.database.DatabaseInternal;
import com.arcadedb.query.sql.executor.ResultSet;
import org.junit.jupiter.api.Test;

import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

/**
 * Reproduction: a full-text search combined with a filter on another property.
 * <p>
 * Corpus: {@value #DOCUMENTS} vertices that all match the query {@code common}. Tenant {@code a} owns most of them and scores
 * high (short text, repeated term). Tenant {@code b} owns {@value #TENANT_B_DOCUMENTS} of them and they score LOWEST (one
 * occurrence in a long text), which is the worst case for a filter that rejects the best matches first.
 * <p>
 * Tests marked "DESIRED" fail today and describe the behaviour the issue asks for. The others pass today and pin the current
 * behaviour the issue describes.
 */
class FilteredFullTextSearchReproTest extends TestHelper {
  private static final int    DOCUMENTS           = 2_000;
  private static final int    TENANT_B_DOCUMENTS  = 20;
  private static final int    WANTED              = 10;
  /** Loads tolerated beyond WANTED: none are needed when the filter is decided before the load. */
  private static final int    SLACK               = 5;
  private static final String SEARCH              = "SEARCH_INDEX('Doc[content]', 'common') = true";

  @Test
  void sqlPermissiveFilterIsAlreadyCheap() {
    createCorpus();
    // Control: the best matches pass the filter, so the lazy load stops after about WANTED records. Not the problem.
    final long loaded = recordsLoadedBy("SELECT FROM Doc WHERE " + SEARCH + " AND tenant = 'a' LIMIT " + WANTED);
    assertThat(loaded).as("records loaded when the best matches pass the filter").isLessThanOrEqualTo(WANTED + SLACK);
  }

  @Test
  void sqlRestrictiveFilterLoadsNearlyEveryMatch() {
    createCorpus();
    final String query = "SELECT FROM Doc WHERE " + SEARCH + " AND tenant = 'b' LIMIT " + WANTED;

    final long loaded = recordsLoadedBy(query);

    // DESIRED: the filter is decided before any record is loaded, so every rejected match is skipped without a load and the
    // count stays close to WANTED. Today the lazy loop loads records until WANTED of them pass the filter, and the rejected
    // ones are loaded too (here about DOCUMENTS - TENANT_B_DOCUMENTS of them): the planner returns a row limit of -1
    // (SelectExecutionPlanner.indexedFunctionRowLimit) and the filter is checked on every loaded record.
    assertThat(loaded).as("records loaded for LIMIT " + WANTED + " of " + DOCUMENTS + " matches, " + TENANT_B_DOCUMENTS
        + " of them pass the filter; plan: " + plan(query).replace('\n', ' ')).isLessThanOrEqualTo(WANTED + SLACK);
  }

  @Test
  void cypherFilterAfterTheProcedureLimitReturnsTooFewRows() {
    createCorpus();
    // CURRENT BEHAVIOUR: the procedure takes its own top-WANTED, then WHERE runs on those rows only.
    final int rows = count("opencypher",
        "CALL db.index.fulltext.queryNodes('Doc[content]', 'common', {limit: " + WANTED + "}) YIELD node "
            + "WHERE node.tenant = 'b' RETURN node", Map.of());
    assertThat(rows).as("rows left after filtering the procedure's top " + WANTED + ", although " + TENANT_B_DOCUMENTS
        + " tenant-b documents match").isEqualTo(0);
  }

  @Test
  void cypherWhereOptionPreFiltersLikeTheVectorProcedure() {
    createCorpus();
    // DESIRED: same shape as db.index.vector.queryNodes(index, k, vector, {where: {...}}) (#10808)
    final int rows = count("opencypher",
        "CALL db.index.fulltext.queryNodes('Doc[content]', 'common', {limit: " + WANTED + ", where: {tenant: 'b'}}) YIELD node "
            + "RETURN node", Map.of());
    assertThat(rows).isEqualTo(WANTED);
  }

  private void createCorpus() {
    database.transaction(() -> {
      database.command("sql", "CREATE VERTEX TYPE Doc BUCKETS 1");
      database.command("sql", "CREATE PROPERTY Doc.content STRING");
      database.command("sql", "CREATE PROPERTY Doc.tenant STRING");
      database.command("sql", "CREATE INDEX ON Doc (content) FULL_TEXT METADATA {\"similarity\": \"BM25\"}");
      for (int i = 0; i < DOCUMENTS; i++) {
        final boolean tenantB = i % (DOCUMENTS / TENANT_B_DOCUMENTS) == 0;
        database.newVertex("Doc").set("tenant", tenantB ? "b" : "a")
            // tenant b: one occurrence in a long text (lowest BM25); tenant a: repeated term in a short text (highest)
            .set("content", tenantB ? "common " + "filler ".repeat(30) : "common common common common").save();
      }
    });
    // The first search computes the index's corpus statistics with a scan of the type, which counts as reads: not the subject
    count("sql", "SELECT FROM Doc WHERE " + SEARCH + " LIMIT 1", Map.of());
  }

  private long recordsLoadedBy(final String query) {
    final long before = readRecords();
    try (final ResultSet rs = database.query("sql", query)) {
      int rows = 0;
      while (rs.hasNext()) {
        rs.next();
        ++rows;
      }
      assertThat(rows).as(query).isEqualTo(WANTED);
    }
    return readRecords() - before;
  }

  private long readRecords() {
    return ((Number) ((DatabaseInternal) database).getStats().get("readRecord")).longValue();
  }

  private String plan(final String query) {
    try (final ResultSet rs = database.query("sql", query)) {
      return rs.getExecutionPlan().orElseThrow().prettyPrint(0, 2);
    }
  }

  private int count(final String language, final String query, final Map<String, Object> params) {
    int rows = 0;
    try (final ResultSet rs = database.query(language, query, params)) {
      while (rs.hasNext()) {
        rs.next();
        ++rows;
      }
    }
    return rows;
  }
}