Merge pull request #108 from sirixdb/claude/vectorized-source-ref-1125

Add source-document identity to the VectorizedExecutor contract

Author: JohannesLichtenberger

src/main/java/io/brackit/query/compiler/optimizer/SourceRef.java

@@ -0,0 +1,131 @@
+/*
+ * [New BSD License]
+ * Copyright (c) 2011-2012, Brackit Project Team <info@brackit.org>
+ * All rights reserved.
+ */
+package io.brackit.query.compiler.optimizer;
+
+/**
+ * Immutable identity of the document a vectorized scan reads from, lifted from the loop variable's
+ * source expression by {@link io.brackit.query.compiler.optimizer.walker.topdown.VectorizedGroupByDetection}
+ * and carried on the annotated {@code PipeExpr} via {@link VectorizedScanAnnotation#SOURCE_REF}.
+ *
+ * <p>{@link VectorizedScanAnnotation#SOURCE_PATH_PREFIX} tells an executor <em>which path</em> inside
+ * a document a scan walks, but never <em>which document</em> it dereferences. A {@link VectorizedExecutor}
+ * is typically bound to a single physical resource/revision (e.g. SirixDB's
+ * {@code SirixVectorizedExecutor}), so a same-shaped query over a <em>different</em> document would be
+ * answered with the bound resource's data — wrong results. This ref closes that gap: the executor
+ * inspects it in {@link VectorizedExecutor#acceptsSource(SourceRef)} and declines a scan it is not
+ * bound to serve, so the translator falls back to the generic (always-correct) pipeline.
+ *
+ * <p>Three kinds, matching what the detection can prove from the AST:
+ * <ul>
+ * <li>{@link Kind#DOCUMENT} — the scan opens a concrete {@code jn:doc}/{@code jn:open} with literal
+ * database/resource arguments; {@link #databaseName()}, {@link #resourceName()} and {@link #revision()}
+ * are populated ({@code revision == } {@link #LATEST_REVISION} when the call names no explicit revision,
+ * i.e. it opens the most-recent one).</li>
+ * <li>{@link Kind#CONTEXT_ITEM} — the scan ranges over the query's context item (the caller's own bound
+ * read transaction); no database/resource is named.</li>
+ * <li>{@link Kind#UNKNOWN} — the source could not be proven to be a single concrete document (a dynamic
+ * {@code jn:doc}, a collection/multi-revision opener, an unresolved variable, or a non-document
+ * source). A resource-bound executor should fail closed on this.</li>
+ * </ul>
+ */
+public final class SourceRef {
+
+  /** Sentinel {@link #revision()} value: the source names no explicit revision (opens the latest). */
+  public static final int LATEST_REVISION = -1;
+
+  /** What the optimizer could prove about the scan's source document. */
+  public enum Kind {
+    DOCUMENT, CONTEXT_ITEM, UNKNOWN
+  }
+
+  private static final SourceRef CONTEXT_ITEM = new SourceRef(Kind.CONTEXT_ITEM, null, null, LATEST_REVISION);
+  private static final SourceRef UNKNOWN = new SourceRef(Kind.UNKNOWN, null, null, LATEST_REVISION);
+
+  private final Kind kind;
+  private final String databaseName;
+  private final String resourceName;
+  private final int revision;
+
+  private SourceRef(final Kind kind, final String databaseName, final String resourceName, final int revision) {
+    this.kind = kind;
+    this.databaseName = databaseName;
+    this.resourceName = resourceName;
+    this.revision = revision;
+  }
+
+  /**
+   * A concrete document scan.
+   *
+   * @param databaseName the literal database name (must not be {@code null})
+   * @param resourceName the literal resource name (must not be {@code null})
+   * @param revision     the explicit revision, or {@link #LATEST_REVISION} when the source opens the
+   *                     most-recent revision
+   */
+  public static SourceRef document(final String databaseName, final String resourceName, final int revision) {
+    if (databaseName == null || resourceName == null) {
+      throw new IllegalArgumentException("databaseName and resourceName must not be null");
+    }
+    return new SourceRef(Kind.DOCUMENT, databaseName, resourceName, revision);
+  }
+
+  /** The query's context item — the caller's own bound read transaction. */
+  public static SourceRef contextItem() {
+    return CONTEXT_ITEM;
+  }
+
+  /** An unprovable / non-single-document source; resource-bound executors should fail closed. */
+  public static SourceRef unknown() {
+    return UNKNOWN;
+  }
+
+  public Kind kind() {
+    return kind;
+  }
+
+  /** {@code true} iff this ref opens a concrete literal document. */
+  public boolean isDocument() {
+    return kind == Kind.DOCUMENT;
+  }
+
+  /** {@code true} iff this ref is the query's context item. */
+  public boolean isContextItem() {
+    return kind == Kind.CONTEXT_ITEM;
+  }
+
+  /** The literal database name for a {@link Kind#DOCUMENT} ref, else {@code null}. */
+  public String databaseName() {
+    return databaseName;
+  }
+
+  /** The literal resource name for a {@link Kind#DOCUMENT} ref, else {@code null}. */
+  public String resourceName() {
+    return resourceName;
+  }
+
+  /**
+   * The explicit revision of a {@link Kind#DOCUMENT} ref, or {@link #LATEST_REVISION} when it opens the
+   * most-recent revision. Always {@link #LATEST_REVISION} for the other kinds.
+   */
+  public int revision() {
+    return revision;
+  }
+
+  /** {@code true} iff a {@link Kind#DOCUMENT} ref names no explicit revision (opens the latest). */
+  public boolean opensLatestRevision() {
+    return revision == LATEST_REVISION;
+  }
+
+  @Override
+  public String toString() {
+    return switch (kind) {
+      case DOCUMENT -> "SourceRef[doc " + databaseName + "/" + resourceName + (revision == LATEST_REVISION
+          ? ""
+          : "@" + revision) + "]";
+      case CONTEXT_ITEM -> "SourceRef[context-item]";
+      case UNKNOWN -> "SourceRef[unknown]";
+    };
+  }
+}

src/main/java/io/brackit/query/compiler/optimizer/VectorizedExecutor.java

@@ -172,4 +172,31 @@ default Sequence executePredicateAggregate(QueryContext ctx, String[] sourcePath
 
   /** Check if this executor can handle the current query context. */
   boolean canExecute(QueryContext ctx);
+
+  /**
+   * Whether this executor may serve a scan over the given source document.
+   *
+   * <p>{@code sourcePath} tells an executor which path inside a document a scan walks, but never which
+   * document. An executor bound to a single physical resource/revision (e.g. SirixDB's
+   * {@code SirixVectorizedExecutor}) would otherwise answer a same-shaped query over a <em>different</em>
+   * document from its own columns — wrong results. The optimizer therefore lifts the scan's source
+   * identity into a {@link SourceRef} (see {@link VectorizedScanAnnotation#SOURCE_REF}) and asks here,
+   * at TRANSLATE time, before substituting the vectorized expression.
+   *
+   * <p>Returning {@code false} is not an error: the translator simply builds the generic (always-correct)
+   * pipeline instead, so declining only ever costs the fast path. A resource-bound executor should fail
+   * closed — accept {@link SourceRef.Kind#DOCUMENT} refs that match its binding (and the query's
+   * {@link SourceRef.Kind#CONTEXT_ITEM}, the caller's own transaction), and decline everything else,
+   * including {@link SourceRef.Kind#UNKNOWN}.
+   *
+   * <p>The default accepts every source — correct for executors that are not bound to one resource
+   * (e.g. bjq's file-backed {@code ParallelGroupByExec}), so the added contract is opt-in and does not
+   * change their behaviour.
+   *
+   * @param source the scan's source identity; never {@code null} when the optimizer annotated the scan
+   * @return {@code true} to allow vectorized serving of this source, {@code false} to fall back
+   */
+  default boolean acceptsSource(SourceRef source) {
+    return true;
+  }
 }

src/main/java/io/brackit/query/compiler/optimizer/VectorizedScanAnnotation.java

@@ -79,6 +79,19 @@ public final class VectorizedScanAnnotation {
    */
   public static final String SOURCE_PATH_PREFIX = "VECTORIZED_SOURCE_PATH_PREFIX";
 
+  /**
+   * Identity of the document the scan reads from. Value is a {@link SourceRef}. Where
+   * {@link #SOURCE_PATH_PREFIX} captures <em>which path</em> a scan walks, this captures <em>which
+   * document</em> it dereferences — the concrete {@code jn:doc}/{@code jn:open} resource/revision, the
+   * query's context item, or {@link SourceRef.Kind#UNKNOWN} when identity can't be proven.
+   *
+   * <p>Set on every vectorizable {@code PipeExpr} the walker annotates. The translator hands it to
+   * {@link VectorizedExecutor#acceptsSource(SourceRef)} so a resource-bound executor can decline a scan
+   * over a document it is not bound to — falling back to the generic pipeline rather than answering with
+   * the wrong resource's data. Executors that are not resource-bound (the default) accept every source.
+   */
+  public static final String SOURCE_REF = "VECTORIZED_SOURCE_REF";
+
   // ---- Order-by ----
   /** Order field name (String). */
   public static final String ORDER_FIELD = "VECTORIZED_ORDER_FIELD";

src/main/java/io/brackit/query/compiler/optimizer/walker/topdown/VectorizedGroupByDetection.java

@@ -9,12 +9,18 @@
 import io.brackit.query.compiler.AST;
 import io.brackit.query.compiler.XQ;
 import io.brackit.query.compiler.optimizer.PredicateNode;
+import io.brackit.query.compiler.optimizer.SourceRef;
 import io.brackit.query.compiler.optimizer.Stage;
 import io.brackit.query.compiler.optimizer.VectorizedScanAnnotation;
+import io.brackit.query.function.json.JSONFun;
 import io.brackit.query.module.StaticContext;
 
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * Optimizer stage that detects FLWOR patterns eligible for vectorized execution.
@@ -43,26 +49,34 @@
  */
 public final class VectorizedGroupByDetection implements Stage {
 
+  /** Guard against pathological ASTs when resolving a scan source through variable bindings. */
+  private static final int MAX_UNWRAP_STEPS = 64;
+
   @Override
   public AST rewrite(StaticContext sctx, AST ast) {
-    walkAndAnnotate(ast);
+    // Resolve a scan's source document (a `for $u in $doc[]` reaches its `jn:doc(...)` only through the
+    // `let $doc := ...` binding), so collect every visible for/let binding once up front and thread the
+    // map into the per-PipeExpr annotation.
+    final Map<Object, AST> variableBindings = new HashMap<>();
+    collectVariableBindings(ast, variableBindings);
+    walkAndAnnotate(ast, variableBindings);
     return ast;
   }
 
-  private void walkAndAnnotate(AST node) {
+  private void walkAndAnnotate(AST node, Map<Object, AST> variableBindings) {
     if (node == null)
       return;
     if (node.getType() == XQ.PipeExpr) {
-      tryAnnotate(node);
+      tryAnnotate(node, variableBindings);
     }
     for (int i = 0; i < node.getChildCount(); i++) {
-      walkAndAnnotate(node.getChild(i));
+      walkAndAnnotate(node.getChild(i), variableBindings);
     }
   }
 
   // ==================== Main pattern matcher ====================
 
-  private void tryAnnotate(AST pipeExpr) {
+  private void tryAnnotate(AST pipeExpr, Map<Object, AST> variableBindings) {
     if (pipeExpr.getChildCount() < 1)
       return;
     AST chain = pipeExpr.getChild(0);
@@ -237,6 +251,11 @@ private void tryAnnotate(AST pipeExpr) {
       if (sourcePath != null) {
         pipeExpr.setProperty(VectorizedScanAnnotation.SOURCE_PATH_PREFIX, sourcePath);
       }
+      // Document identity of the scan source. Set unconditionally (the translator only consults it once
+      // a vectorized claim exists) so a resource-bound executor can decline a scan over a document it is
+      // not bound to — see VectorizedExecutor#acceptsSource.
+      pipeExpr.setProperty(VectorizedScanAnnotation.SOURCE_REF,
+                           resolveSourceRef(forBind.getChild(1), variableBindings));
     }
 
     // Sorted scan emits FULL RECORDS sorted by ONE direct `$loopVar.field` key — only
@@ -564,6 +583,124 @@ private static String qnmLocalName(final Object value) {
     return null;
   }
 
+  // ==================== Source-document identity extraction ====================
+
+  /**
+   * Collect every {@link XQ#ForBind}/{@link XQ#LetBind} binding in the tree into {@code out}, keyed by
+   * the declared variable's QNm. First (outermost) binding wins on shadowing — a heuristic that only
+   * ever costs precision (a misresolved source yields {@link SourceRef#unknown()}, which fails closed).
+   */
+  private static void collectVariableBindings(final AST node, final Map<Object, AST> out) {
+    if ((node.getType() == XQ.ForBind || node.getType() == XQ.LetBind) && node.getChildCount() >= 2) {
+      final Object varKey = bindingVariableKey(node.getChild(0));
+      if (varKey != null) {
+        out.putIfAbsent(varKey, node.getChild(1));
+      }
+    }
+    for (int i = 0, n = node.getChildCount(); i < n; i++) {
+      collectVariableBindings(node.getChild(i), out);
+    }
+  }
+
+  /**
+   * The variable QNm bound by a {@code For}/{@code LetBind}'s first child (a
+   * {@link XQ#TypedVariableBinding} whose own first child, the {@code Variable}, carries the QNm that a
+   * {@link XQ#VariableRef} later resolves against). Falls back to the node's own value defensively.
+   */
+  private static Object bindingVariableKey(final AST typedVariableBinding) {
+    if (typedVariableBinding.getChildCount() > 0) {
+      return typedVariableBinding.getChild(0).getValue();
+    }
+    return typedVariableBinding.getValue();
+  }
+
+  /**
+   * Resolve a loop variable's source expression down to the document it reads from, following
+   * deref/array/filter layers and variable bindings, and classify it as a {@link SourceRef}. Never
+   * {@code null}: an unresolvable, dynamic, cyclic, collection, or non-document source resolves to
+   * {@link SourceRef#unknown()} so a resource-bound executor fails closed.
+   */
+  private SourceRef resolveSourceRef(final AST binding, final Map<Object, AST> variableBindings) {
+    final Set<Object> resolvingVars = new HashSet<>(4);
+    AST current = binding;
+    for (int step = 0; current != null && step < MAX_UNWRAP_STEPS; step++) {
+      switch (current.getType()) {
+        case XQ.DerefExpr, XQ.ArrayAccess, XQ.FilterExpr -> {
+          if (current.getChildCount() < 1) {
+            return SourceRef.unknown();
+          }
+          current = current.getChild(0);
+        }
+        case XQ.VariableRef -> {
+          final Object varKey = current.getValue();
+          if (varKey == null || !resolvingVars.add(varKey)) {
+            return SourceRef.unknown(); // unresolved or cyclic — cannot prove a single document
+          }
+          final AST resolved = variableBindings.get(varKey);
+          if (resolved == null) {
+            return SourceRef.unknown(); // a for-loop / outer variable, not a document binding
+          }
+          current = resolved;
+        }
+        case XQ.ContextItemExpr -> {
+          return SourceRef.contextItem(); // the caller's own bound read transaction
+        }
+        case XQ.FunctionCall -> {
+          return functionCallSourceRef(current);
+        }
+        default -> {
+          return SourceRef.unknown();
+        }
+      }
+    }
+    return SourceRef.unknown();
+  }
+
+  /**
+   * Classify a {@link XQ#FunctionCall} scan source. A {@code jn:doc}/{@code jn:open} with literal
+   * database and resource arguments (and, if present, a literal integer revision) yields a concrete
+   * {@link SourceRef#document}; a dynamic argument, any other {@code jn:} opener (collection /
+   * multi-revision — it spans more than one resource/revision), or a non-JSON function yields
+   * {@link SourceRef#unknown()}.
+   */
+  private SourceRef functionCallSourceRef(final AST call) {
+    if (!(call.getValue() instanceof QNm qnm) || !JSONFun.JSON_NSURI.equals(qnm.getNamespaceURI())) {
+      return SourceRef.unknown();
+    }
+    final String local = qnm.getLocalName();
+    if (!"doc".equals(local) && !"open".equals(local)) {
+      return SourceRef.unknown();
+    }
+    if (call.getChildCount() < 2) {
+      return SourceRef.unknown();
+    }
+    final String databaseName = stringLiteralValue(call.getChild(0));
+    final String resourceName = stringLiteralValue(call.getChild(1));
+    if (databaseName == null || resourceName == null) {
+      return SourceRef.unknown(); // dynamic (non-literal) database/resource — unprovable
+    }
+    if (call.getChildCount() == 2) {
+      return SourceRef.document(databaseName, resourceName, SourceRef.LATEST_REVISION);
+    }
+    final Integer revision = literalRevision(call.getChild(2));
+    if (revision == null) {
+      return SourceRef.unknown(); // dynamic revision — unprovable
+    }
+    return SourceRef.document(databaseName, resourceName, revision);
+  }
+
+  /** The exact int of a literal integer revision argument; {@code null} for anything non-literal/lossy. */
+  private static Integer literalRevision(final AST node) {
+    if (node == null || node.getType() != XQ.Int) {
+      return null;
+    }
+    final Long lv = exactLongOf(node.getValue());
+    if (lv == null || lv < Integer.MIN_VALUE || lv > Integer.MAX_VALUE) {
+      return null;
+    }
+    return lv.intValue();
+  }
+
   // ==================== Generic predicate-tree extraction ====================
 
   /**