Extend use_index hint with table/column-aware format

hibernate-dialect/CHANGELOG.md

@@ -4,6 +4,7 @@
 - Support Hibernate 7
 - Support OFFSET without LIMIT
 - Dropped `current time` in HQL
+- Extended `use_index:` hint with table-and-column-aware format `use_index:<index-name>:<table-name>(<column>[,<column>...])` to pin a secondary index per `FROM`/`JOIN` when the same table is joined multiple times under different aliases
 
 ## 1.6.0 ##
 

hibernate-dialect/README.md

@@ -74,6 +74,56 @@ Use this custom dialect just like any other Hibernate dialect.
 Map your entity classes to database tables and use Hibernate's
 session factory to perform database operations.
 
+## Query hints
+
+The dialect parses the JPA `HibernateHints.HINT_COMMENT` value (or the
+`org.hibernate.query.Query#addQueryHint` value) and rewrites the generated SQL
+before it is sent to YDB. Multiple hints can be combined in one comment by
+separating them with `;`.
+
+| Hint | Effect |
+| --- | --- |
+| `use_index:<index-name>` | Rewrites a simple `select … from <table> … where …` so that the FROM table uses the given secondary index (`VIEW <index-name>`). |
+| `use_index:<index-name>:<table-name>(<column>[,<column>…])` | Table- and column-aware hint. For every `FROM`/`JOIN` occurrence of `<table-name>`, the dialect rewrites the segment to `<table-name> view <index-name> <alias>` **only if every column listed in the hint is referenced by that alias** inside the relevant scope (the `ON` clause for a `JOIN`, the `WHERE` clause for the `FROM` table). When several hints fully match the same table reference the most specific one (largest number of listed columns) wins. |
+| `use_scan` | Wraps the query in `scan`. |
+| `add_pragma:<pragma>` | Prepends `PRAGMA <pragma>;` to the query. |
+
+### Why the table/column form?
+
+YDB does not always pick a secondary index automatically when the same table
+is joined several times under different aliases (e.g. fetching parent and
+child accounts in one query). The auto-picker does not currently consider
+JOIN conditions, so heavy queries can fall back to a full scan. The
+table/column form lets you pin a specific index per JOIN without rewriting
+the JPQL query:
+
+```java
+@QueryHints({
+    @QueryHint(name = HibernateHints.HINT_COMMENT,
+               value = "use_index:bank_account_code_idx:bank_account(code);"
+                     + "use_index:bank_account_parent_idx:bank_account(parent)")
+})
+Optional<Account> findById(Long id);
+```
+
+For a Hibernate-generated query like
+
+```sql
+left join bank_account a1_0 on a1_0.code=source.code
+left join bank_account a3_0 on a3_0.parent=source.parent
+```
+
+the dialect rewrites the joins to
+
+```sql
+left join bank_account view bank_account_code_idx a1_0 on a1_0.code=source.code
+left join bank_account view bank_account_parent_idx a3_0 on a3_0.parent=source.parent
+```
+
+Multiple columns can be listed inside the parentheses (`bank_account(code,parent)`)
+to apply the same index hint to whichever of those columns appears in the
+join condition.
+
 ## Integration with Spring Data JPA
 
 Configure Spring Data JPA with Hibernate to use custom YDB dialect

hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/YdbDialect.java

@@ -4,6 +4,7 @@
 import java.time.Instant;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import org.hibernate.boot.model.FunctionContributions;
@@ -274,17 +275,19 @@ public String addSqlHintOrComment(String sql, QueryOptions queryOptions, boolean
         }
 
         if (queryOptions.getComment() != null) {
-            boolean commentIsHint = false;
-
-            var hints = queryOptions.getComment().split(";");
+            List<String> comments = Arrays.stream(queryOptions.getComment().split(";"))
+                    .map(String::trim)
+                    .filter(comment -> !comment.isEmpty())
+                    .toList();
 
+            boolean commentIsHint = false;
             for (var queryHintHandler : QUERY_HINT_HANDLERS) {
-                for (var hint : hints) {
-                    hint = hint.trim();
-                    if (queryHintHandler.commentIsHint(hint)) {
-                        commentIsHint = true;
-                        sql = queryHintHandler.addQueryHints(sql, List.of(hint));
-                    }
+                List<String> handlerHints = comments.stream()
+                        .filter(queryHintHandler::commentIsHint)
+                        .toList();
+                if (!handlerHints.isEmpty()) {
+                    commentIsHint = true;
+                    sql = queryHintHandler.addQueryHints(sql, handlerHints);
                 }
             }
 

hibernate-dialect/src/main/java/tech/ydb/hibernate/dialect/hint/IndexQueryHintHandler.java

@@ -1,7 +1,10 @@
 package tech.ydb.hibernate.dialect.hint;
 
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -11,11 +14,25 @@
 public class IndexQueryHintHandler implements QueryHintHandler {
     public static final IndexQueryHintHandler INSTANCE = new IndexQueryHintHandler();
 
-    private static final Pattern SELECT_FROM_WHERE_QUERY_PATTERN = Pattern
-            .compile("^\\s*(select.+?from\\s+\\w+)(.+where.+)$", Pattern.CASE_INSENSITIVE);
-
     private static final String HINT_USE_INDEX = "use_index:";
 
+
+    private static final String IDENT = "(?:\\w+|`[^`]+`|\"[^\"]+\")";
+
+
+    private static final Pattern TYPED_HINT_BODY = Pattern.compile("^([^:]+):([^(]+)\\(([^)]+)\\)$");
+
+
+    private static final Pattern TABLE_ALIAS = Pattern.compile("\\s*(" + IDENT + ")\\s+(\\w+)");
+
+
+    private static final Pattern SELECT_FROM_WHERE = Pattern
+            .compile("^\\s*(select.+?from\\s+" + IDENT + ")(.+where.+)$", Pattern.CASE_INSENSITIVE);
+
+
+    private static final Pattern FROM_HAS_VIEW = Pattern
+            .compile("\\bfrom\\s+" + IDENT + "\\s+view\\b", Pattern.CASE_INSENSITIVE);
+
     private IndexQueryHintHandler() {
     }
 
@@ -26,23 +43,244 @@ public boolean commentIsHint(String comment) {
 
     @Override
     public String addQueryHints(String query, List<String> hints) {
-        var useIndexes = new ArrayList<String>();
-        hints.forEach(hint -> {
-            if (hint.startsWith(HINT_USE_INDEX)) {
-                useIndexes.add(hint.substring(HINT_USE_INDEX.length()));
+        ParsedHints parsed = parseHints(hints);
+        if (parsed.empty()) {
+            return query;
+        }
+
+        String result = query;
+        if (!parsed.typed.isEmpty()) {
+            result = applyTypedHints(result, parsed.typed);
+        }
+        if (!parsed.simple.isEmpty() && !FROM_HAS_VIEW.matcher(result).find()) {
+            result = applyShortHints(result, parsed.simple);
+        }
+        return result;
+    }
+
+    private record ParsedHints(List<String> simple, Map<String, List<IndexHint>> typed) {
+        boolean empty() {
+            return simple.isEmpty() && typed.isEmpty();
+        }
+    }
+
+    private record IndexHint(String indexName, List<String> columns) {
+    }
+
+    private static ParsedHints parseHints(List<String> hints) {
+        List<String> simple = new ArrayList<>();
+        Map<String, List<IndexHint>> typed = new LinkedHashMap<>();
+        for (String hint : hints) {
+            if (!hint.startsWith(HINT_USE_INDEX)) {
+                continue;
+            }
+            String body = hint.substring(HINT_USE_INDEX.length()).trim();
+            Matcher m = TYPED_HINT_BODY.matcher(body);
+            if (m.matches()) {
+                String indexName = m.group(1).trim();
+                String tableName = unquote(m.group(2).trim());
+                List<String> columns = splitColumns(m.group(3));
+                if (!indexName.isEmpty() && !tableName.isEmpty() && !columns.isEmpty()) {
+                    typed.computeIfAbsent(tableName, key -> new ArrayList<>())
+                            .add(new IndexHint(indexName, columns));
+                }
+            } else if (!body.isEmpty()) {
+                simple.add(body);
+            }
+        }
+        return new ParsedHints(simple, typed);
+    }
+
+    private static List<String> splitColumns(String raw) {
+        return Arrays.stream(raw.split(","))
+                .map(String::trim)
+                .map(IndexQueryHintHandler::unquote)
+                .filter(s -> !s.isEmpty())
+                .toList();
+    }
+
+    private record TableRef(String table, String alias, int tableEnd, int aliasStart, String columnScope) {
+    }
+
+    private record Replacement(int start, int end, String text) {
+    }
+
+    private static String applyTypedHints(String query, Map<String, List<IndexHint>> typed) {
+        List<TableRef> refs = collectTableRefs(query);
+        List<Replacement> replacements = new ArrayList<>();
+        for (TableRef ref : refs) {
+            List<IndexHint> hintsForTable = typed.get(unquote(ref.table));
+            if (hintsForTable == null) {
+                continue;
+            }
+            IndexHint chosen = pickBestIndex(hintsForTable, ref);
+            if (chosen != null) {
+                replacements.add(new Replacement(ref.tableEnd, ref.aliasStart, " view " + chosen.indexName + " "));
+            }
+        }
+        return applyReplacements(query, replacements);
+    }
+
+
+    private static IndexHint pickBestIndex(List<IndexHint> hints, TableRef ref) {
+        IndexHint best = null;
+        for (IndexHint hint : hints) {
+            if (!allColumnsReferenced(ref.columnScope, ref.alias, hint.columns)) {
+                continue;
+            }
+            if (best == null || hint.columns.size() > best.columns.size()) {
+                best = hint;
+            }
+        }
+        return best;
+    }
+
+    private static boolean allColumnsReferenced(String scope, String alias, List<String> columns) {
+        for (String column : columns) {
+            if (!aliasColumnInScope(scope, alias, column)) {
+                return false;
             }
-        });
+        }
+        return true;
+    }
+
+
+    private static boolean aliasColumnInScope(String scope, String alias, String column) {
+        int from = 0;
+        while (true) {
+            int idx = scope.indexOf(alias, from);
+            if (idx < 0) {
+                return false;
+            }
+            if (isAliasColumnMatchAt(scope, idx, alias, column)) {
+                return true;
+            }
+            from = idx + 1;
+        }
+    }
 
-        if (!useIndexes.isEmpty()) {
-            Matcher matcher = SELECT_FROM_WHERE_QUERY_PATTERN.matcher(query);
-            if (matcher.matches() && matcher.groupCount() > 1) {
-                String startToken = matcher.group(1);
-                String endToken = matcher.group(2);
+    private static boolean isAliasColumnMatchAt(String scope, int aliasIdx, String alias, String column) {
+        if (aliasIdx > 0 && isIdentChar(scope.charAt(aliasIdx - 1))) {
+            return false;
+        }
+        int after = aliasIdx + alias.length();
+        if (after >= scope.length() || scope.charAt(after) != '.') {
+            return false;
+        }
+        int columnPos = after + 1;
+        while (columnPos < scope.length() && Character.isWhitespace(scope.charAt(columnPos))) {
+            columnPos++;
+        }
+        if (columnPos >= scope.length()) {
+            return false;
+        }
+        char first = scope.charAt(columnPos);
+        if (first == '`' || first == '"') {
+            int colStart = columnPos + 1;
+            int colEnd = colStart + column.length();
+            return colEnd < scope.length()
+                    && scope.regionMatches(colStart, column, 0, column.length())
+                    && scope.charAt(colEnd) == first;
+        }
+        int colEnd = columnPos + column.length();
+        if (colEnd > scope.length() || !scope.regionMatches(columnPos, column, 0, column.length())) {
+            return false;
+        }
+        return colEnd == scope.length() || !isIdentChar(scope.charAt(colEnd));
+    }
+
+    private static boolean isIdentChar(char ch) {
+        return ch == '_' || Character.isLetterOrDigit(ch);
+    }
+
+
+    private static List<TableRef> collectTableRefs(String query) {
+        List<SqlTopLevelClauseScanner.Clause> clauses = SqlTopLevelClauseScanner.scan(query);
+        String whereScope = whereClauseScope(query, clauses);
+
+        List<TableRef> refs = new ArrayList<>();
+        for (int i = 0; i < clauses.size(); i++) {
+            SqlTopLevelClauseScanner.Clause clause = clauses.get(i);
+            if (!"from".equals(clause.keyword()) && !"join".equals(clause.keyword())) {
+                continue;
+            }
+            int segmentEnd = nextClauseStart(clauses, i, query.length());
+            TableRef ref = parseTableAliasAt(query, clause, segmentEnd, whereScope);
+            if (ref != null) {
+                refs.add(ref);
+            }
+        }
+        return refs;
+    }
 
-                return startToken + " view " + String.join(", ", useIndexes) + endToken;
+    private static String whereClauseScope(String query, List<SqlTopLevelClauseScanner.Clause> clauses) {
+        for (int i = 0; i < clauses.size(); i++) {
+            if (!"where".equals(clauses.get(i).keyword())) {
+                continue;
             }
+            int start = clauses.get(i).end();
+            int end = nextClauseStart(clauses, i, query.length());
+            return query.substring(start, end);
+        }
+        return "";
+    }
+
+    private static int nextClauseStart(List<SqlTopLevelClauseScanner.Clause> clauses, int currentIndex, int defaultEnd) {
+        return (currentIndex + 1 < clauses.size()) ? clauses.get(currentIndex + 1).start() : defaultEnd;
+    }
+
+    private static TableRef parseTableAliasAt(
+            String query,
+            SqlTopLevelClauseScanner.Clause clause,
+            int segmentEnd,
+            String whereScope
+    ) {
+        Matcher m = TABLE_ALIAS.matcher(query.substring(clause.end(), segmentEnd));
+        if (!m.lookingAt()) {
+            return null;
         }
+        String alias = m.group(2);
+        if ("view".equalsIgnoreCase(alias)) {
+            return null;
+        }
+        int tableEnd = clause.end() + m.end(1);
+        int aliasStart = clause.end() + m.start(2);
+        int aliasEnd = clause.end() + m.end(2);
+        String scope = "join".equals(clause.keyword())
+                ? query.substring(aliasEnd, segmentEnd)
+                : whereScope;
+        return new TableRef(m.group(1), alias, tableEnd, aliasStart, scope);
+    }
 
-        return query;
+    private static String applyShortHints(String query, List<String> indexNames) {
+        Matcher m = SELECT_FROM_WHERE.matcher(query);
+        if (!m.matches() || m.groupCount() < 2) {
+            return query;
+        }
+        return m.group(1) + " view " + String.join(", ", indexNames) + m.group(2);
+    }
+
+    private static String unquote(String identifier) {
+        if (identifier.length() < 2) {
+            return identifier;
+        }
+        char first = identifier.charAt(0);
+        char last = identifier.charAt(identifier.length() - 1);
+        if ((first == '`' && last == '`') || (first == '"' && last == '"')) {
+            return identifier.substring(1, identifier.length() - 1);
+        }
+        return identifier;
+    }
+
+    private static String applyReplacements(String query, List<Replacement> replacements) {
+        if (replacements.isEmpty()) {
+            return query;
+        }
+        replacements.sort((a, b) -> Integer.compare(b.start, a.start));
+        StringBuilder sb = new StringBuilder(query);
+        for (Replacement r : replacements) {
+            sb.replace(r.start, r.end, r.text);
+        }
+        return sb.toString();
     }
 }