refactor: add sqlglot type conversions

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