do not use numpy for offsets

Author: tswast

bigframes/core/compile/compiler.py

@@ -22,10 +22,9 @@
 import bigframes_vendored.ibis.expr.api as ibis_api
 import bigframes_vendored.ibis.expr.datatypes as ibis_dtypes
 import bigframes_vendored.ibis.expr.types as ibis_types
-import pyarrow as pa
 
 from bigframes import dtypes, operations
-from bigframes.core import expression
+from bigframes.core import expression, pyarrow_utils
 import bigframes.core.compile.compiled as compiled
 import bigframes.core.compile.concat as concat_impl
 import bigframes.core.compile.configs as configs
@@ -172,9 +171,7 @@ def compile_readlocal(node: nodes.ReadLocalNode, *args):
     pa_table = pa_table.rename_columns([item.id.sql for item in node.scan_list.items])
 
     if offsets:
-        pa_table = pa_table.append_column(
-            offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
-        )
+        pa_table = pyarrow_utils.append_offsets(pa_table, offsets)
     return compiled.UnorderedIR.from_polars(pa_table, bq_schema)
 
 

bigframes/core/compile/sqlglot/compiler.py

@@ -18,10 +18,9 @@
 import typing
 
 from google.cloud import bigquery
-import pyarrow as pa
 import sqlglot.expressions as sge
 
-from bigframes.core import expression, guid, identifiers, nodes, rewrite
+from bigframes.core import expression, guid, identifiers, nodes, pyarrow_utils, rewrite
 from bigframes.core.compile import configs
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 import bigframes.core.compile.sqlglot.sqlglot_ir as ir
@@ -155,9 +154,7 @@ def compile_readlocal(self, node: nodes.ReadLocalNode, *args) -> ir.SQLGlotIR:
 
         offsets = node.offsets_col.sql if node.offsets_col else None
         if offsets:
-            pa_table = pa_table.append_column(
-                offsets, pa.array(range(pa_table.num_rows), type=pa.int64())
-            )
+            pa_table = pyarrow_utils.append_offsets(pa_table, offsets)
 
         return ir.SQLGlotIR.from_pyarrow(pa_table, node.schema, uid_gen=self.uid_gen)
 

bigframes/core/pyarrow_utils.py

@@ -85,3 +85,12 @@ def truncate_pyarrow_iterable(
         else:
             yield batch
             total_yielded += batch.num_rows
+
+
+def append_offsets(
+    pa_table: pa.Table,
+    offsets_col: str,
+) -> pa.Table:
+    return pa_table.append_column(
+        offsets_col, pa.array(range(pa_table.num_rows), type=pa.int64())
+    )

bigframes/session/_io/bigquery/read_gbq_query.py

@@ -20,12 +20,10 @@
 
 from google.cloud import bigquery
 import google.cloud.bigquery.table
-import numpy
 import pandas
-import pyarrow
 
 from bigframes import dataframe
-from bigframes.core import local_data
+from bigframes.core import local_data, pyarrow_utils
 import bigframes.core as core
 import bigframes.core.blocks as blocks
 import bigframes.core.guid
@@ -67,10 +65,7 @@ def create_dataframe_from_row_iterator(
     # TODO(tswast): Use array_value.promote_offsets() instead once that node is
     # supported by the local engine.
     offsets_col = bigframes.core.guid.generate_guid()
-    pa_table = pa_table.append_column(
-        pyarrow.field(offsets_col, pyarrow.int64()),
-        [numpy.arange(pa_table.num_rows)],
-    )
+    pa_table = pyarrow_utils.append_offsets(pa_table, offsets_col=offsets_col)
 
     # We use the ManagedArrowTable constructor directly, because the
     # results of to_arrow() should be the source of truth with regards