@dxos/echo-pipeline 0.8.3 → 0.8.4-main.05e74ebcff

package/src/query/query-planner.ts

@@ -2,13 +2,31 @@
 // Copyright 2025 DXOS.org
 //
 
-import { type QueryAST } from '@dxos/echo-protocol';
+import { Order, Query } from '@dxos/echo';
+import { QueryAST } from '@dxos/echo-protocol';
 import { invariant } from '@dxos/invariant';
-import type { DXN, SpaceId } from '@dxos/keys';
+import type { DXN } from '@dxos/keys';
 
 import { QueryError } from './errors';
 import { QueryPlan } from './plan';
 
+/**
+ * Creates a QueryError with "Query too complex" message and includes the prettified query in the context.
+ */
+const queryTooComplexError = (query: QueryAST.Query | null): QueryError => {
+  if (query === null) {
+    return new QueryError({
+      message: 'Query too complex',
+      context: { query: null },
+    });
+  }
+  const prettyQuery = Query.pretty(Query.fromAst(query));
+  return new QueryError({
+    message: `Query too complex: ${prettyQuery}`,
+    context: { query, prettyQuery },
+  });
+};
+
 export type QueryPlannerOptions = {
   defaultTextSearchKind: QueryPlan.TextSearchKind;
 };
@@ -32,9 +50,12 @@ export class QueryPlanner {
   }
 
   createPlan(query: QueryAST.Query): QueryPlan.Plan {
+    this._validateQueryScoped(query);
     let plan = this._generate(query, { ...DEFAULT_CONTEXT, originalQuery: query });
     plan = this._optimizeEmptyFilters(plan);
     plan = this._optimizeSoloUnions(plan);
+    plan = this._ensureOrderStep(plan);
+    plan = this._optimizeLimits(plan);
     return plan;
   }
 
@@ -54,12 +75,21 @@ export class QueryPlanner {
         return this._generateRelationTraversalClause(query, context);
       case 'reference-traversal':
         return this._generateReferenceTraversalClause(query, context);
+      case 'hierarchy-traversal':
+        return this._generateHierarchyTraversalClause(query, context);
       case 'union':
         return this._generateUnionClause(query, context);
       case 'set-difference':
         return this._generateSetDifferenceClause(query, context);
+      case 'order':
+        return this._generateOrderClause(query, context);
+      case 'limit':
+        return this._generateLimitClause(query, context);
+      case 'from':
+        return this._generateFromClause(query, context);
       default:
-        throw new QueryError(`Unsupported query type: ${(query as any).type}`, {
+        throw new QueryError({
+          message: `Unsupported query type: ${(query as any).type}`,
           context: { query: context.originalQuery },
         });
     }
@@ -69,15 +99,25 @@ export class QueryPlanner {
     const newContext = {
       ...context,
     };
-    if (query.options.spaceIds) {
-      newContext.selectionSpaces = query.options.spaceIds as readonly SpaceId[];
-    }
     if (query.options.deleted) {
       newContext.deletedHandling = query.options.deleted;
     }
     return this._generate(query.query, newContext);
   }
 
+  private _generateFromClause(query: QueryAST.QueryFromClause, context: GenerationContext): QueryPlan.Plan {
+    if (query.from._tag === 'scope') {
+      return this._generate(query.query, { ...context, scope: query.from.scope });
+    }
+
+    // Subquery from: flatten by rewriting the outer query's leaf select into a filter on the subquery.
+    const subquery = query.from.query;
+    const rewritten = QueryAST.map(query.query, (node) =>
+      node.type === 'select' ? { type: 'filter', selection: subquery, filter: node.filter } : node,
+    );
+    return this._generate(rewritten, context);
+  }
+
   private _generateSelectClause(query: QueryAST.QuerySelectClause, context: GenerationContext): QueryPlan.Plan {
     return this._generateSelectionFromFilter(query.filter, context);
   }
@@ -86,6 +126,7 @@ export class QueryPlanner {
   // TODO(dmaretskyi): If the tip of the query ast is a [select, ...filter] shape we can reorder the filters so the query is most efficient.
   private _generateSelectionFromFilter(filter: QueryAST.Filter, context: GenerationContext): QueryPlan.Plan {
     switch (filter.type) {
+      // Props
       case 'object': {
         if (
           context.selectionInverted &&
@@ -102,7 +143,7 @@ export class QueryPlanner {
           ]);
         }
         if (context.selectionInverted) {
-          throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+          throw queryTooComplexError(context.originalQuery);
         }
 
         // Try to utilize indexes during selection, prioritizing selecting by id, then by typename.
@@ -111,7 +152,7 @@ export class QueryPlanner {
           return QueryPlan.Plan.make([
             {
               _tag: 'SelectStep',
-              spaces: context.selectionSpaces,
+              scope: context.scope,
               selector: {
                 _tag: 'IdSelector',
                 objectIds: filter.id,
@@ -127,7 +168,7 @@ export class QueryPlanner {
           return QueryPlan.Plan.make([
             {
               _tag: 'SelectStep',
-              spaces: context.selectionSpaces,
+              scope: context.scope,
               selector: {
                 _tag: 'TypeSelector',
                 typename: [filter.typename as DXN.String],
@@ -135,16 +176,17 @@ export class QueryPlanner {
               },
             },
             ...this._generateDeletedHandlingSteps(context),
+            // TODO(dmaretskyi): Normally we could skip filtering by typename here, but since the index does not separate schema versions, we need to do additional filter to only select the correct version.
             {
               _tag: 'FilterStep',
-              filter: { ...filter, typename: null },
+              filter: { ...filter },
             },
           ]);
         } else {
           return QueryPlan.Plan.make([
             {
               _tag: 'SelectStep',
-              spaces: context.selectionSpaces,
+              scope: context.scope,
               selector: {
                 _tag: 'WildcardSelector',
               },
@@ -157,44 +199,223 @@ export class QueryPlanner {
           ]);
         }
       }
+
+      // Tag
+      case 'tag': {
+        return QueryPlan.Plan.make([
+          {
+            _tag: 'SelectStep',
+            scope: context.scope,
+            selector: {
+              _tag: 'WildcardSelector',
+            },
+          },
+          ...this._generateDeletedHandlingSteps(context),
+          {
+            _tag: 'FilterStep',
+            filter: { ...filter },
+          },
+        ]);
+      }
+
+      // Text
       case 'text-search': {
         return QueryPlan.Plan.make([
           {
             _tag: 'SelectStep',
-            spaces: context.selectionSpaces,
+            scope: context.scope,
             selector: {
               _tag: 'TextSelector',
               text: filter.text,
               searchKind: filter.searchKind ?? this._options.defaultTextSearchKind,
+              typename: null,
             },
           },
           ...this._generateDeletedHandlingSteps(context),
         ]);
       }
+
+      // Timestamp
+      case 'timestamp': {
+        if (context.selectionInverted) {
+          throw new QueryError({
+            message:
+              'Negated timestamp filters are not supported. Use Filter.updated/Filter.created with the inverted range instead.',
+            context: { query: context.originalQuery },
+          });
+        }
+        return QueryPlan.Plan.make([
+          {
+            _tag: 'SelectStep',
+            scope: context.scope,
+            selector: _timestampFilterToSelector(filter),
+          },
+          ...this._generateDeletedHandlingSteps(context),
+        ]);
+      }
+
+      // ChildOf
+      case 'child-of': {
+        return QueryPlan.Plan.make([
+          {
+            _tag: 'SelectStep',
+            scope: context.scope,
+            selector: {
+              _tag: 'WildcardSelector',
+            },
+          },
+          ...this._generateDeletedHandlingSteps(context),
+          {
+            _tag: 'FilterStep',
+            filter: { ...filter },
+          },
+        ]);
+      }
+
+      // Compare
       case 'compare':
-        throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+        throw queryTooComplexError(context.originalQuery);
       case 'in':
-        throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+        throw queryTooComplexError(context.originalQuery);
       case 'range':
-        throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+        throw queryTooComplexError(context.originalQuery);
+
+      // Boolean
       case 'not':
         return this._generateSelectionFromFilter(filter.filter, {
           ...context,
           selectionInverted: !context.selectionInverted,
         });
-      case 'and':
-        throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+      case 'and': {
+        const flatFilters = _flattenAnd(filter.filters);
+        const timestampFilters = flatFilters.filter((f): f is QueryAST.FilterTimestamp => f.type === 'timestamp');
+        const childOfFilters = flatFilters.filter((f): f is QueryAST.FilterChildOf => f.type === 'child-of');
+        const otherFilters = flatFilters.filter((f) => f.type !== 'timestamp' && f.type !== 'child-of');
+
+        if (timestampFilters.length > 0 && context.selectionInverted) {
+          throw new QueryError({
+            message:
+              'Negated timestamp filters are not supported. Use Filter.updated/Filter.created with the inverted range instead.',
+            context: { query: context.originalQuery },
+          });
+        }
+
+        if (timestampFilters.length > 0 && otherFilters.length <= 1 && childOfFilters.length === 0) {
+          const innerFilter = otherFilters[0];
+          const innerPlan = innerFilter
+            ? this._generateSelectionFromFilter(innerFilter, context)
+            : QueryPlan.Plan.make([
+                {
+                  _tag: 'SelectStep',
+                  scope: context.scope,
+                  selector: { _tag: 'WildcardSelector' },
+                },
+                ...this._generateDeletedHandlingSteps(context),
+              ]);
+
+          const selectIdx = innerPlan.steps.findIndex((s) => s._tag === 'SelectStep');
+          if (selectIdx !== -1) {
+            const newSteps = [...innerPlan.steps];
+            newSteps.splice(selectIdx + 1, 0, {
+              _tag: 'FilterStep',
+              filter: { type: 'and', filters: timestampFilters },
+            });
+            return QueryPlan.Plan.make(newSteps);
+          }
+
+          const selector = _mergeTimestampSelectors(timestampFilters.map(_timestampFilterToSelector));
+          return QueryPlan.Plan.make([
+            {
+              _tag: 'SelectStep',
+              scope: context.scope,
+              selector,
+            },
+            ...this._generateDeletedHandlingSteps(context),
+          ]);
+        }
+
+        if (timestampFilters.length > 0 && childOfFilters.length === 0) {
+          throw new QueryError({
+            message:
+              'Timestamp filters can only be combined with a single type or property filter via AND. Split complex filters into a subquery.',
+            context: { query: context.originalQuery },
+          });
+        }
+
+        if (childOfFilters.length > 0) {
+          const remainingFilters = flatFilters.filter((f) => f.type !== 'child-of');
+          const innerPlan =
+            remainingFilters.length === 1
+              ? this._generateSelectionFromFilter(remainingFilters[0], context)
+              : remainingFilters.length > 1
+                ? this._generateSelectionFromFilter({ type: 'and', filters: remainingFilters }, context)
+                : QueryPlan.Plan.make([
+                    {
+                      _tag: 'SelectStep',
+                      scope: context.scope,
+                      selector: { _tag: 'WildcardSelector' },
+                    },
+                    ...this._generateDeletedHandlingSteps(context),
+                  ]);
+
+          const childOfSteps: QueryPlan.Step[] = childOfFilters.map((f) => ({
+            _tag: 'FilterStep' as const,
+            filter: f,
+          }));
+
+          return QueryPlan.Plan.make([...innerPlan.steps, ...childOfSteps]);
+        }
+
+        // Simple AND: all filters are root-executable against in-memory objects after a wildcard select.
+        // Supports combining true root selectors (`object`, `tag`) plus negations / unions of those.
+        // Property-level filters (`compare`, `in`, `range`, `contains`) are intentionally excluded because
+        // `filterMatchObject` only handles them as nested predicates inside an `object` filter, not at the root.
+        const isRootExecutable = (filter: QueryAST.Filter): boolean => {
+          switch (filter.type) {
+            case 'object':
+            case 'tag':
+              return true;
+            case 'not':
+              return isRootExecutable(filter.filter);
+            case 'or':
+              return filter.filters.every(isRootExecutable);
+            default:
+              return false;
+          }
+        };
+        if (flatFilters.every(isRootExecutable)) {
+          const innerFilter: QueryAST.Filter = { type: 'and', filters: flatFilters };
+          const plannedFilter: QueryAST.Filter = context.selectionInverted
+            ? { type: 'not', filter: innerFilter }
+            : innerFilter;
+          return QueryPlan.Plan.make([
+            {
+              _tag: 'SelectStep',
+              scope: context.scope,
+              selector: { _tag: 'WildcardSelector' },
+            },
+            ...this._generateDeletedHandlingSteps(context),
+            {
+              _tag: 'FilterStep',
+              filter: plannedFilter,
+            },
+          ]);
+        }
+
+        throw queryTooComplexError(context.originalQuery);
+      }
       case 'or':
         // Optimized case
         if (filter.filters.every(isTrivialTypenameFilter)) {
-          const typenames = filter.filters.map((f) => {
-            invariant(f.type === 'object' && f.typename !== null);
-            return f.typename;
+          const typenames = filter.filters.map((filter) => {
+            invariant(filter.type === 'object' && filter.typename !== null);
+            return filter.typename;
           });
+
           return QueryPlan.Plan.make([
             {
               _tag: 'SelectStep',
-              spaces: context.selectionSpaces,
+              scope: context.scope,
               selector: {
                 _tag: 'TypeSelector',
                 typename: typenames as DXN.String[],
@@ -202,13 +423,18 @@ export class QueryPlanner {
               },
             },
             ...this._generateDeletedHandlingSteps(context),
+            {
+              _tag: 'FilterStep',
+              filter: { ...filter },
+            },
           ]);
         } else {
-          throw new QueryError('Query too complex', { context: { query: context.originalQuery } });
+          throw queryTooComplexError(context.originalQuery);
         }
 
       default:
-        throw new QueryError(`Unsupported filter type: ${(filter as any).type}`, {
+        throw new QueryError({
+          message: `Unsupported filter type: ${(filter as any).type}`,
           context: { query: context.originalQuery },
         });
     }
@@ -286,7 +512,7 @@ export class QueryPlanner {
         traversal: {
           _tag: 'ReferenceTraversal',
           direction: 'incoming',
-          property: query.property,
+          property: query.property ?? null,
         },
       },
       ...this._generateDeletedHandlingSteps(context),
@@ -392,6 +618,38 @@ export class QueryPlanner {
     ]);
   }
 
+  private _generateHierarchyTraversalClause(
+    query: QueryAST.QueryHierarchyTraversalClause,
+    context: GenerationContext,
+  ): QueryPlan.Plan {
+    return QueryPlan.Plan.make([
+      ...this._generate(query.anchor, context).steps,
+      {
+        _tag: 'TraverseStep',
+        traversal: {
+          _tag: 'HierarchyTraversal',
+          direction: query.direction,
+        },
+      },
+      ...this._generateDeletedHandlingSteps(context),
+    ]);
+  }
+
+  private _validateQueryScoped(query: QueryAST.Query): void {
+    let hasScope = false;
+    QueryAST.visit(query, (node) => {
+      if (node.type === 'from') {
+        hasScope = true;
+      }
+    });
+    if (!hasScope) {
+      throw new QueryError({
+        message: 'Query must be scoped with a from() clause',
+        context: { query },
+      });
+    }
+  }
+
   /**
    * Removes filter steps that have no predicates.
    */
@@ -425,6 +683,176 @@ export class QueryPlanner {
     // TODO(dmaretskyi): Implement this.
     return plan;
   }
+
+  private _generateOrderClause(query: QueryAST.QueryOrderClause, context: GenerationContext): QueryPlan.Plan {
+    return QueryPlan.Plan.make([
+      ...this._generate(query.query, context).steps,
+      {
+        _tag: 'OrderStep',
+        order: query.order,
+      },
+    ]);
+  }
+
+  private _generateLimitClause(query: QueryAST.QueryLimitClause, context: GenerationContext): QueryPlan.Plan {
+    return QueryPlan.Plan.make([
+      ...this._generate(query.query, context).steps,
+      {
+        _tag: 'LimitStep',
+        limit: query.limit,
+      },
+    ]);
+  }
+
+  // After complete plan is built, inspect it from the end:
+  //   - Walk backwards until hitting an object set changer.
+  //   - If an order step is found, skip.
+  //   - Otherwise insert natural order before any limit steps (since limit should be applied after ordering).
+  // This method is recursive and also processes sub-plans in unions and set differences.
+  private _ensureOrderStep(plan: QueryPlan.Plan): QueryPlan.Plan {
+    // First, recursively process any sub-plans.
+    const processedSteps = plan.steps.map((step): QueryPlan.Step => {
+      if (step._tag === 'UnionStep') {
+        return {
+          _tag: 'UnionStep',
+          plans: step.plans.map((subPlan) => this._ensureOrderStep(subPlan)),
+        };
+      } else if (step._tag === 'SetDifferenceStep') {
+        return {
+          _tag: 'SetDifferenceStep',
+          source: this._ensureOrderStep(step.source),
+          exclude: this._ensureOrderStep(step.exclude),
+        };
+      }
+      return step;
+    });
+
+    const processedPlan = QueryPlan.Plan.make(processedSteps);
+
+    const OBJECT_SET_CHANGERS = new Set(['SelectStep', 'TraverseStep', 'UnionStep', 'SetDifferenceStep']);
+    for (let i = processedPlan.steps.length - 1; i >= 0; i--) {
+      const step = processedPlan.steps[i];
+      if (step._tag === 'OrderStep') {
+        return processedPlan;
+      }
+      if (OBJECT_SET_CHANGERS.has(step._tag)) {
+        break;
+      }
+    }
+
+    // Find the position to insert the order step (before any limit steps).
+    let insertIndex = processedPlan.steps.length;
+    for (let i = processedPlan.steps.length - 1; i >= 0; i--) {
+      if (processedPlan.steps[i]._tag === 'LimitStep') {
+        insertIndex = i;
+      } else {
+        break;
+      }
+    }
+
+    const newSteps = [...processedPlan.steps];
+    newSteps.splice(insertIndex, 0, {
+      _tag: 'OrderStep',
+      order: [Order.natural.ast],
+    });
+
+    return QueryPlan.Plan.make(newSteps);
+  }
+
+  /**
+   * Propagates limits from LimitStep to SelectStep and OrderStep when possible.
+   * Limits can only be propagated if there are no unions or traversals between the LimitStep and SelectStep/OrderStep.
+   */
+  private _optimizeLimits(plan: QueryPlan.Plan): QueryPlan.Plan {
+    // First, recursively process any sub-plans.
+    const processedSteps = plan.steps.map((step): QueryPlan.Step => {
+      if (step._tag === 'UnionStep') {
+        return {
+          _tag: 'UnionStep',
+          plans: step.plans.map((subPlan) => this._optimizeLimits(subPlan)),
+        };
+      } else if (step._tag === 'SetDifferenceStep') {
+        return {
+          _tag: 'SetDifferenceStep',
+          source: this._optimizeLimits(step.source),
+          exclude: this._optimizeLimits(step.exclude),
+        };
+      }
+      return step;
+    });
+
+    // Find if there's a LimitStep at the end (possibly after OrderStep).
+    let limitStepIndex = -1;
+    let limitValue: number | undefined;
+
+    for (let i = processedSteps.length - 1; i >= 0; i--) {
+      const step = processedSteps[i];
+      if (step._tag === 'LimitStep') {
+        limitStepIndex = i;
+        limitValue = step.limit;
+        break;
+      }
+      // Only allow OrderStep after LimitStep when looking backwards.
+      if (step._tag !== 'OrderStep') {
+        break;
+      }
+    }
+
+    if (limitStepIndex === -1 || limitValue === undefined) {
+      return QueryPlan.Plan.make(processedSteps);
+    }
+
+    // Check if we can propagate the limit to a SelectStep and/or OrderStep.
+    // We can only do this if there are no unions, traversals, or set differences between them.
+    const BLOCKERS = new Set(['UnionStep', 'TraverseStep', 'SetDifferenceStep']);
+
+    let selectStepIndex = -1;
+    let orderStepIndex = -1;
+    for (let i = 0; i < limitStepIndex; i++) {
+      const step = processedSteps[i];
+      if (step._tag === 'SelectStep') {
+        selectStepIndex = i;
+      }
+      if (step._tag === 'OrderStep') {
+        orderStepIndex = i;
+      }
+      if (BLOCKERS.has(step._tag)) {
+        // Found a blocker after the select/order step - can't propagate.
+        selectStepIndex = -1;
+        orderStepIndex = -1;
+      }
+    }
+
+    if (selectStepIndex === -1 && orderStepIndex === -1) {
+      return QueryPlan.Plan.make(processedSteps);
+    }
+
+    // Create a mutable copy of steps to modify.
+    const newSteps = [...processedSteps];
+
+    // Propagate the limit to the SelectStep if found.
+    if (selectStepIndex !== -1) {
+      const selectStep = newSteps[selectStepIndex] as QueryPlan.SelectStep;
+      newSteps[selectStepIndex] = {
+        ...selectStep,
+        limit: limitValue,
+      };
+    }
+
+    // Propagate the limit to the OrderStep if found.
+    if (orderStepIndex !== -1) {
+      const orderStep = newSteps[orderStepIndex] as QueryPlan.OrderStep;
+      newSteps[orderStepIndex] = {
+        ...orderStep,
+        limit: limitValue,
+      };
+    }
+
+    // Remove the LimitStep since limit is now propagated.
+    newSteps.splice(limitStepIndex, 1);
+
+    return QueryPlan.Plan.make(newSteps);
+  }
 }
 
 /**
@@ -437,9 +865,9 @@ type GenerationContext = {
   originalQuery: QueryAST.Query | null;
 
   /**
-   * Which spaces to select from.
+   * The scope to select from (space IDs, queues, etc.).
    */
-  selectionSpaces: readonly SpaceId[];
+  scope: QueryAST.Scope;
 
   /**
    * How to handle deleted objects.
@@ -454,7 +882,7 @@ type GenerationContext = {
 
 const DEFAULT_CONTEXT: GenerationContext = {
   originalQuery: null,
-  selectionSpaces: [],
+  scope: {},
   deletedHandling: 'exclude',
   selectionInverted: false,
 };
@@ -474,6 +902,54 @@ const createRelationTraversalStep = (direction: QueryPlan.RelationTraversal['dir
   },
 });
 
+/**
+ * Recursively flattens nested `and` filters into a single list.
+ */
+const _flattenAnd = (filters: readonly QueryAST.Filter[]): QueryAST.Filter[] => {
+  const result: QueryAST.Filter[] = [];
+  for (const f of filters) {
+    if (f.type === 'and') {
+      result.push(..._flattenAnd(f.filters));
+    } else {
+      result.push(f);
+    }
+  }
+  return result;
+};
+
+const _timestampFilterToSelector = (filter: QueryAST.FilterTimestamp): QueryPlan.TimestampSelector => {
+  const selector: QueryPlan.TimestampSelector = { _tag: 'TimestampSelector' };
+  const key =
+    filter.field === 'updatedAt'
+      ? filter.operator === 'gte' || filter.operator === 'gt'
+        ? 'updatedAfter'
+        : 'updatedBefore'
+      : filter.operator === 'gte' || filter.operator === 'gt'
+        ? 'createdAfter'
+        : 'createdBefore';
+  selector[key] = filter.value;
+  return selector;
+};
+
+const _mergeTimestampSelectors = (selectors: QueryPlan.TimestampSelector[]): QueryPlan.TimestampSelector => {
+  const merged: QueryPlan.TimestampSelector = { _tag: 'TimestampSelector' };
+  for (const s of selectors) {
+    if (s.updatedAfter != null) {
+      merged.updatedAfter = Math.max(merged.updatedAfter ?? 0, s.updatedAfter);
+    }
+    if (s.updatedBefore != null) {
+      merged.updatedBefore = Math.min(merged.updatedBefore ?? Infinity, s.updatedBefore);
+    }
+    if (s.createdAfter != null) {
+      merged.createdAfter = Math.max(merged.createdAfter ?? 0, s.createdAfter);
+    }
+    if (s.createdBefore != null) {
+      merged.createdBefore = Math.min(merged.createdBefore ?? Infinity, s.createdBefore);
+    }
+  }
+  return merged;
+};
+
 const isTrivialTypenameFilter = (filter: QueryAST.Filter): boolean => {
   return (
     filter.type === 'object' &&