refactor: Move ibis type mappings to compile module (#745)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -60,6 +60,7 @@ class ArrayValue:
 
     node: nodes.BigFrameNode
 
+    # DO NOT use, on deprecation path
     @classmethod
     def from_ibis(
         cls,
@@ -69,11 +70,13 @@ def from_ibis(
         hidden_ordering_columns: Sequence[ibis_types.Value],
         ordering: orderings.ExpressionOrdering,
     ):
+        import bigframes.core.compile.ibis_types
+
         node = nodes.ReadGbqNode(
             table=table,
             table_session=session,
             columns=tuple(
-                bigframes.dtypes.ibis_value_to_canonical_type(column)
+                bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(column)
                 for column in columns
             ),
             hidden_ordering_columns=tuple(hidden_ordering_columns),

bigframes/core/compile/aggregate_compiler.py

@@ -22,10 +22,10 @@
 import pandas as pd
 
 import bigframes.constants as constants
+import bigframes.core.compile.ibis_types as compile_ibis_types
 import bigframes.core.compile.scalar_op_compiler as scalar_compilers
 import bigframes.core.expression as ex
 import bigframes.core.window_spec as window_spec
-import bigframes.dtypes as dtypes
 import bigframes.operations.aggregations as agg_ops
 
 scalar_compiler = scalar_compilers.scalar_op_compiler
@@ -323,7 +323,7 @@ def _(
             for this_bin in range(op.bins - 1):
                 out = out.when(
                     x <= (col_min + (this_bin + 1) * bin_width),
-                    dtypes.literal_to_ibis_scalar(
+                    compile_ibis_types.literal_to_ibis_scalar(
                         this_bin, force_dtype=pd.Int64Dtype()
                     ),
                 )
@@ -352,8 +352,8 @@ def _(
                     out = out.when(x.notnull(), interval_struct)
     else:  # Interpret as intervals
         for interval in op.bins:
-            left = dtypes.literal_to_ibis_scalar(interval[0])
-            right = dtypes.literal_to_ibis_scalar(interval[1])
+            left = compile_ibis_types.literal_to_ibis_scalar(interval[0])
+            right = compile_ibis_types.literal_to_ibis_scalar(interval[1])
             condition = (x > left) & (x <= right)
             interval_struct = ibis.struct(
                 {"left_exclusive": left, "right_inclusive": right}
@@ -370,7 +370,7 @@ def _(
     window=None,
 ) -> ibis_types.IntegerValue:
     if isinstance(self.quantiles, int):
-        quantiles_ibis = dtypes.literal_to_ibis_scalar(self.quantiles)
+        quantiles_ibis = compile_ibis_types.literal_to_ibis_scalar(self.quantiles)
         percent_ranks = cast(
             ibis_types.FloatingColumn,
             _apply_window_if_present(column.percent_rank(), window),
@@ -383,13 +383,19 @@ def _(
             _apply_window_if_present(column.percent_rank(), window),
         )
         out = ibis.case()
-        first_ibis_quantile = dtypes.literal_to_ibis_scalar(self.quantiles[0])
+        first_ibis_quantile = compile_ibis_types.literal_to_ibis_scalar(
+            self.quantiles[0]
+        )
         out = out.when(percent_ranks < first_ibis_quantile, None)
         for bucket_n in range(len(self.quantiles) - 1):
-            ibis_quantile = dtypes.literal_to_ibis_scalar(self.quantiles[bucket_n + 1])
+            ibis_quantile = compile_ibis_types.literal_to_ibis_scalar(
+                self.quantiles[bucket_n + 1]
+            )
             out = out.when(
                 percent_ranks <= ibis_quantile,
-                dtypes.literal_to_ibis_scalar(bucket_n, force_dtype=pd.Int64Dtype()),
+                compile_ibis_types.literal_to_ibis_scalar(
+                    bucket_n, force_dtype=pd.Int64Dtype()
+                ),
             )
         out = out.else_(None)
         return out.end()  # type: ignore

bigframes/core/compile/compiled.py

@@ -28,6 +28,7 @@
 import pandas
 
 import bigframes.core.compile.aggregate_compiler as agg_compiler
+import bigframes.core.compile.ibis_types
 import bigframes.core.compile.scalar_op_compiler as op_compilers
 import bigframes.core.expression as ex
 import bigframes.core.guid
@@ -157,16 +158,19 @@ def _get_ibis_column(self, key: str) -> ibis_types.Value:
             )
         return typing.cast(
             ibis_types.Value,
-            bigframes.dtypes.ibis_value_to_canonical_type(self._column_names[key]),
+            bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                self._column_names[key]
+            ),
         )
 
     def get_column_type(self, key: str) -> bigframes.dtypes.Dtype:
         ibis_type = typing.cast(
-            bigframes.dtypes.IbisDtype, self._get_ibis_column(key).type()
+            bigframes.core.compile.ibis_types.IbisDtype,
+            self._get_ibis_column(key).type(),
         )
         return typing.cast(
             bigframes.dtypes.Dtype,
-            bigframes.dtypes.ibis_dtype_to_bigframes_dtype(ibis_type),
+            bigframes.core.compile.ibis_types.ibis_dtype_to_bigframes_dtype(ibis_type),
         )
 
     def _aggregate_base(
@@ -332,7 +336,8 @@ def _to_ibis_expr(
         # Make sure all dtypes are the "canonical" ones for BigFrames. This is
         # important for operations like UNION where the schema must match.
         table = self._table.select(
-            bigframes.dtypes.ibis_value_to_canonical_type(column) for column in columns
+            bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(column)
+            for column in columns
         )
         base_table = table
         if self._reduced_predicate is not None:
@@ -579,7 +584,10 @@ def from_pandas(
         ibis_values = ibis_values.assign(**{ORDER_ID_COLUMN: range(len(pd_df))})
         # derive the ibis schema from the original pandas schema
         ibis_schema = [
-            (name, bigframes.dtypes.bigframes_dtype_to_ibis_dtype(dtype))
+            (
+                name,
+                bigframes.core.compile.ibis_types.bigframes_dtype_to_ibis_dtype(dtype),
+            )
             for name, dtype in zip(schema.names, schema.dtypes)
         ]
         ibis_schema.append((ORDER_ID_COLUMN, ibis_dtypes.int64))
@@ -993,7 +1001,9 @@ def _to_ibis_expr(
         # Make sure all dtypes are the "canonical" ones for BigFrames. This is
         # important for operations like UNION where the schema must match.
         table = table.select(
-            bigframes.dtypes.ibis_value_to_canonical_type(table[column])
+            bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                table[column]
+            )
             for column in table.columns
         )
         base_table = table

bigframes/core/compile/compiler.py

@@ -26,11 +26,11 @@
 import bigframes.core.compile.compiled as compiled
 import bigframes.core.compile.concat as concat_impl
 import bigframes.core.compile.default_ordering as default_ordering
+import bigframes.core.compile.ibis_types
 import bigframes.core.compile.schema_translator
 import bigframes.core.compile.single_column
 import bigframes.core.nodes as nodes
 import bigframes.core.ordering as bf_ordering
-import bigframes.dtypes as bigframes_dtypes
 
 if typing.TYPE_CHECKING:
     import bigframes.core
@@ -112,7 +112,9 @@ def compile_cached_table(node: nodes.CachedTableNode, ordered: bool = True):
         return compiled.OrderedIR(
             ibis_table,
             columns=tuple(
-                bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col])
+                bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                    ibis_table[col]
+                )
                 for col in node.schema.names
             ),
             ordering=node.ordering,
@@ -123,7 +125,9 @@ def compile_cached_table(node: nodes.CachedTableNode, ordered: bool = True):
         return compiled.UnorderedIR(
             ibis_table,
             columns=tuple(
-                bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col])
+                bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                    ibis_table[col]
+                )
                 for col in node.schema.names
             ),
         )
@@ -166,7 +170,9 @@ def compile_read_table_unordered(node: nodes.ReadTableNode):
     return compiled.UnorderedIR(
         ibis_table,
         tuple(
-            bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col])
+            bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                ibis_table[col]
+            )
             for col in node.schema.names
         ),
     )
@@ -202,7 +208,9 @@ def compile_read_table_ordered(node: nodes.ReadTableNode):
     return compiled.OrderedIR(
         ibis_table,
         columns=tuple(
-            bigframes_dtypes.ibis_value_to_canonical_type(ibis_table[col])
+            bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(
+                ibis_table[col]
+            )
             for col in node.schema.names
         ),
         ordering=ordering,