perf: inline small Series and DataFrames in query text (#45)

The prevents unnecessary load and query jobs.

Towards internal issue 296474170 🦕

Author: tswast

bigframes/core/__init__.py

@@ -144,21 +144,56 @@ def mem_expr_from_pandas(
         """
         Builds an in-memory only (SQL only) expr from a pandas dataframe.
 
-        Caution: If session is None, only a subset of expr functionality will be available (null Session is usually not supported).
+        Caution: If session is None, only a subset of expr functionality will
+        be available (null Session is usually not supported).
         """
-        # must set non-null column labels. these are not the user-facing labels
-        pd_df = pd_df.set_axis(
-            [column or bigframes.core.guid.generate_guid() for column in pd_df.columns],
-            axis="columns",
-        )
+        # We can't include any hidden columns in the ArrayValue constructor, so
+        # grab the column names before we add the hidden ordering column.
+        column_names = [str(column) for column in pd_df.columns]
+        # Make sure column names are all strings.
+        pd_df = pd_df.set_axis(column_names, axis="columns")
         pd_df = pd_df.assign(**{ORDER_ID_COLUMN: range(len(pd_df))})
+
         # ibis memtable cannot handle NA, must convert to None
         pd_df = pd_df.astype("object")  # type: ignore
         pd_df = pd_df.where(pandas.notnull(pd_df), None)
+
+        # NULL type isn't valid in BigQuery, so retry with an explicit schema in these cases.
         keys_memtable = ibis.memtable(pd_df)
+        schema = keys_memtable.schema()
+        new_schema = []
+        for column_index, column in enumerate(schema):
+            if column == ORDER_ID_COLUMN:
+                new_type: ibis_dtypes.DataType = ibis_dtypes.int64
+            else:
+                column_type = schema[column]
+                # The autodetected type might not be one we can support, such
+                # as NULL type for empty rows, so convert to a type we do
+                # support.
+                new_type = bigframes.dtypes.bigframes_dtype_to_ibis_dtype(
+                    bigframes.dtypes.ibis_dtype_to_bigframes_dtype(column_type)
+                )
+                # TODO(swast): Ibis memtable doesn't use backticks in struct
+                # field names, so spaces and other characters aren't allowed in
+                # the memtable context. Blocked by
+                # https://github.com/ibis-project/ibis/issues/7187
+                column = f"col_{column_index}"
+            new_schema.append((column, new_type))
+
+        # must set non-null column labels. these are not the user-facing labels
+        pd_df = pd_df.set_axis(
+            [column for column, _ in new_schema],
+            axis="columns",
+        )
+        keys_memtable = ibis.memtable(pd_df, schema=ibis.schema(new_schema))
+
         return cls(
             session,  # type: ignore # Session cannot normally be none, see "caution" above
             keys_memtable,
+            columns=[
+                keys_memtable[f"col_{column_index}"].name(column)
+                for column_index, column in enumerate(column_names)
+            ],
             ordering=ExpressionOrdering(
                 ordering_value_columns=[OrderingColumnReference(ORDER_ID_COLUMN)],
                 total_ordering_columns=frozenset([ORDER_ID_COLUMN]),
@@ -426,11 +461,16 @@ def shape(self) -> typing.Tuple[int, int]:
         width = len(self.columns)
         count_expr = self._to_ibis_expr(ordering_mode="unordered").count()
         sql = self._session.ibis_client.compile(count_expr)
-        row_iterator, _ = self._session._start_query(
-            sql=sql,
-            max_results=1,
-        )
-        length = next(row_iterator)[0]
+
+        # Support in-memory engines for hermetic unit tests.
+        if not isinstance(sql, str):
+            length = self._session.ibis_client.execute(count_expr)
+        else:
+            row_iterator, _ = self._session._start_query(
+                sql=sql,
+                max_results=1,
+            )
+            length = next(row_iterator)[0]
         return (length, width)
 
     def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue:

bigframes/core/blocks.py

@@ -44,6 +44,7 @@
 import bigframes.dtypes
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
+import third_party.bigframes_vendored.pandas.io.common as vendored_pandas_io_common
 
 # Type constraint for wherever column labels are used
 Label = typing.Hashable
@@ -1522,37 +1523,35 @@ def _is_monotonic(
         return result
 
 
-def block_from_local(data, session=None, use_index=True) -> Block:
-    # TODO(tbergeron): Handle duplicate column labels
+def block_from_local(data, session=None) -> Block:
     pd_data = pd.DataFrame(data)
+    columns = pd_data.columns
 
-    column_labels = list(pd_data.columns)
-    if not all((label is None) or isinstance(label, str) for label in column_labels):
-        raise NotImplementedError(
-            f"Only string column labels supported. {constants.FEEDBACK_LINK}"
-        )
+    # Make a flattened version to treat as a table.
+    if len(pd_data.columns.names) > 1:
+        pd_data.columns = columns.to_flat_index()
 
-    if use_index:
-        if pd_data.index.nlevels > 1:
-            raise NotImplementedError(
-                f"multi-indices not supported. {constants.FEEDBACK_LINK}"
-            )
-        index_label = pd_data.index.name
-
-        index_id = guid.generate_guid()
-        pd_data = pd_data.reset_index(names=index_id)
-        keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session)
-        return Block(
-            keys_expr,
-            column_labels=column_labels,
-            index_columns=[index_id],
-            index_labels=[index_label],
-        )
-    else:
-        keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session)
-        keys_expr, offsets_id = keys_expr.promote_offsets()
-        # Constructor will create default range index
-        return Block(keys_expr, index_columns=[offsets_id], column_labels=column_labels)
+    index_labels = list(pd_data.index.names)
+    # The ArrayValue layer doesn't know about indexes, so make sure indexes
+    # are real columns with unique IDs.
+    pd_data = pd_data.reset_index(
+        names=[f"level_{level}" for level in range(len(index_labels))]
+    )
+    pd_data = pd_data.set_axis(
+        vendored_pandas_io_common.dedup_names(
+            list(pd_data.columns), is_potential_multiindex=False
+        ),
+        axis="columns",
+    )
+    index_ids = pd_data.columns[: len(index_labels)]
+
+    keys_expr = core.ArrayValue.mem_expr_from_pandas(pd_data, session)
+    return Block(
+        keys_expr,
+        column_labels=columns,
+        index_columns=index_ids,
+        index_labels=index_labels,
+    )
 
 
 def _align_block_to_schema(

bigframes/dataframe.py

@@ -65,9 +65,7 @@
 
 # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type.
 # TODO(tbergeron): Convert to bytes-based limit
-# TODO(swast): Address issues with string escaping and empty tables before
-# re-enabling inline data (ibis.memtable) feature.
-MAX_INLINE_DF_SIZE = -1
+MAX_INLINE_DF_SIZE = 5000
 
 LevelType = typing.Union[str, int]
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]

bigframes/dtypes.py

@@ -96,13 +96,13 @@
     ),
 )
 
-BIGFRAMES_TO_IBIS: Dict[Dtype, IbisDtype] = {
+BIGFRAMES_TO_IBIS: Dict[Dtype, ibis_dtypes.DataType] = {
     pandas: ibis for ibis, pandas in BIDIRECTIONAL_MAPPINGS
 }
 
-IBIS_TO_BIGFRAMES: Dict[
-    Union[IbisDtype, ReadOnlyIbisDtype], Union[Dtype, np.dtype[Any]]
-] = {ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS}
+IBIS_TO_BIGFRAMES: Dict[ibis_dtypes.DataType, Union[Dtype, np.dtype[Any]]] = {
+    ibis: pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS
+}
 # Allow REQUIRED fields to map correctly.
 IBIS_TO_BIGFRAMES.update(
     {ibis.copy(nullable=False): pandas for ibis, pandas in BIDIRECTIONAL_MAPPINGS}
@@ -130,7 +130,7 @@
 
 
 def ibis_dtype_to_bigframes_dtype(
-    ibis_dtype: Union[IbisDtype, ReadOnlyIbisDtype]
+    ibis_dtype: ibis_dtypes.DataType,
 ) -> Union[Dtype, np.dtype[Any]]:
     """Converts an Ibis dtype to a BigQuery DataFrames dtype
 
@@ -155,6 +155,9 @@ def ibis_dtype_to_bigframes_dtype(
 
     if ibis_dtype in IBIS_TO_BIGFRAMES:
         return IBIS_TO_BIGFRAMES[ibis_dtype]
+    elif isinstance(ibis_dtype, ibis_dtypes.Null):
+        # Fallback to STRING for NULL values for most flexibility in SQL.
+        return IBIS_TO_BIGFRAMES[ibis_dtypes.string]
     else:
         raise ValueError(
             f"Unexpected Ibis data type {ibis_dtype}. {constants.FEEDBACK_LINK}"
@@ -185,8 +188,8 @@ def ibis_table_to_canonical_types(table: ibis_types.Table) -> ibis_types.Table:
 
 
 def bigframes_dtype_to_ibis_dtype(
-    bigframes_dtype: Union[DtypeString, Dtype]
-) -> IbisDtype:
+    bigframes_dtype: Union[DtypeString, Dtype, np.dtype[Any]]
+) -> ibis_dtypes.DataType:
     """Converts a BigQuery DataFrames supported dtype to an Ibis dtype.
 
     Args:
@@ -281,7 +284,9 @@ def literal_to_ibis_scalar(
     return scalar_expr
 
 
-def cast_ibis_value(value: ibis_types.Value, to_type: IbisDtype) -> ibis_types.Value:
+def cast_ibis_value(
+    value: ibis_types.Value, to_type: ibis_dtypes.DataType
+) -> ibis_types.Value:
     """Perform compatible type casts of ibis values
 
     Args:

bigframes/operations/base.py

@@ -30,9 +30,7 @@
 
 # BigQuery has 1 MB query size limit, 5000 items shouldn't take more than 10% of this depending on data type.
 # TODO(tbergeron): Convert to bytes-based limit
-# TODO(swast): Address issues with string escaping and empty tables before
-# re-enabling inline data (ibis.memtable) feature.
-MAX_INLINE_SERIES_SIZE = -1
+MAX_INLINE_SERIES_SIZE = 5000
 
 
 class SeriesMethods:

setup.py

@@ -44,12 +44,12 @@
     "google-cloud-resource-manager >=1.10.3",
     "google-cloud-storage >=2.0.0",
     # TODO: Relax upper bound once we have fixed `system_prerelease` tests.
-    "ibis-framework[bigquery] >=6.0.0,<=6.1.0",
+    "ibis-framework[bigquery] >=6.2.0,<7.0.0dev",
     "pandas >=1.5.0",
     "pydata-google-auth >=1.8.2",
     "requests >=2.27.1",
     "scikit-learn >=1.2.2",
-    "sqlalchemy >=1.4,<3.0",
+    "sqlalchemy >=1.4,<3.0dev",
     "ipywidgets >=7.7.1",
     "humanize >= 4.6.0",
 ]

testing/constraints-3.9.txt

@@ -45,7 +45,7 @@ greenlet==2.0.2
 grpc-google-iam-v1==0.12.6
 grpcio==1.53.0
 grpcio-status==1.48.2
-ibis-framework==6.0.0
+ibis-framework==6.2.0
 humanize==4.6.0
 identify==2.5.22
 idna==3.4

tests/unit/core/__init__.py

@@ -0,0 +1,13 @@
+# 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.

tests/unit/core/test_blocks.py

@@ -0,0 +1,85 @@
+# 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.
+
+import pandas
+import pandas.testing
+import pytest
+
+import bigframes.core.blocks as blocks
+
+from .. import resources
+
+
+@pytest.mark.parametrize(
+    ("data",),
+    (
+        pytest.param(
+            {"test 1": [1, 2, 3], "test 2": [0.25, 0.5, 0.75]},
+            id="dict_spaces_in_column_names",
+        ),
+        pytest.param(
+            [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
+            id="nested_list",
+        ),
+        pytest.param(
+            pandas.concat(
+                [
+                    pandas.Series([1, 2, 3], name="some col"),
+                    pandas.Series([2, 3, 4], name="some col"),
+                ],
+                axis="columns",
+            ),
+            id="duplicate_column_names",
+        ),
+        pytest.param(
+            pandas.DataFrame(
+                {"test": [1, 2, 3]},
+                index=pandas.Index(["a", "b", "c"], name="string index"),
+            ),
+            id="string_index",
+        ),
+        pytest.param(
+            pandas.DataFrame(
+                [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]],
+                columns=pandas.MultiIndex.from_tuples(
+                    [(1, 1), (1, 2), (0, 0), (0, 1)],
+                    names=["some level", "another level"],
+                ),
+            ),
+            marks=[
+                pytest.mark.skipif(
+                    tuple(pandas.__version__.split()) < ("2", "0", "0"),
+                    reason="pandas 1.5.3 treats column MultiIndex as Index of tuples",
+                ),
+            ],
+            id="multiindex_columns",
+        ),
+        pytest.param(
+            pandas.DataFrame(
+                {"test": [1, 2, 3]},
+                index=pandas.MultiIndex.from_tuples([(1, 1), (1, 2), (0, 0)]),
+            ),
+            id="multiindex_rows",
+        ),
+    ),
+)
+def test_block_from_local(data):
+    expected = pandas.DataFrame(data)
+    session = resources.create_pandas_session({})
+
+    block = blocks.block_from_local(data, session=session)
+
+    pandas.testing.assert_index_equal(block.column_labels, expected.columns)
+    assert tuple(block.index_labels) == tuple(expected.index.names)
+    assert block.shape == expected.shape