diff --git a/src/languages/spark/spark.formatter.ts b/src/languages/spark/spark.formatter.ts index 192675eadb..96689efd1d 100644 --- a/src/languages/spark/spark.formatter.ts +++ b/src/languages/spark/spark.formatter.ts @@ -144,7 +144,9 @@ export const spark: DialectOptions = { identTypes: ['``'], identChars: { allowFirstCharNumber: true }, variableTypes: [{ quote: '{}', prefixes: ['$'], requirePrefix: true }], - operators: ['%', '~', '^', '|', '&', '<=>', '==', '!', '||', '->'], + // ':' is optional between STRUCT field name and type: + // https://spark.apache.org/docs/latest/sql-ref-datatypes.html + operators: ['%', '~', '^', '|', '&', '<=>', '==', '!', '||', '->', ':'], postProcess, }, formatOptions: { @@ -154,24 +156,94 @@ export const spark: DialectOptions = { }; function postProcess(tokens: Token[]) { - return tokens.map((token, i) => { - const prevToken = tokens[i - 1] || EOF_TOKEN; - const nextToken = tokens[i + 1] || EOF_TOKEN; - - // [WINDOW](...) - if (isToken.WINDOW(token) && nextToken.type === TokenType.OPEN_PAREN) { - // This is a function call, treat it as a reserved function name - return { ...token, type: TokenType.RESERVED_FUNCTION_NAME }; - } + return combineParameterizedTypes( + tokens.map((token, i) => { + const prevToken = tokens[i - 1] || EOF_TOKEN; + const nextToken = tokens[i + 1] || EOF_TOKEN; + + // [WINDOW](...) + if (isToken.WINDOW(token) && nextToken.type === TokenType.OPEN_PAREN) { + // This is a function call, treat it as a reserved function name + return { ...token, type: TokenType.RESERVED_FUNCTION_NAME }; + } - // TODO: deprecate this once ITEMS is merged with COLLECTION - if (token.text === 'ITEMS' && token.type === TokenType.RESERVED_KEYWORD) { - if (!(prevToken.text === 'COLLECTION' && nextToken.text === 'TERMINATED')) { - // this is a word and not COLLECTION ITEMS - return { ...token, type: TokenType.IDENTIFIER, text: token.raw }; + // TODO: deprecate this once ITEMS is merged with COLLECTION + if (token.text === 'ITEMS' && token.type === TokenType.RESERVED_KEYWORD) { + if (!(prevToken.text === 'COLLECTION' && nextToken.text === 'TERMINATED')) { + // this is a word and not COLLECTION ITEMS + return { ...token, type: TokenType.IDENTIFIER, text: token.raw }; + } } + + return token; + }) + ); +} + +// Combines ARRAY, MAP, STRUCT into a single token. +// Spark STRUCT fields allow an optional colon: STRUCT or STRUCT. +// https://spark.apache.org/docs/latest/sql-ref-datatypes.html +function combineParameterizedTypes(tokens: Token[]) { + const processed: Token[] = []; + for (let i = 0; i < tokens.length; i++) { + const token = tokens[i]; + + if (isComplexTypeName(token) && tokens[i + 1]?.text === '<') { + const endIndex = findClosingAngleBracketIndex(tokens, i + 1); + if (endIndex === -1) { + processed.push(token); + continue; + } + const typeDefTokens = tokens.slice(i, endIndex + 1); + processed.push({ + type: TokenType.RESERVED_DATA_TYPE, + raw: typeDefTokens.map(formatTypeDefToken('raw')).join(''), + text: typeDefTokens.map(formatTypeDefToken('text')).join(''), + start: token.start, + }); + i = endIndex; + } else { + processed.push(token); + } + } + return processed; +} + +function isComplexTypeName(token: Token): boolean { + return ( + (token.text === 'ARRAY' || token.text === 'STRUCT' || token.text === 'MAP') && + (token.type === TokenType.RESERVED_DATA_TYPE || token.type === TokenType.RESERVED_FUNCTION_NAME) + ); +} + +const formatTypeDefToken = + (key: Extract) => + (token: Token, i: number, tokens: Token[]): string => { + if (token.text === ':' || token.type === TokenType.COMMA) { + return token[key] + ' '; + } + if (isTypeDefFieldName(token) && tokens[i + 1]?.text !== ':') { + return token[key] + ' '; } + return token[key]; + }; - return token; - }); +function isTypeDefFieldName(token: Token): boolean { + return token.type === TokenType.IDENTIFIER || token.type === TokenType.QUOTED_IDENTIFIER; +} + +function findClosingAngleBracketIndex(tokens: Token[], startIndex: number): number { + let level = 0; + for (let i = startIndex; i < tokens.length; i++) { + const token = tokens[i]; + if (token.text === '<') { + level++; + } else if (token.text === '>') { + level--; + } + if (level === 0) { + return i; + } + } + return -1; } diff --git a/test/spark.test.ts b/test/spark.test.ts index 62a9786eea..49b72ca3fd 100644 --- a/test/spark.test.ts +++ b/test/spark.test.ts @@ -152,4 +152,51 @@ describe('SparkFormatter', () => { ALTER COLUMN FirstName COMMENT "new comment"; `); }); + + // regression test for #900 + // Spark SHOW CREATE TABLE emits STRUCT; the colon is optional. + it('supports optional colon in STRUCT field specs', () => { + expect( + format(`CREATE TABLE t_a ( +a bigint, +b struct +)`) + ).toBe(dedent` + CREATE TABLE t_a (a bigint, b struct) + `); + }); + + it('supports STRUCT field specs without colon', () => { + expect(format(`CREATE TABLE t_a (a bigint, b struct )`)).toBe( + dedent` + CREATE TABLE t_a (a bigint, b struct) + ` + ); + }); + + it('supports nested ARRAY, MAP and STRUCT types with optional colons', () => { + expect( + format(`CREATE TABLE family ( +name STRING, +friends ARRAY, +children MAP, +address STRUCT +)`) + ).toBe(dedent` + CREATE TABLE family ( + name STRING, + friends ARRAY, + children MAP, + address STRUCT + ) + `); + }); + + it('applies dataTypeCase to STRUCT field types', () => { + expect(format('CREATE TABLE t (b struct)', { dataTypeCase: 'upper' })).toBe( + dedent` + CREATE TABLE t (b STRUCT) + ` + ); + }); });