Fix(hive)!: improve transpilability of GET_JSON_OBJECT by parsing json path (#4980)

Author: georgesittas

sqlglot/dialects/databricks.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import typing as t
 
 from sqlglot import exp, transforms, jsonpath
 from sqlglot.dialects.dialect import (
@@ -13,13 +12,6 @@
 from sqlglot.tokens import TokenType
 
 
-def _build_json_extract(args: t.List) -> exp.JSONExtract:
-    # Transform GET_JSON_OBJECT(expr, '$.<path>') -> expr:<path>
-    this = args[0]
-    path = args[1].name.lstrip("$.")
-    return exp.JSONExtract(this=this, expression=path)
-
-
 def _jsonextract_sql(
     self: Databricks.Generator, expression: exp.JSONExtract | exp.JSONExtractScalar
 ) -> str:
@@ -46,7 +38,6 @@ class Parser(Spark.Parser):
             "DATE_ADD": build_date_delta(exp.DateAdd),
             "DATEDIFF": build_date_delta(exp.DateDiff),
             "DATE_DIFF": build_date_delta(exp.DateDiff),
-            "GET_JSON_OBJECT": _build_json_extract,
             "TO_DATE": build_formatted_time(exp.TsOrDsToDate, "databricks"),
         }
 

sqlglot/dialects/hive.py

@@ -307,7 +307,9 @@ class Parser(parser.Parser):
             "FIRST": _build_with_ignore_nulls(exp.First),
             "FIRST_VALUE": _build_with_ignore_nulls(exp.FirstValue),
             "FROM_UNIXTIME": build_formatted_time(exp.UnixToStr, "hive", True),
-            "GET_JSON_OBJECT": exp.JSONExtractScalar.from_arg_list,
+            "GET_JSON_OBJECT": lambda args, dialect: exp.JSONExtractScalar(
+                this=seq_get(args, 0), expression=dialect.to_json_path(seq_get(args, 1))
+            ),
             "LAST": _build_with_ignore_nulls(exp.Last),
             "LAST_VALUE": _build_with_ignore_nulls(exp.LastValue),
             "MAP": parser.build_var_map,

tests/dialects/test_databricks.py

@@ -68,6 +68,24 @@ def test_databricks(self):
             "FROM_UTC_TIMESTAMP(CAST(x AS TIMESTAMP), tz)",
         )
 
+        self.validate_all(
+            "SELECT c1:item[1].price",
+            read={
+                "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            },
+            write={
+                "databricks": "SELECT c1:item[1].price",
+                "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            },
+        )
+
+        self.validate_all(
+            "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            write={
+                "databricks": "SELECT c1:item[1].price",
+                "spark": "SELECT GET_JSON_OBJECT(c1, '$.item[1].price')",
+            },
+        )
         self.validate_all(
             "CREATE TABLE foo (x INT GENERATED ALWAYS AS (YEAR(y)))",
             write={