feat: Initial support for biglake iceberg tables (#2409)

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -17,9 +17,8 @@
 importdatetime
 importfunctools
 importtyping
-fromtypingimportIterable, List, Mapping, Optional, Sequence, Tuple
+fromtypingimportIterable, List, Mapping, Optional, Sequence, Tuple, Union
 
-importgoogle.cloud.bigquery
 importpandas
 importpyarrowaspa
 
@@ -91,7 +90,7 @@ def from_range(cls, start, end, step):
 @classmethod
 deffrom_table(
 cls,
-table: google.cloud.bigquery.Table,
+table: Union[bq_data.BiglakeIcebergTable, bq_data.GbqNativeTable],
 session: Session,
 *,
 columns: Optional[Sequence[str]] =None,
@@ -103,8 +102,6 @@ def from_table(
  ):
 ifoffsets_colandprimary_key:
 raiseValueError("must set at most one of 'offests', 'primary_key'")
-# define data source only for needed columns, this makes row-hashing cheaper
-table_def=bq_data.GbqTable.from_table(table, columns=columnsor ())
 
 # create ordering from info
 ordering=None
@@ -115,7 +112,9 @@ def from_table(
  [ids.ColumnId(key_part) forkey_partinprimary_key]
  )
 
-bf_schema=schemata.ArraySchema.from_bq_table(table, columns=columns)
+bf_schema=schemata.ArraySchema.from_bq_schema(
+table.physical_schema, columns=columns
+ )
 # Scan all columns by default, we define this list as it can be pruned while preserving source_def
 scan_list=nodes.ScanList(
 tuple(
@@ -124,7 +123,7 @@ def from_table(
  )
  )
 source_def=bq_data.BigqueryDataSource(
-table=table_def,
+table=table,
 schema=bf_schema,
 at_time=at_time,
 sql_predicate=predicate,

bigframes/core/bq_data.py

@@ -22,74 +22,214 @@
 importqueue
 importthreading
 importtyping
-fromtypingimportAny, Iterator, Optional, Sequence, Tuple
+fromtypingimportAny, Iterator, List, Literal, Optional, Sequence, Tuple, Union
 
 fromgoogle.cloudimportbigquery_storage_v1
 importgoogle.cloud.bigqueryasbq
 importgoogle.cloud.bigquery_storage_v1.typesasbq_storage_types
 fromgoogle.protobufimporttimestamp_pb2
 importpyarrowaspa
 
+importbigframes.constants
 frombigframes.coreimportpyarrow_utils
 importbigframes.core.schema
 
 iftyping.TYPE_CHECKING:
 importbigframes.core.orderingasorderings
 
 
+def_resolve_standard_gcp_region(bq_region: str):
+"""
+ Resolve bq regions to standardized
+ """
+ifbq_region.casefold() =="US":
+return"us-central1"
+elifbq_region.casefold() =="EU":
+return"europe-west4"
+returnbq_region
+
+
+defis_irc_table(table_id: str):
+"""
+ Determines if a table id should be resolved through the iceberg rest catalog.
+ """
+returnlen(table_id.split(".")) ==4
+
+
+defis_compatible(
+data_region: Union[GcsRegion, BigQueryRegion], session_location: str
+) ->bool:
+# based on https://docs.cloud.google.com/bigquery/docs/locations#storage-location-considerations
+ifisinstance(data_region, BigQueryRegion):
+returndata_region.name==session_location
+else:
+assertisinstance(data_region, GcsRegion)
+# TODO(b/463675088): Multi-regions don't yet support rest catalog tables
+ifsession_locationinbigframes.constants.BIGQUERY_MULTIREGIONS:
+returnFalse
+return_resolve_standard_gcp_region(session_location) indata_region.included
+
+
+defget_default_bq_region(data_region: Union[GcsRegion, BigQueryRegion]) ->str:
+ifisinstance(data_region, BigQueryRegion):
+returndata_region.name
+elifisinstance(data_region, GcsRegion):
+# should maybe try to track and prefer primary replica?
+returndata_region.included[0]
+
+
+@dataclasses.dataclass(frozen=True)
+classBigQueryRegion:
+name: str
+
+
+@dataclasses.dataclass(frozen=True)
+classGcsRegion:
+# this is the name of gcs regions, which may be names for multi-regions, so shouldn't be compared with non-gcs locations
+storage_regions: tuple[str, ...]
+# this tracks all the included standard, specific regions (eg us-east1), and should be comparable to bq regions (except non-standard US, EU, omni regions)
+included: tuple[str, ...]
+
+
+# what is the line between metadata and core fields? Mostly metadata fields are optional or unreliable, but its fuzzy
 @dataclasses.dataclass(frozen=True)
-classGbqTable:
+classTableMetadata:
+# this size metadata might be stale, don't use where strict correctness is needed
+location: Union[BigQueryRegion, GcsRegion]
+type: Literal["TABLE", "EXTERNAL", "VIEW", "MATERIALIZE_VIEW", "SNAPSHOT"]
+numBytes: Optional[int] =None
+numRows: Optional[int] =None
+created_time: Optional[datetime.datetime] =None
+modified_time: Optional[datetime.datetime] =None
+
+
+@dataclasses.dataclass(frozen=True)
+classGbqNativeTable:
 project_id: str=dataclasses.field()
 dataset_id: str=dataclasses.field()
 table_id: str=dataclasses.field()
 physical_schema: Tuple[bq.SchemaField, ...] =dataclasses.field()
-is_physically_stored: bool=dataclasses.field()
-cluster_cols: typing.Optional[Tuple[str, ...]]
+metadata: TableMetadata=dataclasses.field()
+partition_col: Optional[str] =None
+cluster_cols: typing.Optional[Tuple[str, ...]] =None
+primary_key: Optional[Tuple[str, ...]] =None
 
 @staticmethod
-deffrom_table(table: bq.Table, columns: Sequence[str] = ()) ->GbqTable:
+deffrom_table(table: bq.Table, columns: Sequence[str] = ()) ->GbqNativeTable:
 # Subsetting fields with columns can reduce cost of row-hash default ordering
 ifcolumns:
 schema=tuple(itemforitemintable.schemaifitem.nameincolumns)
 else:
 schema=tuple(table.schema)
-returnGbqTable(
+
+metadata=TableMetadata(
+numBytes=table.num_bytes,
+numRows=table.num_rows,
+location=BigQueryRegion(table.location), # type: ignore
+type=table.table_typeor"TABLE", # type: ignore
+created_time=table.created,
+modified_time=table.modified,
+ )
+partition_col=None
+iftable.range_partitioning:
+partition_col=table.range_partitioning.field
+eliftable.time_partitioning:
+partition_col=table.time_partitioning.field
+
+returnGbqNativeTable(
 project_id=table.project,
 dataset_id=table.dataset_id,
 table_id=table.table_id,
 physical_schema=schema,
-is_physically_stored=(table.table_typein ["TABLE", "MATERIALIZED_VIEW"]),
+partition_col=partition_col,
 cluster_cols=None
-iftable.clustering_fieldsisNone
+if(table.clustering_fieldsisNone)
 elsetuple(table.clustering_fields),
+primary_key=tuple(_get_primary_keys(table)),
+metadata=metadata,
  )
 
 @staticmethod
 deffrom_ref_and_schema(
 table_ref: bq.TableReference,
 schema: Sequence[bq.SchemaField],
+location: str,
+table_type: Literal["TABLE"] ="TABLE",
 cluster_cols: Optional[Sequence[str]] =None,
- ) ->GbqTable:
-returnGbqTable(
+ ) ->GbqNativeTable:
+returnGbqNativeTable(
 project_id=table_ref.project,
 dataset_id=table_ref.dataset_id,
 table_id=table_ref.table_id,
+metadata=TableMetadata(location=BigQueryRegion(location), type=table_type),
 physical_schema=tuple(schema),
-is_physically_stored=True,
 cluster_cols=tuple(cluster_cols) ifcluster_colselseNone,
  )
 
+@property
+defis_physically_stored(self) ->bool:
+returnself.metadata.typein ["TABLE", "MATERIALIZED_VIEW"]
+
 defget_table_ref(self) ->bq.TableReference:
 returnbq.TableReference(
 bq.DatasetReference(self.project_id, self.dataset_id), self.table_id
  )
 
+defget_full_id(self, quoted: bool=False) ->str:
+ifquoted:
+returnf"`{self.project_id}`.`{self.dataset_id}`.`{self.table_id}`"
+returnf"{self.project_id}.{self.dataset_id}.{self.table_id}"
+
 @property
 @functools.cache
 defschema_by_id(self):
 return {col.name: colforcolinself.physical_schema}
 
 
+@dataclasses.dataclass(frozen=True)
+classBiglakeIcebergTable:
+project_id: str=dataclasses.field()
+catalog_id: str=dataclasses.field()
+namespace_id: str=dataclasses.field()
+table_id: str=dataclasses.field()
+physical_schema: Tuple[bq.SchemaField, ...] =dataclasses.field()
+cluster_cols: typing.Optional[Tuple[str, ...]]
+metadata: TableMetadata
+
+defget_full_id(self, quoted: bool=False) ->str:
+ifquoted:
+returnf"`{self.project_id}`.`{self.catalog_id}`.`{self.namespace_id}`.`{self.table_id}`"
+return (
+f"{self.project_id}.{self.catalog_id}.{self.namespace_id}.{self.table_id}"
+ )
+
+@property
+@functools.cache
+defschema_by_id(self):
+return {col.name: colforcolinself.physical_schema}
+
+@property
+defpartition_col(self) ->Optional[str]:
+# TODO: Use iceberg partition metadata
+returnNone
+
+@property
+defdataset_id(self) ->str:
+"""
+ Not a true dataset, but serves as the dataset component of the identifer in sql queries
+ """
+returnf"{self.catalog_id}.{self.namespace_id}"
+
+@property
+defprimary_key(self) ->Optional[Tuple[str, ...]]:
+returnNone
+
+defget_table_ref(self) ->bq.TableReference:
+returnbq.TableReference(
+bq.DatasetReference(self.project_id, self.dataset_id), self.table_id
+ )
+
+
 @dataclasses.dataclass(frozen=True)
 classBigqueryDataSource:
 """
@@ -104,13 +244,13 @@ def __post_init__(self):
 self.schema.names
  )
 
-table: GbqTable
+table: Union[GbqNativeTable, BiglakeIcebergTable]
 schema: bigframes.core.schema.ArraySchema
 at_time: typing.Optional[datetime.datetime] =None
 # Added for backwards compatibility, not validated
 sql_predicate: typing.Optional[str] =None
 ordering: typing.Optional[orderings.RowOrdering] =None
-# Optimization field
+# Optimization field, must be correct if set, don't put maybe-stale number here
 n_rows: Optional[int] =None
 
 
@@ -188,6 +328,8 @@ def get_arrow_batches(
 project_id: str,
 sample_rate: Optional[float] =None,
 ) ->ReadResult:
+assertisinstance(data.table, GbqNativeTable)
+
 table_mod_options= {}
 read_options_dict: dict[str, Any] = {"selected_fields": list(columns)}
 
@@ -245,3 +387,21 @@ def process_batch(pa_batch):
 returnReadResult(
 batches, session.estimated_row_count, session.estimated_total_bytes_scanned
  )
+
+
+def_get_primary_keys(
+table: bq.Table,
+) ->List[str]:
+"""Get primary keys from table if they are set."""
+
+primary_keys: List[str] = []
+if (
+ (table_constraints:=getattr(table, "table_constraints", None)) isnotNone
+and (primary_key:=table_constraints.primary_key) isnotNone
+# This will be False for either None or empty list.
+# We want primary_keys = None if no primary keys are set.
+and (columns:=primary_key.columns)
+ ):
+primary_keys=columnsifcolumnsisnotNoneelse []
+
+returnprimary_keys

bigframes/core/compile/ibis_compiler/ibis_compiler.py

@@ -215,9 +215,7 @@ def _table_to_ibis(
 source: bq_data.BigqueryDataSource,
 scan_cols: typing.Sequence[str],
 ) ->ibis_types.Table:
-full_table_name= (
-f"{source.table.project_id}.{source.table.dataset_id}.{source.table.table_id}"
- )
+full_table_name=source.table.get_full_id(quoted=False)
 # Physical schema might include unused columns, unsupported datatypes like JSON
 physical_schema=ibis_bigquery.BigQuerySchema.to_ibis(
 list(source.table.physical_schema)

bigframes/core/nodes.py

@@ -825,9 +825,7 @@ def variables_introduced(self) -> int:
 
 @property
 defrow_count(self) ->typing.Optional[int]:
-ifself.source.sql_predicateisNoneandself.source.table.is_physically_stored:
-returnself.source.n_rows
-returnNone
+returnself.source.n_rows
 
 @property
 defnode_defined_ids(self) ->Tuple[identifiers.ColumnId, ...]:

bigframes/core/schema.py

@@ -17,7 +17,7 @@
 fromdataclassesimportdataclass
 importfunctools
 importtyping
-fromtypingimportDict, List, Optional, Sequence
+fromtypingimportDict, Optional, Sequence
 
 importgoogle.cloud.bigquery
 importpyarrow
@@ -40,31 +40,16 @@ class ArraySchema:
 def__iter__(self):
 yieldfromself.items
 
-@classmethod
-deffrom_bq_table(
-cls,
-table: google.cloud.bigquery.Table,
-column_type_overrides: Optional[
-typing.Dict[str, bigframes.dtypes.Dtype]
- ] =None,
-columns: Optional[Sequence[str]] =None,
- ):
-ifnotcolumns:
-fields=table.schema
-else:
-lookup= {field.name: fieldforfieldintable.schema}
-fields= [lookup[col] forcolincolumns]
-
-returnArraySchema.from_bq_schema(
-fields, column_type_overrides=column_type_overrides
- )
-
 @classmethod
 deffrom_bq_schema(
 cls,
-schema: List[google.cloud.bigquery.SchemaField],
+schema: Sequence[google.cloud.bigquery.SchemaField],
 column_type_overrides: Optional[Dict[str, bigframes.dtypes.Dtype]] =None,
+columns: Optional[Sequence[str]] =None,
  ):
+ifcolumns:
+lookup= {field.name: fieldforfieldinschema}
+schema= [lookup[col] forcolincolumns]
 ifcolumn_type_overridesisNone:
 column_type_overrides= {}
 items=tuple(