googleapis · gcf-merge-on-green · Nov 9, 2023 · Oct 27, 2023 · Oct 28, 2023 · Oct 28, 2023
@@ -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
        )

@@ -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

@@ -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",
 ]