Skip to content

Commit 79e92ad

Browse files
VaggelisDVaggelisD
authored
feat(spark): Default naming of STRUCT fields (tobymao#3991)
1 parent ee9dc39 commit 79e92ad

4 files changed

Lines changed: 31 additions & 5 deletions

File tree

‎sqlglot/dialects/hive.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -436,6 +436,14 @@ def _parse_parameter(self) -> exp.Parameter:
436436
self._match(TokenType.R_BRACE)
437437
return self.expression(exp.Parameter, this=this, expression=expression)
438438

439+
def _to_prop_eq(self, expression: exp.Expression, index: int) -> exp.Expression:
440+
if isinstance(expression, exp.Column):
441+
key = expression.this
442+
else:
443+
key = exp.to_identifier(f"col{index + 1}")
444+
445+
return self.expression(exp.PropertyEQ, this=key, expression=expression)
446+
439447
class Generator(generator.Generator):
440448
LIMIT_FETCH = "LIMIT"
441449
TABLESAMPLE_WITH_METHOD = False

‎sqlglot/parser.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5115,10 +5115,13 @@ def _parse_function_call(
51155115
self._match_r_paren(this)
51165116
return self._parse_window(this)
51175117

5118+
def _to_prop_eq(self, expression: exp.Expression, index: int) -> exp.Expression:
5119+
return expression
5120+
51185121
def _kv_to_prop_eq(self, expressions: t.List[exp.Expression]) -> t.List[exp.Expression]:
51195122
transformed = []
51205123

5121-
for e in expressions:
5124+
for index, e in enumerate(expressions):
51225125
if isinstance(e, self.KEY_VALUE_DEFINITIONS):
51235126
if isinstance(e, exp.Alias):
51245127
e = self.expression(exp.PropertyEQ, this=e.args.get("alias"), expression=e.this)
@@ -5130,6 +5133,8 @@ def _kv_to_prop_eq(self, expressions: t.List[exp.Expression]) -> t.List[exp.Expr
51305133

51315134
if isinstance(e.this, exp.Column):
51325135
e.this.replace(e.this.this)
5136+
else:
5137+
e = self._to_prop_eq(e, index)
51335138

51345139
transformed.append(e)
51355140

‎tests/dialects/test_presto.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -716,9 +716,6 @@ def test_presto(self):
716716
)
717717
self.validate_all(
718718
"SELECT ROW(1, 2)",
719-
read={
720-
"spark": "SELECT STRUCT(1, 2)",
721-
},
722719
write={
723720
"presto": "SELECT ROW(1, 2)",
724721
"spark": "SELECT STRUCT(1, 2)",

‎tests/dialects/test_spark.py‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -485,7 +485,7 @@ def test_spark(self):
485485
)
486486
self.validate_all(
487487
"SELECT CAST(STRUCT('fooo') AS STRUCT<a: VARCHAR(2)>)",
488-
write={"spark": "SELECT CAST(STRUCT('fooo') AS STRUCT<a: STRING>)"},
488+
write={"spark": "SELECT CAST(STRUCT('fooo' AS col1) AS STRUCT<a: STRING>)"},
489489
)
490490
self.validate_all(
491491
"SELECT CAST(123456 AS VARCHAR(3))",
@@ -718,6 +718,22 @@ def test_spark(self):
718718
},
719719
)
720720

721+
self.validate_all(
722+
"SELECT STRUCT(1, 2)",
723+
write={
724+
"spark": "SELECT STRUCT(1 AS col1, 2 AS col2)",
725+
"presto": "SELECT CAST(ROW(1, 2) AS ROW(col1 INTEGER, col2 INTEGER))",
726+
"duckdb": "SELECT {'col1': 1, 'col2': 2}",
727+
},
728+
)
729+
self.validate_all(
730+
"SELECT STRUCT(x, 1, y AS col3, STRUCT(5)) FROM t",
731+
write={
732+
"spark": "SELECT STRUCT(x AS x, 1 AS col2, y AS col3, STRUCT(5 AS col1) AS col4) FROM t",
733+
"duckdb": "SELECT {'x': x, 'col2': 1, 'col3': y, 'col4': {'col1': 5}} FROM t",
734+
},
735+
)
736+
721737
def test_bool_or(self):
722738
self.validate_all(
723739
"SELECT a, LOGICAL_OR(b) FROM table GROUP BY a",

0 commit comments

Comments
 (0)