fix: PIVOT FOR...IN AS alias parsing, COPY PARTITION BY roundtrip

- Add pivotInValue grammar rule to support aliased values in PIVOT IN
  lists, e.g. FOR symbol IN ('BTC-USD' AS bitcoin). Both parsing and
  toSql round-trip now handle the AS alias correctly.
- Fix COPY FROM visitor dropping the BY token from PARTITION BY,
  which caused toSql to emit invalid SQL rejected by QuestDB.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: emrberk

src/parser/ast.ts

@@ -777,9 +777,15 @@ export interface PivotForClause extends AstNode {
   in: PivotInSource
 }
 
+export interface PivotInValue extends AstNode {
+  type: "pivotInValue"
+  expression: Expression
+  alias?: string
+}
+
 export interface PivotInSource extends AstNode {
   type: "pivotIn"
-  values?: Expression[]
+  values?: PivotInValue[]
   select?: SelectStatement
 }
 

src/parser/cst-types.d.ts

@@ -1769,12 +1769,23 @@ export type PivotForClauseCstChildren = {
   columnRef: ColumnRefCstNode[];
   In: IToken[];
   LParen: IToken[];
-  expression?: (ExpressionCstNode)[];
+  pivotInValue?: (PivotInValueCstNode)[];
   Comma?: IToken[];
   selectStatement?: SelectStatementCstNode[];
   RParen: IToken[];
 };
 
+export interface PivotInValueCstNode extends CstNode {
+  name: "pivotInValue";
+  children: PivotInValueCstChildren;
+}
+
+export type PivotInValueCstChildren = {
+  expression: ExpressionCstNode[];
+  As?: IToken[];
+  identifier?: IdentifierCstNode[];
+};
+
 export interface ExpressionCstNode extends CstNode {
   name: "expression";
   children: ExpressionCstChildren;
@@ -2428,6 +2439,7 @@ export interface ICstNodeVisitor<IN, OUT> extends ICstVisitor<IN, OUT> {
   pivotBody(children: PivotBodyCstChildren, param?: IN): OUT;
   pivotAggregation(children: PivotAggregationCstChildren, param?: IN): OUT;
   pivotForClause(children: PivotForClauseCstChildren, param?: IN): OUT;
+  pivotInValue(children: PivotInValueCstChildren, param?: IN): OUT;
   expression(children: ExpressionCstChildren, param?: IN): OUT;
   orExpression(children: OrExpressionCstChildren, param?: IN): OUT;
   andExpression(children: AndExpressionCstChildren, param?: IN): OUT;

src/parser/parser.ts

@@ -3201,10 +3201,10 @@ class QuestDBParser extends CstParser {
     this.OR([
       {
         ALT: () => {
-          this.SUBRULE1(this.expression)
+          this.SUBRULE1(this.pivotInValue)
           this.MANY(() => {
             this.CONSUME(Comma)
-            this.SUBRULE2(this.expression)
+            this.SUBRULE2(this.pivotInValue)
           })
         },
       },
@@ -3213,6 +3213,16 @@ class QuestDBParser extends CstParser {
     this.CONSUME(RParen)
   })
 
+  // A single value inside PIVOT FOR ... IN (...), optionally aliased:
+  //   'BTC-USD' AS bitcoin
+  private pivotInValue = this.RULE("pivotInValue", () => {
+    this.SUBRULE(this.expression)
+    this.OPTION(() => {
+      this.CONSUME(As)
+      this.SUBRULE(this.identifier)
+    })
+  })
+
   // ==========================================================================
   // Expressions
   // ==========================================================================

src/parser/toSql.ts

@@ -1346,7 +1346,15 @@ function pivotClauseToSql(pivot: AST.PivotClause): string {
     if (p.in.select) {
       forSql += selectToSql(p.in.select)
     } else if (p.in.values) {
-      forSql += p.in.values.map(expressionToSql).join(", ")
+      forSql += p.in.values
+        .map((v) => {
+          let sql = expressionToSql(v.expression)
+          if (v.alias) {
+            sql += ` AS ${escapeIdentifier(v.alias)}`
+          }
+          return sql
+        })
+        .join(", ")
     }
     forSql += ")"
     pivotParts.push(forSql)
@@ -1398,7 +1406,15 @@ function pivotToSql(stmt: AST.PivotStatement): string {
     if (pivot.in.select) {
       forSql += selectToSql(pivot.in.select)
     } else if (pivot.in.values) {
-      forSql += pivot.in.values.map(expressionToSql).join(", ")
+      forSql += pivot.in.values
+        .map((v) => {
+          let sql = expressionToSql(v.expression)
+          if (v.alias) {
+            sql += ` AS ${escapeIdentifier(v.alias)}`
+          }
+          return sql
+        })
+        .join(", ")
     }
     forSql += ")"
     pivotParts.push(forSql)

src/parser/visitor.ts

@@ -113,6 +113,7 @@ import type {
   PivotAggregationCstChildren,
   PivotBodyCstChildren,
   PivotForClauseCstChildren,
+  PivotInValueCstChildren,
   PivotStatementCstChildren,
   PrimaryExpressionCstChildren,
   QualifiedNameCstChildren,
@@ -2343,6 +2344,7 @@ class QuestDBVisitor extends BaseVisitor {
 
     let key = keyToken?.image ?? "OPTION"
     if (ctx.On && ctx.Error) key = "ON ERROR"
+    else if (ctx.Partition && ctx.By) key = "PARTITION BY"
     const result: AST.CopyOption = {
       type: "copyOption",
       key,
@@ -2775,14 +2777,25 @@ class QuestDBVisitor extends BaseVisitor {
     }
     if (ctx.selectStatement) {
       result.in.select = this.visit(ctx.selectStatement) as AST.SelectStatement
-    } else if (ctx.expression && ctx.expression.length > 0) {
-      result.in.values = ctx.expression.map(
-        (e: CstNode) => this.visit(e) as AST.Expression,
+    } else if (ctx.pivotInValue && ctx.pivotInValue.length > 0) {
+      result.in.values = ctx.pivotInValue.map(
+        (e: CstNode) => this.visit(e) as AST.PivotInValue,
       )
     }
     return result
   }
 
+  pivotInValue(ctx: PivotInValueCstChildren): AST.PivotInValue {
+    const result: AST.PivotInValue = {
+      type: "pivotInValue",
+      expression: this.visit(ctx.expression) as AST.Expression,
+    }
+    if (ctx.As && ctx.identifier) {
+      result.alias = this.extractIdentifierName(ctx.identifier[0].children)
+    }
+    return result
+  }
+
   // ==========================================================================
   // Expressions
   // ==========================================================================

tests/parser.test.ts

@@ -2125,6 +2125,37 @@ describe("QuestDB Parser", () => {
         expect(pivot.groupBy).toHaveLength(1)
       }
     })
+
+    it("should parse PIVOT FOR...IN with AS aliases on values", () => {
+      const result = parseToAst(
+        "trades PIVOT (avg(price) FOR symbol IN ('BTC-USD' AS bitcoin, 'ETH-USD' AS ethereum))",
+      )
+      expect(result.errors).toHaveLength(0)
+      expect(result.ast).toHaveLength(1)
+      const pivot = result.ast[0]
+      expect(pivot.type).toBe("pivot")
+      if (pivot.type === "pivot") {
+        expect(pivot.pivots).toHaveLength(1)
+        const values = pivot.pivots[0].in.values!
+        expect(values).toHaveLength(2)
+        expect(values[0].alias).toBe("bitcoin")
+        expect(values[1].alias).toBe("ethereum")
+      }
+    })
+
+    it("should parse PIVOT FOR...IN with mixed aliased and unaliased values", () => {
+      const result = parseToAst(
+        "trades PIVOT (avg(price) FOR symbol IN ('BTC-USD' AS bitcoin, 'ETH-USD'))",
+      )
+      expect(result.errors).toHaveLength(0)
+      const pivot = result.ast[0]
+      if (pivot.type === "pivot") {
+        const values = pivot.pivots[0].in.values!
+        expect(values).toHaveLength(2)
+        expect(values[0].alias).toBe("bitcoin")
+        expect(values[1].alias).toBeUndefined()
+      }
+    })
   })
 
   describe("PIVOT multi-statement boundary", () => {
@@ -2168,6 +2199,8 @@ orders PIVOT (sum(amount) FOR status IN ('open'))`
       "(SELECT * FROM trades) PIVOT (sum(amount) FOR category IN ('food', 'drinks'))",
       "trades WHERE price > 100 PIVOT (sum(amount) FOR category IN ('food', 'drinks'))",
       "trades PIVOT (sum(amount) FOR category IN ('food', 'drinks')) AS p",
+      "trades PIVOT (avg(price) FOR symbol IN ('BTC-USD' AS bitcoin, 'ETH-USD' AS ethereum))",
+      "trades PIVOT (avg(price) FOR symbol IN ('BTC-USD' AS bitcoin, 'ETH-USD'))",
     ]
 
     for (const query of queries) {
@@ -4745,6 +4778,14 @@ orders PIVOT (sum(amount) FOR status IN ('open'))`
       expect(result.ast[0].type).toBe("copyFrom")
     })
 
+    it("should round-trip COPY FROM with PARTITION BY", () => {
+      const sql = "COPY weather FROM '/tmp/weather.csv' WITH PARTITION BY DAY"
+      const { ast, errors } = parseToAst(sql)
+      expect(errors).toHaveLength(0)
+      const output = toSql(ast)
+      expect(output).toContain("PARTITION BY DAY")
+    })
+
     // BACKUP ABORT
     it("should parse BACKUP ABORT", () => {
       const result = parseToAst("BACKUP ABORT")