compiled.py

# Copyright 2023 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 itertools
import typing
from typing import Literal, Optional, Sequence

import bigframes_vendored.ibis
import bigframes_vendored.ibis.backends.bigquery.backend as ibis_bigquery
import bigframes_vendored.ibis.common.deferred as ibis_deferred  # type: ignore
import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
import bigframes_vendored.ibis.expr.operations as ibis_ops
import bigframes_vendored.ibis.expr.types as ibis_types
from google.cloud import bigquery
import pyarrow as pa

from bigframes.core import agg_expressions, rewrite
import bigframes.core.agg_expressions as ex_types
import bigframes.core.compile.googlesql
import bigframes.core.compile.ibis_compiler.aggregate_compiler as agg_compiler
import bigframes.core.compile.ibis_compiler.scalar_op_compiler as op_compilers
import bigframes.core.compile.ibis_types
import bigframes.core.expression as ex
from bigframes.core.ordering import OrderingExpression
import bigframes.core.sql
from bigframes.core.window_spec import WindowSpec
import bigframes.dtypes

op_compiler = op_compilers.scalar_op_compiler


# Ibis Implementations
class UnorderedIR:
    def __init__(
        self,
        table: ibis_types.Table,
        columns: Sequence[ibis_types.Value],
    ):
        self._table = table
        # Allow creating a DataFrame directly from an Ibis table expression.
        # TODO(swast): Validate that each column references the same table (or
        # no table for literal values).
        self._columns = tuple(
            column.resolve(table)  # type:ignore
            # TODO(https://github.com/ibis-project/ibis/issues/7613): use
            # public API to refer to Deferred type.
            if isinstance(column, ibis_deferred.Deferred) else column
            for column in columns
        )
        # To allow for more efficient lookup by column name, create a
        # dictionary mapping names to column values.
        self._column_names = {column.get_name(): column for column in self._columns}

    def to_sql(
        self,
        order_by: Sequence[OrderingExpression],
        limit: Optional[int],
        selections: tuple[tuple[ex.DerefOp, str], ...],
    ) -> str:
        ibis_table = self._to_ibis_expr()
        # This set of output transforms maybe should be its own output node??

        selection_strings = tuple((ref.id.sql, name) for ref, name in selections)

        names_preserved = tuple(name for _, name in selections) == tuple(
            self.column_ids
        )
        is_noop_selection = (
            all((i[0] == i[1] for i in selection_strings)) and names_preserved
        )

        if order_by or limit or not is_noop_selection:
            sql = ibis_bigquery.Backend().compile(ibis_table)
            sql = (
                bigframes.core.compile.googlesql.Select()
                .from_(sql)
                .select(selection_strings)
                .sql()
            )

            # Single row frames may not have any ordering columns
            if len(order_by) > 0:
                order_by_clause = bigframes.core.sql.ordering_clause(order_by)
                sql += f"\n{order_by_clause}"
            if limit is not None:
                if not isinstance(limit, int):
                    raise TypeError(f"Limit param: {limit} must be an int.")
                sql += f"\nLIMIT {limit}"
        else:
            sql = ibis_bigquery.Backend().compile(self._to_ibis_expr())
        return typing.cast(str, sql)

    @property
    def columns(self) -> tuple[ibis_types.Value, ...]:
        return self._columns

    @property
    def column_ids(self) -> typing.Sequence[str]:
        return tuple(self._column_names.keys())

    @property
    def _ibis_bindings(self) -> dict[str, ibis_types.Value]:
        return {col: self._get_ibis_column(col) for col in self.column_ids}

    def projection(
        self,
        expression_id_pairs: tuple[tuple[ex.Expression, str], ...],
    ) -> UnorderedIR:
        """Apply an expression to the ArrayValue and assign the output to a column."""
        cannot_inline = any(expr.expensive for expr, _ in expression_id_pairs)

        bindings = {col: self._get_ibis_column(col) for col in self.column_ids}
        new_values = [
            op_compiler.compile_expression(expression, bindings).name(id)
            for expression, id in expression_id_pairs
        ]
        result = UnorderedIR(self._table, (*self._columns, *new_values))
        if cannot_inline:
            return result._reproject_to_table()
        else:
            # Cheap ops can defer "SELECT" and inline into later ops
            return result

    def selection(
        self,
        input_output_pairs: tuple[tuple[ex.DerefOp, str], ...],
    ) -> UnorderedIR:
        """Apply an expression to the ArrayValue and assign the output to a column."""
        bindings = {col: self._get_ibis_column(col) for col in self.column_ids}
        values = [
            op_compiler.compile_expression(input, bindings).name(id)
            for input, id in input_output_pairs
        ]
        return UnorderedIR(self._table, tuple(values))

    def _get_ibis_column(self, key: str) -> ibis_types.Value:
        """Gets the Ibis expression for a given column."""
        if key not in self.column_ids:
            raise ValueError(
                "Column name {} not in set of values: {}".format(key, self.column_ids)
            )
        return typing.cast(ibis_types.Value, self._column_names[key])

    def get_column_type(self, key: str) -> bigframes.dtypes.Dtype:
        ibis_type = typing.cast(
            bigframes.core.compile.ibis_types.IbisDtype,
            self._get_ibis_column(key).type(),
        )
        return typing.cast(
            bigframes.dtypes.Dtype,
            bigframes.core.compile.ibis_types.ibis_dtype_to_bigframes_dtype(ibis_type),
        )

    def _to_ibis_expr(
        self,
        *,
        fraction: Optional[float] = None,
    ):
        """
        Creates an Ibis table expression representing the DataFrame.

        Args:
            expose_hidden_cols:
                If True, include the hidden ordering columns in the results.

        Returns:
            An ibis expression representing the data help by the ArrayValue object.
        """
        # Special case for empty tables, since we can't create an empty
        # projection.
        if not self._columns:
            return self._table.select([bigframes_vendored.ibis.literal(1)])

        table = self._table.select(self._columns)
        if fraction is not None:
            table = table.filter(
                bigframes_vendored.ibis.random() < ibis_types.literal(fraction)
            )
        return table

    def filter(self, predicate: ex.Expression) -> UnorderedIR:
        table = self._to_ibis_expr()
        condition = op_compiler.compile_expression(predicate, table)
        table = table.filter(condition)
        return UnorderedIR(
            table, tuple(table[column_name] for column_name in self._column_names)
        )

    def aggregate(
        self,
        aggregations: typing.Sequence[tuple[ex_types.Aggregation, str]],
        by_column_ids: typing.Sequence[ex.DerefOp] = (),
        order_by: typing.Sequence[OrderingExpression] = (),
    ) -> UnorderedIR:
        """
        Apply aggregations to the expression.
        Arguments:
            aggregations: input_column_id, operation, output_column_id tuples
            by_column_ids: column ids of the aggregation key, this is preserved through
              the transform
            dropna: whether null keys should be dropped
        Returns:
            OrderedIR: the grouping key is a unique-valued column and has ordering
              information.
        """
        table = self._to_ibis_expr()
        bindings = {col: table[col] for col in self.column_ids}
        stats = {
            col_out: agg_compiler.compile_aggregate(
                aggregate,
                bindings,
                order_by=op_compiler._convert_row_ordering_to_table_values(
                    table, order_by
                ),
            )
            for aggregate, col_out in aggregations
        }
        if by_column_ids:
            result = table.group_by((ref.id.sql for ref in by_column_ids)).aggregate(
                **stats
            )
            return UnorderedIR(
                result, columns=tuple(result[key] for key in result.columns)
            )
        else:
            result = table.aggregate(**stats)
            return UnorderedIR(
                result,
                columns=[result[col_id] for col_id in [*stats.keys()]],
            )

    def _uniform_sampling(self, fraction: float) -> UnorderedIR:
        """Sampling the table on given fraction.

        .. warning::
            The row numbers of result is non-deterministic, avoid to use.
        """
        table = self._to_ibis_expr(fraction=fraction)
        columns = [table[column_name] for column_name in self._column_names]
        return UnorderedIR(
            table,
            columns=columns,
        )

    ## Helpers
    def _reproject_to_table(self) -> UnorderedIR:
        """
        Internal operators that projects the internal representation into a
        new ibis table expression where each value column is a direct
        reference to a column in that table expression. Needed after
        some operations such as window operations that cannot be used
        recursively in projections.
        """
        table = self._to_ibis_expr()
        columns = [table[column_name] for column_name in self._column_names]
        return UnorderedIR(
            table,
            columns=columns,
        )

    @classmethod
    def from_polars(
        cls, pa_table: pa.Table, schema: Sequence[bigquery.SchemaField]
    ) -> UnorderedIR:
        """Builds an in-memory only (SQL only) expr from a pyarrow table."""
        import bigframes_vendored.ibis.backends.bigquery.datatypes as third_party_ibis_bqtypes

        # derive the ibis schema from the original pandas schema
        keys_memtable = bigframes_vendored.ibis.memtable(
            pa_table,
            schema=third_party_ibis_bqtypes.BigQuerySchema.to_ibis(list(schema)),
        )
        return cls(
            keys_memtable,
            columns=tuple(keys_memtable[key] for key in keys_memtable.columns),
        )

    def join(
        self: UnorderedIR,
        right: UnorderedIR,
        conditions: tuple[tuple[str, str], ...],
        type: Literal["inner", "outer", "left", "right", "cross"],
        *,
        join_nulls: bool = True,
    ) -> UnorderedIR:
        """Join two expressions by column equality.

        Arguments:
            left: Expression for left table to join.
            left_column_ids: Column IDs (not label) to join by.
            right: Expression for right table to join.
            right_column_ids: Column IDs (not label) to join by.
            how: The type of join to perform.
            join_nulls (bool):
                If True, will joins NULL keys to each other.
        Returns:
            The joined expression. The resulting columns will be, in order,
            first the coalesced join keys, then, all the left columns, and
            finally, all the right columns.
        """
        # Shouldn't need to select the column ids explicitly, but it seems that ibis has some
        # bug resolving column ids otherwise, potentially because of the "JoinChain" op
        left_table = self._to_ibis_expr().select(self.column_ids)
        right_table = right._to_ibis_expr().select(right.column_ids)

        join_conditions = [
            _join_condition(
                left_table[left_index], right_table[right_index], nullsafe=join_nulls
            )
            for left_index, right_index in conditions
        ]

        combined_table = bigframes_vendored.ibis.join(
            left_table,
            right_table,
            predicates=join_conditions,
            how=type,  # type: ignore
        )
        columns = [combined_table[col.get_name()] for col in self.columns] + [
            combined_table[col.get_name()] for col in right.columns
        ]
        return UnorderedIR(
            combined_table,
            columns=columns,
        )

    def isin_join(
        self: UnorderedIR,
        right: UnorderedIR,
        indicator_col: str,
        conditions: tuple[str, str],
        *,
        join_nulls: bool = True,
    ) -> UnorderedIR:
        """Join two expressions by column equality.

        Arguments:
            left: Expression for left table to join.
            right: Expression for right table to join.
            conditions: Id pairs to compare
        Returns:
            The joined expression.
        """
        left_table = self._to_ibis_expr()
        right_table = right._to_ibis_expr()
        if join_nulls:  # nullsafe isin join must actually use "exists" subquery
            new_column = (
                (
                    _join_condition(
                        left_table[conditions[0]],
                        right_table[conditions[1]],
                        nullsafe=True,
                    )
                )
                .any()
                .name(indicator_col)
            )

        else:  # Can do simpler "in" subquery
            new_column = (
                (left_table[conditions[0]])
                .isin((right_table[conditions[1]]))
                .name(indicator_col)
            )

        columns = tuple(
            itertools.chain(
                (left_table[col.get_name()] for col in self.columns), (new_column,)
            )
        )

        return UnorderedIR(
            left_table,
            columns=columns,
        )

    def project_window_op(
        self,
        expression: ex_types.Aggregation,
        window_spec: WindowSpec,
        output_name: str,
    ) -> UnorderedIR:
        """
        Creates a new expression based on this expression with unary operation applied to one column.
        column_name: the id of the input column present in the expression
        op: the windowable operator to apply to the input column
        window_spec: a specification of the window over which to apply the operator
        output_name: the id to assign to the output of the operator
        """
        # Cannot nest analytic expressions, so reproject to cte first if needed.
        # Also ibis cannot window literals, so need to reproject those (even though this is legal in googlesql)
        # See: https://github.com/ibis-project/ibis/issues/9773
        used_exprs = map(
            self._compile_expression,
            map(
                ex.DerefOp,
                itertools.chain(
                    expression.column_references, window_spec.all_referenced_columns
                ),
            ),
        )
        can_directly_window = not any(
            map(lambda x: is_literal(x) or is_window(x), used_exprs)
        )
        if not can_directly_window:
            return self._reproject_to_table().project_window_op(
                expression,
                window_spec,
                output_name,
            )

        rewritten_expr = rewrite.simplify_complex_windows(
            agg_expressions.WindowExpression(expression, window_spec)
        )

        ibis_expr = op_compiler.compile_expression(rewritten_expr, self._ibis_bindings)

        return UnorderedIR(self._table, (*self.columns, ibis_expr.name(output_name)))

    def _compile_expression(self, expr: ex.Expression):
        return op_compiler.compile_expression(expr, self._ibis_bindings)


def is_literal(column: ibis_types.Value) -> bool:
    # Unfortunately, Literals in ibis are not "Columns"s and therefore can't be aggregated.
    return not isinstance(column, ibis_types.Column)


def is_window(column: ibis_types.Value) -> bool:
    matches = (
        (column)
        .op()
        .find_topmost(
            lambda x: isinstance(x, (ibis_ops.WindowFunction, ibis_ops.Relation))
        )
    )
    return any(isinstance(op, ibis_ops.WindowFunction) for op in matches)


def _string_cast_join_cond(
    lvalue: ibis_types.Column, rvalue: ibis_types.Column
) -> ibis_types.BooleanColumn:
    result = (
        lvalue.cast(ibis_dtypes.str).fill_null(ibis_types.literal("0"))
        == rvalue.cast(ibis_dtypes.str).fill_null(ibis_types.literal("0"))
    ) & (
        lvalue.cast(ibis_dtypes.str).fill_null(ibis_types.literal("1"))
        == rvalue.cast(ibis_dtypes.str).fill_null(ibis_types.literal("1"))
    )
    return typing.cast(ibis_types.BooleanColumn, result)


def _numeric_join_cond(
    lvalue: ibis_types.Column, rvalue: ibis_types.Column
) -> ibis_types.BooleanColumn:
    lvalue1 = lvalue.fill_null(ibis_types.literal(0))
    lvalue2 = lvalue.fill_null(ibis_types.literal(1))
    rvalue1 = rvalue.fill_null(ibis_types.literal(0))
    rvalue2 = rvalue.fill_null(ibis_types.literal(1))
    if lvalue.type().is_floating() and rvalue.type().is_floating():
        # NaN aren't equal so need to coalesce as well with diff constants
        lvalue1 = (
            typing.cast(ibis_types.FloatingColumn, lvalue)
            .isnan()
            .ifelse(ibis_types.literal(2), lvalue1)
        )
        lvalue2 = (
            typing.cast(ibis_types.FloatingColumn, lvalue)
            .isnan()
            .ifelse(ibis_types.literal(3), lvalue2)
        )
        rvalue1 = (
            typing.cast(ibis_types.FloatingColumn, rvalue)
            .isnan()
            .ifelse(ibis_types.literal(2), rvalue1)
        )
        rvalue2 = (
            typing.cast(ibis_types.FloatingColumn, rvalue)
            .isnan()
            .ifelse(ibis_types.literal(3), rvalue2)
        )
    result = (lvalue1 == rvalue1) & (lvalue2 == rvalue2)
    return typing.cast(ibis_types.BooleanColumn, result)


def _join_condition(
    lvalue: ibis_types.Column, rvalue: ibis_types.Column, nullsafe: bool
) -> ibis_types.BooleanColumn:
    if (lvalue.type().is_floating()) and (lvalue.type().is_floating()):
        # Need to always make safe join condition to handle nan, even if no nulls
        return _numeric_join_cond(lvalue, rvalue)
    if nullsafe:
        # TODO: Define more coalesce constants for non-numeric types to avoid cast
        if (lvalue.type().is_numeric()) and (lvalue.type().is_numeric()):
            return _numeric_join_cond(lvalue, rvalue)
        else:
            return _string_cast_join_cond(lvalue, rvalue)
    return typing.cast(ibis_types.BooleanColumn, lvalue == rvalue)