blocks.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.

"""Block is a 2D data structure that supports data mutability and views.

These data structures are shared by DataFrame and Series. This allows views to
link in both directions (DataFrame to Series and vice versa) and prevents
circular dependencies.
"""

from __future__ import annotations

import ast
import dataclasses
import datetime
import functools
import itertools
import random
import typing
from typing import (
    Iterable,
    Iterator,
    List,
    Literal,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Union,
)
import warnings

import bigframes_vendored.constants as constants
import google.cloud.bigquery as bigquery
import numpy
import pandas as pd
import pyarrow as pa

from bigframes import session
from bigframes._config import sampling_options
import bigframes.constants
from bigframes.core import agg_expressions, local_data
import bigframes.core as core
import bigframes.core.agg_expressions as ex_types
import bigframes.core.expression as ex
import bigframes.core.expression as scalars
import bigframes.core.guid as guid
import bigframes.core.identifiers
import bigframes.core.join_def as join_defs
import bigframes.core.ordering as ordering
import bigframes.core.pyarrow_utils as pyarrow_utils
import bigframes.core.utils as utils
import bigframes.core.window_spec as windows
import bigframes.dtypes
import bigframes.exceptions as bfe
import bigframes.operations as ops
import bigframes.operations.aggregations as agg_ops
from bigframes.session import dry_runs, execution_spec
from bigframes.session import executor as executors
from bigframes.session._io import pandas as io_pandas

# Type constraint for wherever column labels are used
Label = typing.Hashable

# Bytes to Megabyte Conversion
_BYTES_TO_KILOBYTES = 1024
_BYTES_TO_MEGABYTES = _BYTES_TO_KILOBYTES * 1024

# This is the max limit of physical columns in BQ
# May choose to set smaller limit for number of block columns to allow overhead for ordering, etc.
_BQ_MAX_COLUMNS = 10000

# All sampling method
_HEAD = "head"
_UNIFORM = "uniform"
_SAMPLING_METHODS = (_HEAD, _UNIFORM)

# Monotonic Cache Names
_MONOTONIC_INCREASING = "monotonic_increasing"
_MONOTONIC_DECREASING = "monotonic_decreasing"


LevelType = typing.Hashable
LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]


class PandasBatches(Iterator[pd.DataFrame]):
    """Interface for mutable objects with state represented by a block value object."""

    def __init__(
        self,
        pandas_batches: Iterator[pd.DataFrame],
        total_rows: Optional[int] = 0,
        *,
        total_bytes_processed: Optional[int] = 0,
    ):
        self._dataframes: Iterator[pd.DataFrame] = pandas_batches
        self._total_rows: Optional[int] = total_rows
        self._total_bytes_processed: Optional[int] = total_bytes_processed

    @property
    def total_rows(self) -> Optional[int]:
        return self._total_rows

    @property
    def total_bytes_processed(self) -> Optional[int]:
        return self._total_bytes_processed

    def __next__(self) -> pd.DataFrame:
        return next(self._dataframes)


@dataclasses.dataclass()
class MaterializationOptions:
    downsampling: sampling_options.SamplingOptions = dataclasses.field(
        default_factory=sampling_options.SamplingOptions
    )
    allow_large_results: Optional[bool] = None
    ordered: bool = True


class Block:
    """A immutable 2D data structure."""

    def __init__(
        self,
        expr: core.ArrayValue,
        index_columns: Iterable[str],
        column_labels: typing.Union[pd.Index, typing.Iterable[Label]],
        index_labels: typing.Union[pd.Index, typing.Iterable[Label], None] = None,
        *,
        transpose_cache: Optional[Block] = None,
    ):
        """Construct a block object, will create default index if no index columns specified."""
        index_columns = list(index_columns)
        if index_labels is not None:
            index_labels = list(index_labels)
            if len(index_labels) != len(index_columns):
                raise ValueError(
                    f"'index_columns' (size {len(index_columns)}) and 'index_labels' (size {len(index_labels)}) must have equal length"
                )

        self._index_columns = tuple(index_columns)
        # Index labels don't need complicated hierarchical access so can store as tuple
        self._index_labels = (
            tuple(index_labels)
            if index_labels
            else tuple([None for _ in index_columns])
        )
        self._expr = self._normalize_expression(expr, self._index_columns)
        # Use pandas index to more easily replicate column indexing, especially for hierarchical column index
        self._column_labels = (
            column_labels.copy()
            if isinstance(column_labels, pd.Index)
            else pd.Index(column_labels)
        )
        if len(self.value_columns) != len(self._column_labels):
            raise ValueError(
                f"'value_columns' (size {len(self.value_columns)}) and 'column_labels' (size {len(self._column_labels)}) must have equal length"
            )
        # col_id -> [stat_name -> scalar]
        # TODO: Preserve cache under safe transforms (eg. drop column, reorder)
        self._stats_cache: dict[str, dict[str, typing.Any]] = {
            col_id: {} for col_id in self.value_columns
        }
        # TODO(kemppeterson) Add a cache for corr to parallel the single-column stats.

        self._stats_cache[" ".join(self.index_columns)] = {}
        self._transpose_cache: Optional[Block] = transpose_cache
        self._view_ref: Optional[bigquery.TableReference] = None
        self._view_ref_dry_run: Optional[bigquery.TableReference] = None

    @classmethod
    def from_pyarrow(
        cls,
        data: pa.Table,
        session: bigframes.Session,
    ) -> Block:
        column_labels = data.column_names

        # TODO(tswast): Use array_value.promote_offsets() instead once that node is
        # supported by the local engine.
        offsets_col = bigframes.core.guid.generate_guid()
        index_ids = [offsets_col]
        index_labels = [None]

        # TODO(https://github.com/googleapis/python-bigquery-dataframes/issues/859):
        # Allow users to specify the "total ordering" column(s) or allow multiple
        # such columns.
        data = pyarrow_utils.append_offsets(data, offsets_col=offsets_col)

        # from_pyarrow will normalize the types for us.
        managed_data = local_data.ManagedArrowTable.from_pyarrow(data)
        array_value = core.ArrayValue.from_managed(managed_data, session=session)
        block = cls(
            array_value,
            column_labels=column_labels,
            index_columns=index_ids,
            index_labels=index_labels,
        )
        return block

    @classmethod
    def from_local(
        cls,
        data: pd.DataFrame,
        session: bigframes.Session,
        *,
        cache_transpose: bool = True,
    ) -> Block:
        # Assumes caller has already converted datatypes to bigframes ones.
        pd_data = data
        column_labels = pd_data.columns
        index_labels = list(pd_data.index.names)

        # unique internal ids
        column_ids = [f"column_{i}" for i in range(len(pd_data.columns))]
        index_ids = [f"level_{level}" for level in range(pd_data.index.nlevels)]

        pd_data = pd_data.set_axis(column_ids, axis=1)
        pd_data = pd_data.reset_index(names=index_ids)
        managed_data = local_data.ManagedArrowTable.from_pandas(pd_data)
        array_value = core.ArrayValue.from_managed(managed_data, session=session)
        block = cls(
            array_value,
            column_labels=column_labels,
            index_columns=index_ids,
            index_labels=index_labels,
        )
        if cache_transpose:
            try:
                # this cache will help when aligning on axis=1
                block = block.with_transpose_cache(
                    cls.from_local(data.T, session, cache_transpose=False)
                )
            except Exception:
                pass
        return block

    @property
    def has_index(self) -> bool:
        return len(self._index_columns) > 0

    @property
    def index(self) -> BlockIndexProperties:
        """Row identities for values in the Block."""
        return BlockIndexProperties(self)

    @functools.cached_property
    def shape(self) -> typing.Tuple[int, int]:
        """Returns dimensions as (length, width) tuple."""
        # Support zero-query for hermetic unit tests.
        if self.expr.session is None and self.expr.node.row_count:
            try:
                return self.expr.node.row_count
            except Exception:
                pass

        row_count = (
            self.session._executor.execute(
                self.expr.row_count(),
                execution_spec.ExecutionSpec(promise_under_10gb=True, ordered=False),
            )
            .batches()
            .to_py_scalar()
        )
        return (row_count, len(self.value_columns))

    @property
    def index_columns(self) -> Sequence[str]:
        """Column(s) to use as row labels."""
        return self._index_columns

    @property
    def value_columns(self) -> Sequence[str]:
        """All value columns, mutually exclusive with index columns."""
        return [
            column
            for column in self._expr.column_ids
            if column not in self.index_columns
        ]

    @property
    def column_labels(self) -> pd.Index:
        return self._column_labels

    @property
    def expr(self) -> core.ArrayValue:
        """Expression representing all columns, including index columns."""
        return self._expr

    @property
    def dtypes(
        self,
    ) -> Sequence[bigframes.dtypes.Dtype]:
        """Returns the dtypes of the value columns."""
        return [self.expr.get_column_type(col) for col in self.value_columns]

    @property
    def session(self) -> session.Session:
        return self._expr.session

    @functools.cached_property
    def col_id_to_label(self) -> typing.Mapping[str, Label]:
        """Get column label for value columns, or index name for index columns"""
        return {
            col_id: label
            for col_id, label in zip(self.value_columns, self._column_labels)
        }

    @functools.cached_property
    def label_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]:
        """Get column label for value columns, or index name for index columns"""
        mapping: typing.Dict[Label, typing.Sequence[str]] = {}
        for id, label in self.col_id_to_label.items():
            mapping[label] = (*mapping.get(label, ()), id)
        return mapping

    def resolve_label_exact(self, label: Label) -> Optional[str]:
        """Returns the column id matching the label if there is exactly
        one such column. If there are multiple columns with the same name,
        raises an error. If there is no such a column, returns None."""
        matches = self.label_to_col_id.get(label, [])
        if len(matches) > 1:
            raise ValueError(
                f"Multiple columns matching id {label} were found. {constants.FEEDBACK_LINK}"
            )
        return matches[0] if len(matches) != 0 else None

    def resolve_label_exact_or_error(self, label: Label) -> str:
        """Returns the column id matching the label if there is exactly
        one such column. If there are multiple columns with the same name,
        raises an error. If there is no such a column, raises an error too."""
        col_id = self.resolve_label_exact(label)
        if col_id is None:
            raise ValueError(f"Label {label} not found. {constants.FEEDBACK_LINK}")
        return col_id

    @functools.cached_property
    def col_id_to_index_name(self) -> typing.Mapping[str, Label]:
        """Get column label for value columns, or index name for index columns"""
        return {
            col_id: label
            for col_id, label in zip(self.index_columns, self._index_labels)
        }

    @functools.cached_property
    def index_name_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]:
        """Get column label for value columns, or index name for index columns"""
        mapping: typing.Dict[Label, typing.Sequence[str]] = {}
        for id, label in self.col_id_to_index_name.items():
            mapping[label] = (*mapping.get(label, ()), id)
        return mapping

    @property
    def explicitly_ordered(self) -> bool:
        return self.expr.explicitly_ordered

    def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]:
        """
        Unlike label_to_col_id, this works with partial labels for multi-index.

        Only some methods, like __getitem__ can use a partial key to get columns
        from a dataframe. These methods should use cols_matching_label, while
        methods that require exact label matches should use label_to_col_id.
        """
        # TODO(tbergeron): Refactor so that all label lookups use this method
        if partial_label not in self.column_labels:
            return []
        loc = self.column_labels.get_loc(partial_label)
        if isinstance(loc, int):
            return [self.value_columns[loc]]
        if isinstance(loc, slice):
            return self.value_columns[loc]
        return [col for col, is_present in zip(self.value_columns, loc) if is_present]

    def order_by(
        self,
        by: typing.Sequence[ordering.OrderingExpression],
    ) -> Block:
        return Block(
            self._expr.order_by(by),
            index_columns=self.index_columns,
            column_labels=self.column_labels,
            index_labels=self.index.names,
        )

    def reversed(self) -> Block:
        return Block(
            self._expr.reversed(),
            index_columns=self.index_columns,
            column_labels=self.column_labels,
            index_labels=self.index.names,
        )

    def reset_index(
        self,
        level: LevelsType = None,
        drop: bool = True,
        *,
        col_level: Union[str, int] = 0,
        col_fill: typing.Hashable = "",
        allow_duplicates: bool = False,
        replacement: Optional[bigframes.enums.DefaultIndexKind] = None,
    ) -> Block:
        """Reset the index of the block, promoting the old index to a value column.

        Arguments:
            level: the label or index level of the index levels to remove.
            name: this is the column id for the new value id derived from the old index
            allow_duplicates: if false, duplicate col labels will result in error
            replacement: if not null, will override default index replacement type

        Returns:
            A new Block because dropping index columns can break references
            from Index classes that point to this block.
        """
        if level is not None:
            # preserve original order, not user provided order
            level_ids: Sequence[str] = [
                id for id in self.index_columns if id in self.index.resolve_level(level)
            ]
        else:
            level_ids = self.index_columns

        expr = self._expr
        replacement_idx_type = replacement or self.session._default_index_type
        if set(self.index_columns) > set(level_ids):
            new_index_cols = [col for col in self.index_columns if col not in level_ids]
            new_index_labels = [self.col_id_to_index_name[id] for id in new_index_cols]
        elif replacement_idx_type == bigframes.enums.DefaultIndexKind.SEQUENTIAL_INT64:
            expr, new_index_col_id = expr.promote_offsets()
            new_index_cols = [new_index_col_id]
            new_index_labels = [None]
        elif replacement_idx_type == bigframes.enums.DefaultIndexKind.NULL:
            new_index_cols = []
            new_index_labels = []
        else:
            raise ValueError(f"Unrecognized default index kind: {replacement_idx_type}")

        if drop:
            # Even though the index might be part of the ordering, keep that
            # ordering expression as reset_index shouldn't change the row
            # order.
            expr = expr.drop_columns(level_ids)
            return Block(
                expr,
                index_columns=new_index_cols,
                index_labels=new_index_labels,
                column_labels=self.column_labels,
            )
        else:
            # Add index names to column index
            col_level_n = (
                col_level
                if isinstance(col_level, int)
                else self.column_labels.names.index(col_level)
            )
            column_labels_modified = self.column_labels
            for position, level_id in enumerate(level_ids):
                label = self.col_id_to_index_name[level_id]
                if label is None:
                    if "index" not in self.column_labels and self.index.nlevels <= 1:
                        label = "index"
                    else:
                        label = f"level_{self.index_columns.index(level_id)}"

                if (not allow_duplicates) and (label in self.column_labels):
                    raise ValueError(f"cannot insert {label}, already exists")

                if isinstance(self.column_labels, pd.MultiIndex):
                    nlevels = self.column_labels.nlevels
                    label = tuple(
                        label if i == col_level_n else col_fill for i in range(nlevels)
                    )

                # Create index copy with label inserted
                # See: https://pandas.pydata.org/docs/reference/api/pandas.Index.insert.html
                column_labels_modified = column_labels_modified.insert(position, label)

            return Block(
                expr.select_columns((*new_index_cols, *level_ids, *self.value_columns)),
                index_columns=new_index_cols,
                index_labels=new_index_labels,
                column_labels=column_labels_modified,
            )

    def set_index(
        self,
        col_ids: typing.Sequence[str],
        drop: bool = True,
        append: bool = False,
        index_labels: typing.Sequence[Label] = (),
    ) -> Block:
        """Set the index of the block to

        Arguments:
            ids: columns to be converted to index columns
            drop: whether to drop the new index columns as value columns
            append: whether to discard the existing index or add on to it
            index_labels: new index labels

        Returns:
            Block with new index
        """
        expr = self._expr

        new_index_columns = []
        new_index_labels = []
        for col_id in col_ids:
            col_copy_id = guid.generate_guid()
            expr = expr.assign(col_id, col_copy_id)
            new_index_columns.append(col_copy_id)
            new_index_labels.append(self.col_id_to_label[col_id])

        if append:
            new_index_columns = [*self.index_columns, *new_index_columns]
            new_index_labels = [*self._index_labels, *new_index_labels]
        else:
            expr = expr.drop_columns(self.index_columns)

        if index_labels:
            new_index_labels = list(index_labels)

        block = Block(
            expr,
            index_columns=new_index_columns,
            column_labels=self.column_labels,
            index_labels=new_index_labels,
        )
        if drop:
            # These are the value columns, new index uses the copies, so this is safe
            block = block.drop_columns(col_ids)
        return block

    def drop_levels(self, ids: typing.Sequence[str]):
        for id in ids:
            if id not in self.index_columns:
                raise ValueError(f"{id} is not an index column")
        expr = self._expr.drop_columns(ids)
        remaining_index_col_ids = [
            col_id for col_id in self.index_columns if col_id not in ids
        ]
        if len(remaining_index_col_ids) == 0:
            raise ValueError("Cannot drop all index levels, at least 1 must remain.")
        level_names = [
            self.col_id_to_index_name[index_id] for index_id in remaining_index_col_ids
        ]
        return Block(expr, remaining_index_col_ids, self.column_labels, level_names)

    def reorder_levels(self, ids: typing.Sequence[str]):
        if sorted(self.index_columns) != sorted(ids):
            raise ValueError("Cannot drop or duplicate levels using reorder_levels.")
        level_names = [self.col_id_to_index_name[index_id] for index_id in ids]
        return Block(self.expr, ids, self.column_labels, level_names)

    def to_arrow(
        self,
        *,
        ordered: bool = True,
        allow_large_results: Optional[bool] = None,
    ) -> Tuple[pa.Table, Optional[bigquery.QueryJob]]:
        """Run query and download results as a pyarrow Table."""
        under_10gb = (
            (not allow_large_results)
            if (allow_large_results is not None)
            else not bigframes.options._allow_large_results
        )
        execute_result = self.session._executor.execute(
            self.expr,
            execution_spec.ExecutionSpec(
                promise_under_10gb=under_10gb,
                ordered=ordered,
            ),
        )
        pa_table = execute_result.batches().to_arrow_table()

        pa_index_labels = []
        for index_level, index_label in enumerate(self._index_labels):
            if isinstance(index_label, str):
                pa_index_labels.append(index_label)
            else:
                pa_index_labels.append(f"__index_level_{index_level}__")

        # pa.Table.from_pandas puts index columns last, so update to match.
        pa_table = pa_table.select([*self.value_columns, *self.index_columns])
        pa_table = pa_table.rename_columns(list(self.column_labels) + pa_index_labels)
        return pa_table, execute_result.query_job

    def to_pandas(
        self,
        max_download_size: Optional[int] = None,
        sampling_method: Optional[str] = None,
        random_state: Optional[int] = None,
        *,
        ordered: bool = True,
        allow_large_results: Optional[bool] = None,
    ) -> Tuple[pd.DataFrame, Optional[bigquery.QueryJob]]:
        """Run query and download results as a pandas DataFrame.

        Args:
            max_download_size (int, default None):
                Download size threshold in MB. If max_download_size is exceeded when downloading data
                (e.g., to_pandas()), the data will be downsampled if
                bigframes.options.sampling.enable_downsampling is True, otherwise, an error will be
                raised. If set to a value other than None, this will supersede the global config.
            sampling_method (str, default None):
                Downsampling algorithms to be chosen from, the choices are: "head": This algorithm
                returns a portion of the data from the beginning. It is fast and requires minimal
                computations to perform the downsampling; "uniform": This algorithm returns uniform
                random samples of the data. If set to a value other than None, this will supersede
                the global config.
            random_state (int, default None):
                The seed for the uniform downsampling algorithm. If provided, the uniform method may
                take longer to execute and require more computation. If set to a value other than
                None, this will supersede the global config.
            ordered (bool, default True):
                Determines whether the resulting pandas dataframe will be ordered.
                Whether the row ordering is deterministics depends on whether session ordering is strict.

        Returns:
            pandas.DataFrame, QueryJob
        """
        sampling = self._get_sampling_option(
            max_download_size, sampling_method, random_state
        )

        return self._materialize_local(
            materialize_options=MaterializationOptions(
                downsampling=sampling,
                allow_large_results=allow_large_results,
                ordered=ordered,
            )
        )

    def _get_sampling_option(
        self,
        max_download_size: Optional[int] = None,
        sampling_method: Optional[str] = None,
        random_state: Optional[int] = None,
    ) -> sampling_options.SamplingOptions:

        if (sampling_method is not None) and (sampling_method not in _SAMPLING_METHODS):
            raise NotImplementedError(
                f"The downsampling method {sampling_method} is not implemented, "
                f"please choose from {','.join(_SAMPLING_METHODS)}."
            )

        sampling = bigframes.options.sampling.with_max_download_size(max_download_size)
        if sampling_method is None:
            return sampling.with_disabled()

        return sampling.with_method(sampling_method).with_random_state(  # type: ignore
            random_state
        )

    def try_peek(
        self, n: int = 20, force: bool = False, allow_large_results=None
    ) -> typing.Optional[pd.DataFrame]:
        if force or self.expr.supports_fast_peek:
            # really, we should just block insane peek values and always assume <10gb
            under_10gb = (
                (not allow_large_results)
                if (allow_large_results is not None)
                else not bigframes.options._allow_large_results
            )
            result = self.session._executor.execute(
                self.expr,
                execution_spec.ExecutionSpec(promise_under_10gb=under_10gb, peek=n),
            )
            df = result.batches().to_pandas()
            return self._copy_index_to_pandas(df)
        else:
            return None

    def to_pandas_batches(
        self,
        page_size: Optional[int] = None,
        max_results: Optional[int] = None,
        allow_large_results: Optional[bool] = None,
    ) -> PandasBatches:
        """Download results one message at a time.

        page_size and max_results determine the size and number of batches,
        see https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob#google_cloud_bigquery_job_QueryJob_result"""

        under_10gb = (
            (not allow_large_results)
            if (allow_large_results is not None)
            else not bigframes.options._allow_large_results
        )
        execution_result = self.session._executor.execute(
            self.expr,
            execution_spec.ExecutionSpec(
                promise_under_10gb=under_10gb,
                ordered=True,
            ),
        )
        result_batches = execution_result.batches()

        # To reduce the number of edge cases to consider when working with the
        # results of this, always return at least one DataFrame. See:
        # b/428918844.
        try:
            empty_arrow_table = self.expr.schema.to_pyarrow().empty_table()
        except pa.ArrowNotImplementedError:
            # Bug with some pyarrow versions(https://github.com/apache/arrow/issues/45262),
            # empty_table only supports base storage types, not extension types.
            empty_arrow_table = self.expr.schema.to_pyarrow(
                use_storage_types=True
            ).empty_table()
        empty_val = io_pandas.arrow_to_pandas(empty_arrow_table, self.expr.schema)
        dfs = map(
            lambda a: a[0],
            itertools.zip_longest(
                result_batches.to_pandas_batches(page_size, max_results),
                [0],
                fillvalue=empty_val,
            ),
        )
        dfs = iter(map(self._copy_index_to_pandas, dfs))

        total_rows = result_batches.approx_total_rows
        if (total_rows is not None) and (max_results is not None):
            total_rows = min(total_rows, max_results)

        return PandasBatches(
            dfs,
            total_rows,
            total_bytes_processed=execution_result.total_bytes_processed,
        )

    def _copy_index_to_pandas(self, df: pd.DataFrame) -> pd.DataFrame:
        """Set the index on pandas DataFrame to match this block."""
        # Note: If BigQuery DataFrame has null index, a default one will be created for the local materialization.
        new_df = df.copy()
        if len(self.index_columns) > 0:
            new_df.set_index(list(self.index_columns), inplace=True)
            # Pandas names is annotated as list[str] rather than the more
            # general Sequence[Label] that BigQuery DataFrames has.
            # See: https://github.com/pandas-dev/pandas-stubs/issues/804
            new_df.index.names = self.index.names  # type: ignore
        new_df.columns = self.column_labels
        return new_df

    def _materialize_local(
        self, materialize_options: MaterializationOptions = MaterializationOptions()
    ) -> tuple[pd.DataFrame, Optional[bigquery.QueryJob]]:
        """Run query and download results as a pandas DataFrame. Return the total number of results as well."""
        # TODO(swast): Allow for dry run and timeout.
        under_10gb = (
            (not materialize_options.allow_large_results)
            if (materialize_options.allow_large_results is not None)
            else (not bigframes.options._allow_large_results)
        )
        execute_result = self.session._executor.execute(
            self.expr,
            execution_spec.ExecutionSpec(
                promise_under_10gb=under_10gb,
                ordered=materialize_options.ordered,
            ),
        )
        result_batches = execute_result.batches()

        sample_config = materialize_options.downsampling
        if result_batches.approx_total_bytes is not None:
            table_mb = result_batches.approx_total_bytes / _BYTES_TO_MEGABYTES
            max_download_size = sample_config.max_download_size
            fraction = (
                max_download_size / table_mb
                if (max_download_size is not None) and (table_mb != 0)
                else 2
            )
        else:
            # Since we cannot acquire the table size without a query_job,
            # we skip the sampling.
            if sample_config.enable_downsampling:
                msg = bfe.format_message(
                    "Sampling is disabled and there is no download size limit when 'allow_large_results' is set to "
                    "False. To prevent downloading excessive data, it is recommended to use the peek() method, or "
                    "limit the data with methods like .head() or .sample() before proceeding with downloads."
                )
                warnings.warn(msg, category=UserWarning)
            fraction = 2

        # TODO: Maybe materialize before downsampling
        # Some downsampling methods
        if fraction < 1 and (result_batches.approx_total_rows is not None):
            if not sample_config.enable_downsampling:
                raise RuntimeError(
                    f"The data size ({table_mb:.2f} MB) exceeds the maximum download limit of "
                    f"{max_download_size} MB. You can:\n\t* Enable downsampling in global options:\n"
                    "\t\t`bigframes.options.sampling.enable_downsampling = True`\n"
                    "\t* Update the global `max_download_size` option. Please make sure "
                    "there is enough memory available:\n"
                    "\t\t`bigframes.options.sampling.max_download_size = desired_size`"
                    " # Setting it to None will download all the data\n"
                    f"{constants.FEEDBACK_LINK}"
                )
            msg = bfe.format_message(
                f"The data size ({table_mb:.2f} MB) exceeds the maximum download limit of"
                f"({max_download_size} MB). It will be downsampled to {max_download_size} "
                "MB for download.\nPlease refer to the documentation for configuring "
                "the downloading limit."
            )
            warnings.warn(msg, category=UserWarning)
            total_rows = result_batches.approx_total_rows
            # Remove downsampling config from subsequent invocations, as otherwise could result in many
            # iterations if downsampling undershoots
            if sample_config.sampling_method == "head":
                # Just truncates the result iterator without a follow-up query
                raw_df = result_batches.to_pandas(limit=int(total_rows * fraction))
            elif (
                sample_config.sampling_method == "uniform"
                and sample_config.random_state is None
            ):
                # Pushes sample into result without new query
                sampled_batches = execute_result.batches(sample_rate=fraction)
                raw_df = sampled_batches.to_pandas()
            else:  # uniform sample with random state requires a full follow-up query
                down_sampled_block = self.split(
                    fracs=(fraction,),
                    random_state=sample_config.random_state,
                    sort=False,
                )[0]
                return down_sampled_block._materialize_local(
                    MaterializationOptions(ordered=materialize_options.ordered)
                )
        else:
            raw_df = result_batches.to_pandas()
        df = self._copy_index_to_pandas(raw_df)
        df.set_axis(self.column_labels, axis=1, copy=False)
        return df, execute_result.query_job

    def split(
        self,
        ns: Iterable[int] = (),
        fracs: Iterable[float] = (),
        *,
        random_state: Optional[int] = None,
        sort: Optional[bool | Literal["random"]] = "random",
    ) -> List[Block]:
        """Internal function to support splitting Block to multiple parts along index axis.

        At most one of ns and fracs can be passed in. If neither, default to ns = (1,).
        Return a list of sampled Blocks.
        """
        block = self
        if ns and fracs:
            raise ValueError("Only one of 'ns' or 'fracs' parameter must be specified.")

        if not ns and not fracs:
            ns = (1,)

        if ns:
            sample_sizes = ns
        else:
            total_rows = block.shape[0]
            # Round to nearest integer. "round half to even" rule applies.
            # At least to be 1.
            sample_sizes = [round(frac * total_rows) or 1 for frac in fracs]

        if random_state is None:
            random_state = random.randint(-(2**63), 2**63 - 1)

        # Create a new column with random_state value.
        block, random_state_col = block.create_constant(str(random_state))

        # Create an ordering col and convert to string
        block, ordering_col = block.promote_offsets()
        block, string_ordering_col = block.apply_unary_op(
            ordering_col, ops.AsTypeOp(to_type=bigframes.dtypes.STRING_DTYPE)
        )

        # Apply hash method to sum col and order by it.
        block, string_sum_col = block.apply_binary_op(
            string_ordering_col, random_state_col, ops.strconcat_op
        )
        block, hash_string_sum_col = block.apply_unary_op(string_sum_col, ops.hash_op)
        block = block.order_by(
            [ordering.OrderingExpression(ex.deref(hash_string_sum_col))]
        )

        intervals = []
        cur = 0

        for sample_size in sample_sizes:
            intervals.append((cur, cur + sample_size))
            cur += sample_size

        sliced_blocks = [
            typing.cast(Block, block.slice(start=lower, stop=upper))
            for lower, upper in intervals
        ]

        if sort is True:
            sliced_blocks = [
                sliced_block.order_by(
                    [
                        ordering.OrderingExpression(ex.deref(idx_col))
                        for idx_col in sliced_block.index_columns
                    ]
                )
                for sliced_block in sliced_blocks
            ]
        elif sort is False:
            sliced_blocks = [
                sliced_block.order_by(
                    [ordering.OrderingExpression(ex.deref(ordering_col))]
                )
                for sliced_block in sliced_blocks
            ]

        drop_cols = [
            random_state_col,
            ordering_col,
            string_ordering_col,
            string_sum_col,
            hash_string_sum_col,
        ]
        return [sliced_block.drop_columns(drop_cols) for sliced_block in sliced_blocks]

    def _compute_dry_run(
        self,
        value_keys: Optional[Iterable[str]] = None,
        *,
        ordered: bool = True,
        max_download_size: Optional[int] = None,
        sampling_method: Optional[str] = None,
        random_state: Optional[int] = None,
    ) -> typing.Tuple[pd.Series, bigquery.QueryJob]:
        sampling = self._get_sampling_option(
            max_download_size, sampling_method, random_state
        )
        if sampling.enable_downsampling:
            raise NotImplementedError("Dry run with sampling is not supported")

        expr = self._apply_value_keys_to_expr(value_keys=value_keys)
        query_job = self.session._executor.dry_run(expr, ordered)

        column_dtypes = {
            col: self.expr.get_column_type(self.resolve_label_exact_or_error(col))
            for col in self.column_labels
        }

        dry_run_stats = dry_runs.get_query_stats_with_dtypes(
            query_job, column_dtypes, self.index.dtypes, self.expr.node
        )
        return dry_run_stats, query_job

    def _apply_value_keys_to_expr(self, value_keys: Optional[Iterable[str]] = None):
        expr = self._expr
        if value_keys is not None:
            expr = expr.select_columns(itertools.chain(self._index_columns, value_keys))
        return expr

    def with_column_labels(
        self,
        value: typing.Union[pd.Index, typing.Iterable[Label]],
    ) -> Block:
        label_list = value.copy() if isinstance(value, pd.Index) else pd.Index(value)
        if len(label_list) != len(self.value_columns):
            raise ValueError(
                f"The column labels size `{len(label_list)} ` should equal to the value"
                + f"columns size: {len(self.value_columns)}."
            )
        block = Block(
            self._expr,
            index_columns=self.index_columns,
            column_labels=label_list,
            index_labels=self.index.names,
        )
        singleton_label = len(list(value)) == 1 and list(value)[0]
        if singleton_label is not None and self._transpose_cache is not None:
            new_cache, label_id = self._transpose_cache.create_constant(singleton_label)
            new_cache = new_cache.set_index([label_id])
            block = block.with_transpose_cache(new_cache)
        return block

    def with_transpose_cache(self, transposed: Block):
        return Block(
            self._expr,
            index_columns=self.index_columns,
            column_labels=self._column_labels,
            index_labels=self.index.names,
            transpose_cache=transposed,
        )

    def with_index_labels(self, value: typing.Sequence[Label]) -> Block: