refactor: add sqlglot type conversions (#1599)

Author: chelsea-lin

bigframes/core/compile/sqlglot/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

bigframes/core/compile/sqlglot/sqlglot_types.py

@@ -0,0 +1,84 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import typing
+
+import bigframes_vendored.constants as constants
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import sqlglot as sg
+
+import bigframes.dtypes
+
+
+class SQLGlotType:
+    @classmethod
+    def from_bigframes_dtype(
+        cls,
+        bigframes_dtype: typing.Union[
+            bigframes.dtypes.DtypeString, bigframes.dtypes.Dtype, np.dtype[typing.Any]
+        ],
+    ):
+        if bigframes_dtype == bigframes.dtypes.INT_DTYPE:
+            return "INT64"
+        elif bigframes_dtype == bigframes.dtypes.FLOAT_DTYPE:
+            return "FLOAT64"
+        elif bigframes_dtype == bigframes.dtypes.STRING_DTYPE:
+            return "STRING"
+        elif bigframes_dtype == bigframes.dtypes.BOOL_DTYPE:
+            return "BOOLEAN"
+        elif bigframes_dtype == bigframes.dtypes.DATE_DTYPE:
+            return "DATE"
+        elif bigframes_dtype == bigframes.dtypes.TIME_DTYPE:
+            return "TIME"
+        elif bigframes_dtype == bigframes.dtypes.DATETIME_DTYPE:
+            return "DATETIME"
+        elif bigframes_dtype == bigframes.dtypes.TIMESTAMP_DTYPE:
+            return "TIMESTAMP"
+        elif bigframes_dtype == bigframes.dtypes.BYTES_DTYPE:
+            return "BYTES"
+        elif bigframes_dtype == bigframes.dtypes.NUMERIC_DTYPE:
+            return "NUMERIC"
+        elif bigframes_dtype == bigframes.dtypes.BIGNUMERIC_DTYPE:
+            return "BIGNUMERIC"
+        elif bigframes_dtype == bigframes.dtypes.JSON_DTYPE:
+            return "JSON"
+        elif bigframes_dtype == bigframes.dtypes.GEO_DTYPE:
+            return "GEOGRAPHY"
+        elif isinstance(bigframes_dtype, pd.ArrowDtype):
+            if pa.types.is_list(bigframes_dtype.pyarrow_dtype):
+                inner_bigframes_dtype = bigframes.dtypes.arrow_dtype_to_bigframes_dtype(
+                    bigframes_dtype.pyarrow_dtype.value_type
+                )
+                return (
+                    f"ARRAY<{SQLGlotType.from_bigframes_dtype(inner_bigframes_dtype)}>"
+                )
+            elif pa.types.is_struct(bigframes_dtype.pyarrow_dtype):
+                struct_type = typing.cast(pa.StructType, bigframes_dtype.pyarrow_dtype)
+                inner_fields: list[str] = []
+                for i in range(struct_type.num_fields):
+                    field = struct_type.field(i)
+                    key = sg.to_identifier(field.name).sql("bigquery")
+                    dtype = SQLGlotType.from_bigframes_dtype(
+                        bigframes.dtypes.arrow_dtype_to_bigframes_dtype(field.type)
+                    )
+                    inner_fields.append(f"{key} {dtype}")
+                return "STRUCT<{}>".format(", ".join(inner_fields))
+
+        raise ValueError(
+            f"Unsupported type for {bigframes_dtype}. {constants.FEEDBACK_LINK}"
+        )

tests/unit/core/compile/sqlglot/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

tests/unit/core/compile/sqlglot/test_sqlglot_types.py

@@ -0,0 +1,64 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pandas as pd
+import pyarrow as pa
+
+import bigframes.core.compile.sqlglot.sqlglot_types as sgt
+import bigframes.dtypes as dtypes
+
+
+def test_from_bigframes_simple_dtypes():
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.INT_DTYPE) == "INT64"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.FLOAT_DTYPE) == "FLOAT64"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.STRING_DTYPE) == "STRING"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.BOOL_DTYPE) == "BOOLEAN"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.DATE_DTYPE) == "DATE"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.TIME_DTYPE) == "TIME"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.DATETIME_DTYPE) == "DATETIME"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.TIMESTAMP_DTYPE) == "TIMESTAMP"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.BYTES_DTYPE) == "BYTES"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.NUMERIC_DTYPE) == "NUMERIC"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.BIGNUMERIC_DTYPE) == "BIGNUMERIC"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.JSON_DTYPE) == "JSON"
+    assert sgt.SQLGlotType.from_bigframes_dtype(dtypes.GEO_DTYPE) == "GEOGRAPHY"
+
+
+def test_from_bigframes_struct_dtypes():
+    fields = [pa.field("int_col", pa.int64()), pa.field("bool_col", pa.bool_())]
+    struct_type = pd.ArrowDtype(pa.struct(fields))
+    expected = "STRUCT<int_col INT64, bool_col BOOLEAN>"
+    assert sgt.SQLGlotType.from_bigframes_dtype(struct_type) == expected
+
+
+def test_from_bigframes_array_dtypes():
+    int_array_type = pd.ArrowDtype(pa.list_(pa.int64()))
+    assert sgt.SQLGlotType.from_bigframes_dtype(int_array_type) == "ARRAY<INT64>"
+
+    string_array_type = pd.ArrowDtype(pa.list_(pa.string()))
+    assert sgt.SQLGlotType.from_bigframes_dtype(string_array_type) == "ARRAY<STRING>"
+
+
+def test_from_bigframes_multi_nested_dtypes():
+    fields = [
+        pa.field("string_col", pa.string()),
+        pa.field("date_col", pa.date32()),
+        pa.field("array_col", pa.list_(pa.timestamp("us"))),
+    ]
+    array_type = pd.ArrowDtype(pa.list_(pa.struct(fields)))
+
+    expected = (
+        "ARRAY<STRUCT<string_col STRING, date_col DATE, array_col ARRAY<DATETIME>>>"
+    )
+    assert sgt.SQLGlotType.from_bigframes_dtype(array_type) == expected