fix arrays in extend blocks

malloydata · Nov 8, 2024 · 5d83457 · 5d83457
1 parent 5ce9ba6
commit 5d83457
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 16 deletions.
diff --git a/packages/malloy/src/lang/ast/field-space/join-space-field.ts b/packages/malloy/src/lang/ast/field-space/join-space-field.ts
@@ -28,9 +28,8 @@ import {StructSpaceField} from './static-space';
 export class JoinSpaceField extends StructSpaceField {
   constructor(
     readonly parameterSpace: ParameterSpace,
-    readonly join: Join,
-    parentDialect: string
+    readonly join: Join
   ) {
-    super(join.structDef(parameterSpace), parentDialect);
+    super(join.structDef(parameterSpace));
   }
 }
diff --git a/packages/malloy/src/lang/ast/field-space/static-space.ts b/packages/malloy/src/lang/ast/field-space/static-space.ts
@@ -72,7 +72,7 @@ export class StaticSpace implements FieldSpace {
 
   defToSpaceField(from: FieldDef): SpaceField {
     if (isJoined(from)) {
-      return new StructSpaceField(from, this.fromStruct.dialect);
+      return new StructSpaceField(from);
     } else if (isTurtleDef(from)) {
       return new IRViewField(this, from);
     }
@@ -138,7 +138,7 @@ export class StaticSpace implements FieldSpace {
   lookup(path: FieldName[]): LookupResult {
     const head = path[0];
     const rest = path.slice(1);
-    const found = this.entry(head.refString);
+    let found = this.entry(head.refString);
     if (!found) {
       return {
         error: {
@@ -151,9 +151,16 @@ export class StaticSpace implements FieldSpace {
     if (found instanceof SpaceField) {
       const definition = found.fieldDef();
       if (definition) {
+        if (isJoined(definition)) {
+          // We have looked up a field which is a join, but not a StructSpaceField
+          // because it is someting like "dimension: joinedArray is arrayComputation"
+          // so it is a FieldDefinitionValue.
+          // TODO don't make one of these every time you do a lookup
+          found = new StructSpaceField(definition);
+        }
         head.addReference({
           type:
-            found instanceof StructSpaceFieldBase
+            found instanceof StructSpaceFieldBase && path.length > 1
               ? 'joinReference'
               : 'fieldReference',
           definition,
@@ -195,10 +202,8 @@ export class StaticSpace implements FieldSpace {
 }
 
 export class StructSpaceField extends StructSpaceFieldBase {
-  private parentDialect: string;
-  constructor(def: JoinFieldDef, dialect: string) {
+  constructor(def: JoinFieldDef) {
     super(def);
-    this.parentDialect = dialect;
   }
 
   get fieldSpace(): FieldSpace {

diff --git a/packages/malloy/src/lang/ast/source-properties/join.ts b/packages/malloy/src/lang/ast/source-properties/join.ts
@@ -63,7 +63,7 @@ export abstract class Join
     fs.newEntry(
       this.name.refString,
       this,
-      new JoinSpaceField(fs.parameterSpace(), this, fs.dialect)
+      new JoinSpaceField(fs.parameterSpace(), this)
     );
   }
 

diff --git a/test/src/databases/all/composite-atomic.spec.ts b/test/src/databases/all/composite-atomic.spec.ts
@@ -44,18 +44,48 @@ describe.each(runtimes.runtimeList)(
       });
       test('array literal in source', async () => {
         await expect(`
-          run: duckdb.sql("select 1") extend { dimension: d1 is [1,2,3,4,5,6] }
-          -> { select: die_roll is d1.each }
+          run: duckdb.sql("select 1")
+          extend { dimension: d4 is [1,2,3,4] }
+          -> { select: die_roll is d4.each }
         `).malloyResultMatches(runtime, [
           {die_roll: 1},
           {die_roll: 2},
           {die_roll: 3},
           {die_roll: 4},
-          {die_roll: 5},
-          {die_roll: 6},
         ]);
       });
-      test.todo('cross join two arrays');
+      test('array literal in extend block', async () => {
+        await expect(`
+          run: duckdb.sql("select 1") -> {
+            extend: { dimension: d4 is [1,2,3,4] }
+            select: die_roll is d4.each
+          }
+        `).malloyResultMatches(runtime, [
+          {die_roll: 1},
+          {die_roll: 2},
+          {die_roll: 3},
+          {die_roll: 4},
+        ]);
+      });
+      test.skip('cross join arrays', async () => {
+        await expect(`
+          # test.verbose
+          run: duckdb.sql("select 1") extend {
+            dimension: d1 is [1,2,3,4], d2 is [1,2,3,4]
+          } -> {
+            group_by: roll is d1.each + d2.each
+            aggregate: rolls is count()
+          }
+        `).malloyResultMatches(runtime, [
+          {roll: 2, rolls: 1},
+          {roll: 3, rolls: 2},
+          {roll: 4, rolls: 3},
+          {roll: 5, rolls: 4},
+          {roll: 6, rolls: 3},
+          {roll: 7, rolls: 2},
+          {roll: 8, rolls: 1},
+        ]);
+      });
     });
     describe('record', () => {
       const record = 'duckdb.sql("SELECT {s: 0, m: 1, l:2, xl: 3} as record")';
@@ -70,7 +100,7 @@ describe.each(runtimes.runtimeList)(
           'record/xl': 3,
         });
       });
-      test('record literal', async () => {
+      test('record literal can be selected', async () => {
         await expect(`
           run: duckdb.sql("select 1") -> {
             select: record is {s is 0, m is 1, l is 2, xl is 3}
@@ -82,6 +112,19 @@ describe.each(runtimes.runtimeList)(
           'record/xl': 3,
         });
       });
+      test('record literal from a source', async () => {
+        await expect(`
+          run: duckdb.sql("select 1") -> {
+            extend: { dimension: record is {s is 0, m is 1, l is 2, xl is 3} }
+            select: record
+          }
+        `).malloyResultMatches(runtime, {
+          'record/s': 0,
+          'record/m': 1,
+          'record/l': 2,
+          'record/xl': 3,
+        });
+      });
       test.todo('array of records can be selected');
       test.todo('array of records literal');
     });