feat: add unordered sql compilation

bigframes/core/__init__.py

@@ -23,7 +23,8 @@
 import ibis.expr.types as ibis_types
 import pandas
 
-import bigframes.core.compile as compiled
+import bigframes.core.compile.compiled as compiled
+import bigframes.core.compile.compiler as compiler
 import bigframes.core.guid
 import bigframes.core.nodes as nodes
 from bigframes.core.ordering import OrderingColumnReference
@@ -78,7 +79,7 @@ def from_pandas(cls, pd_df: pandas.DataFrame):
 
  @property
  def column_ids(self) -> typing.Sequence[str]:
- return self.compile().column_ids
+ return self._compile_ordered().column_ids
 
  @property
  def session(self) -> Session:
@@ -88,15 +89,18 @@ def session(self) -> Session:
  return self.node.session[0] if required_session else get_global_session()
 
  def get_column_type(self, key: str) -> bigframes.dtypes.Dtype:
- return self.compile().get_column_type(key)
+ return self._compile_ordered().get_column_type(key)
 
- def compile(self) -> compiled.CompiledArrayValue:
- return compiled.compile_node(self.node)
+ def _compile_ordered(self) -> compiled.OrderedIR:
+ return compiler.compile_ordered(self.node)
+
+ def _compile_unordered(self) -> compiled.UnorderedIR:
+ return compiler.compile_unordered(self.node)
 
  def shape(self) -> typing.Tuple[int, int]:
  """Returns dimensions as (length, width) tuple."""
- width = len(self.compile().columns)
- count_expr = self.compile()._to_ibis_expr("unordered").count()
+ width = len(self._compile_unordered().columns)
+ count_expr = self._compile_unordered()._to_ibis_expr().count()
 
  # Support in-memory engines for hermetic unit tests.
  if not self.node.session:
@@ -121,11 +125,14 @@ def to_sql(
  col_id_overrides: typing.Mapping[str, str] = {},
  sorted: bool = False,
  ) -> str:
- return self.compile().to_sql(
- offset_column=offset_column,
- col_id_overrides=col_id_overrides,
- sorted=sorted,
- )
+ if sorted or offset_column:
+ return self._compile_ordered().to_sql(
+ offset_column=offset_column,
+ col_id_overrides=col_id_overrides,
+ sorted=sorted,
+ )
+ else:
+ return self._compile_unordered().to_sql(col_id_overrides=col_id_overrides)
 
  def start_query(
  self,
@@ -154,7 +161,7 @@ def start_query(
 
  def cached(self, cluster_cols: typing.Sequence[str]) -> ArrayValue:
  """Write the ArrayValue to a session table and create a new block object that references it."""
- compiled_value = self.compile()
+ compiled_value = self._compile_ordered()
  ibis_expr = compiled_value._to_ibis_expr(
  ordering_mode="unordered", expose_hidden_cols=True
  )

bigframes/core/blocks.py

@@ -386,6 +386,8 @@ def to_pandas(
  max_download_size: Optional[int] = None,
  sampling_method: Optional[str] = None,
  random_state: Optional[int] = None,
+ *,
+ ordered: bool = True,
  ) -> Tuple[pd.DataFrame, bigquery.QueryJob]:
  """Run query and download results as a pandas DataFrame."""
  if max_download_size is None:
@@ -412,6 +414,7 @@ def to_pandas(
  max_download_size=max_download_size,
  sampling_method=sampling_method,
  random_state=random_state,
+ ordered=ordered,
  )
  return df, query_job
 
@@ -446,12 +449,16 @@ def _compute_and_count(
  max_download_size: Optional[int] = None,
  sampling_method: Optional[str] = None,
  random_state: Optional[int] = None,
+ *,
+ ordered: bool = True,
  ) -> Tuple[pd.DataFrame, int, 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.
  expr = self._apply_value_keys_to_expr(value_keys=value_keys)
 
- results_iterator, query_job = expr.start_query(max_results=max_results)
+ results_iterator, query_job = expr.start_query(
+ max_results=max_results, sorted=ordered
+ )
 
  table_size = (
  expr.session._get_table_size(query_job.destination) / _BYTES_TO_MEGABYTES

bigframes/core/compile/__init__.py

@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from bigframes.core.compile.compiled import CompiledArrayValue
-from bigframes.core.compile.compiler import compile_node
+from bigframes.core.compile.compiled import OrderedIR, UnorderedIR
+from bigframes.core.compile.compiler import compile_ordered, compile_unordered
 
 __all__ = [
- "compile_node",
- "CompiledArrayValue",
+ "compile_ordered",
+ "compile_unordered",
+ "OrderedIR",
+ "UnorderedIR",
 ]