Compute expanded field usage for queries

packages/malloy/src/lang/ast/query-elements/query-arrow.ts

@@ -74,19 +74,23 @@ export class QueryArrow extends QueryBase implements QueryElement {
       inputStruct = lhsQuery.outputStruct;
       fieldSpace = new StaticSourceSpace(lhsQuery.outputStruct, 'public');
     }
-    const {pipeline, annotation, outputStruct, name} =
-      this.view.pipelineComp(fieldSpace);
+    const {
+      pipeline: rhsPipeline,
+      annotation,
+      outputStruct,
+      name,
+    } = this.view.pipelineComp(fieldSpace);
 
     const query = {
       ...queryBase,
       name,
       annotation,
-      pipeline: [...queryBase.pipeline, ...pipeline],
+      pipeline: [...queryBase.pipeline, ...rhsPipeline],
     };
 
     const compositeResolvedSourceDef =
       query.compositeResolvedSourceDef ??
-      this.resolveCompositeSource(inputStruct, pipeline);
+      this.resolveCompositeSource(inputStruct, rhsPipeline);
 
     const segment = query.pipeline[0];
     if (segment !== undefined) {
@@ -107,10 +111,24 @@ export class QueryArrow extends QueryBase implements QueryElement {
       }
     }
 
+    const pipelineWithExpandedFieldUsage = [
+      // The base query (if it exists) will already have its `expandedFieldUsage` computed
+      ...queryBase.pipeline,
+      ...this.expandFieldUsage(
+        this.source instanceof Source
+          ? // If `source ->` then use the composite resolved struct,
+            compositeResolvedSourceDef ?? inputStruct
+          : // Otherwise just use the `inputStruct`
+            inputStruct,
+        rhsPipeline
+      ),
+    ];
+
     return {
       query: {
         ...query,
         compositeResolvedSourceDef,
+        pipeline: pipelineWithExpandedFieldUsage,
       },
       outputStruct,
       inputStruct,

packages/malloy/src/lang/ast/query-elements/query-base.ts

@@ -23,9 +23,9 @@
 
 import {
   hasCompositesAnywhere,
-  emptyFieldUsage,
   resolveCompositeSources,
   logCompositeError,
+  expandFieldUsage,
 } from '../../../model/composite_source_utils';
 import {
   isIndexSegment,
@@ -34,13 +34,30 @@ import {
   type Query,
   type SourceDef,
 } from '../../../model/malloy_types';
+import {ErrorFactory} from '../error-factory';
 import {detectAndRemovePartialStages} from '../query-utils';
 import {MalloyElement} from '../types/malloy-element';
 import type {QueryComp} from '../types/query-comp';
 
 export abstract class QueryBase extends MalloyElement {
   abstract queryComp(isRefOk: boolean): QueryComp;
 
+  protected expandFieldUsage(
+    inputSource: SourceDef,
+    pipeline: PipeSegment[]
+  ): PipeSegment[] {
+    return pipeline.map((segment, index) => {
+      // Theoretically for `index > 0`, this should be the output struct of the prior stage,
+      // but because field references in prior stages can only refer to definitions from `extend:`
+      // statements, we can just use an empty struct here.
+      const stageInput = index === 0 ? inputSource : ErrorFactory.structDef;
+      return {
+        ...segment,
+        expandedFieldUsage: expandFieldUsage(segment, stageInput),
+      };
+    });
+  }
+
   protected resolveCompositeSource(
     inputSource: SourceDef,
     pipeline: PipeSegment[]
@@ -54,8 +71,7 @@ export abstract class QueryBase extends MalloyElement {
       (isQuerySegment(stage1) || isIndexSegment(stage1)) &&
       hasCompositesAnywhere(inputSource)
     ) {
-      const fieldUsage = stage1.fieldUsage ?? emptyFieldUsage();
-      const resolved = resolveCompositeSources(inputSource, stage1, fieldUsage);
+      const resolved = resolveCompositeSources(inputSource, stage1);
       if (resolved.error) {
         logCompositeError(resolved.error, this);
       }

packages/malloy/src/lang/ast/query-elements/query-refine.ts

@@ -61,10 +61,16 @@ export class QueryRefine extends QueryBase implements QueryElement {
       resultPipe
     );
 
+    const pipelineWithExpandedFieldUsage = this.expandFieldUsage(
+      compositeResolvedSourceDef ?? q.inputStruct,
+      resultPipe
+    );
+
     return {
       query: {
         ...query,
         compositeResolvedSourceDef,
+        pipeline: pipelineWithExpandedFieldUsage,
       },
       outputStruct: getFinalStruct(this.refinement, q.inputStruct, resultPipe),
       inputStruct: q.inputStruct,