Supporting BigQuery's PARTITION BY and CLUSTER BY via generators

[hrshdhgd]

I have data in BigQuery tables, does LinkML's gen-sqltables have the possibility of generating queries having PARTITION BY and CLUSTER BY right out of the box?

My initial thought was to create a BigQueryGenerator class inheriting from SQLTableGenerator, along with a gen-bigquery command, and extending the schema to include partition_by and clustered_by attributes.

Slack thread

[cmungall]

SQLTableGenerator delegates to SQLAlchemy for dialect-specific features. SQLAlchemy has excellent support for dialects — including BigQuery

However, SQLAlchemy does not currently support PARTITION BY and CLUSTER BY for BigQuery.

That thread is from 2022, I would say the first thing is to check (including on that thread) if there has been any progress. It will be much easier to support this feature in LinkML if it's in SQLA. Otherwise, we can explore ways to bolt these, on e.g.. by subclassing the generator, but this may be complicated

[cmungall]

It would be good to double check with a quick sqla test to be sure it can be done from sqla.

if that's the case, we can proceed. I recommend we control this via annotations on the schema

[hrshdhgd]

Sounds like a plan!

It would be good to double check with a quick sqla test to be sure it can be done from sqla.

I've got to figure out how sqlalchemy works first .. 😆

[cmungall]

One other possible gotcha - gen-sqltables will normalize your schema for you - i.e if you have multivalued slots or slots that inline objects, these get translated. This is desirable for "pure" RDBMS work.

However, BigQuery is not a "pure" RDBMS, and it has features like STRUCTS and ARRAY. Are you expecting these to be used or are you ok with the normalization step?

[hrshdhgd]

features like STRUCTS and ARRAY. Are you expecting these to be used

Yes, definitely as far as I can tell. Would it make sense to have an independent BigQueryGenerator in that case?

[hrshdhgd]

Alright this works!

from sqlalchemy import create_engine, Table, Column, MetaData
from sqlalchemy.sql.sqltypes import INT, String, DateTime
from sqlalchemy.schema import CreateTable
from sqlalchemy_bigquery import BigQueryDialect  # Native BigQuery support
from google.cloud import bigquery

def test_bigquery_partition_cluster_options():
    metadata = MetaData()

    # Define table with partitioning and clustering
    table = Table(
        "test_table",
        metadata,
        Column("id", INT, primary_key=True),
        Column("name", String),
        Column("created_at", DateTime),
        Column("zipcode", INT),
        bigquery_time_partitioning=bigquery.TimePartitioning(
            field="created_at",
            type_="MONTH",
            expiration_ms=1000 * 60 * 60 * 24 * 30 * 6, # 6 months
        ),
        # bigquery_range_partitioning=bigquery.RangePartitioning(
        #     field="zipcode",
        #     range_=bigquery.PartitionRange(start=0, end=100, interval=10),
        # ),

        bigquery_require_partition_filter=True,
        
        
        bigquery_clustering_fields=["name", "id"],
    )

    # Verify partitioning and clustering options exist
    assert "bigquery" in table.dialect_options
    # assert table.dialect_options["bigquery"]["time_partitioning"] == "DATE(created_at)"
    assert table.dialect_options["bigquery"]["clustering_fields"] == ["name", "id"]

    ddl = str(CreateTable(table).compile(dialect=BigQueryDialect()))
    print(ddl)

if __name__ == "__main__":
    pytest.main([__file__])

prints

CREATE TABLE `test_table` (
        `id` INT64 NOT NULL, 
        `name` STRING, 
        `created_at` DATE, 
        `zipcode` INT64
) PARTITION BY DATE_TRUNC(created_at, MONTH)
CLUSTER BY name, id
OPTIONS(partition_expiration_days=180.0, require_partition_filter=true)

There is also bigquery_range_partitioning for columns that are integers and works the same way.

Does this look reasonable @cmungall ?

[hrshdhgd]

Using SQLA would result in generation of python code. The end user in this situation prefers the generation of SQL queries similar to what we get from gen-sqltables.

[cmungall]

We use SQLA for the generation of DDL/CREATE TABLE statements too.

linkml/linkml/generators/sqltablegen.py

Lines 146 to 235 in 9cf2c4c

	def generate_ddl(self, naming_policy: SqlNamingPolicy = None, **kwargs) -> str:
	ddl_str = ""
	def dump(sql, *multiparams, **params):
	nonlocal ddl_str
	ddl_str += f"{str(sql.compile(dialect=engine.dialect)).rstrip()};"
	engine = create_mock_engine(f"{self.dialect}://./MyDb", strategy="mock", executor=dump)
	schema_metadata = MetaData()
	sqltr = RelationalModelTransformer(SchemaView(self.schema))
	if not self.use_foreign_keys:
	sqltr.foreign_key_policy = ForeignKeyPolicy.NO_FOREIGN_KEYS
	tr_result = sqltr.transform(
	tgt_schema_name=kwargs.get("tgt_schema_name", None), top_class=kwargs.get("top_class", None)
	)
	schema = tr_result.schema
	def sql_name(n: str) -> str:
	if naming_policy is not None:
	if naming_policy == SqlNamingPolicy.underscore:
	return underscore(n)
	elif naming_policy == SqlNamingPolicy.camelcase:
	return camelcase(n)
	elif naming_policy == SqlNamingPolicy.preserve:
	return n
	else:
	raise Exception(f"Unknown: {naming_policy}")
	else:
	return n
	def strip_newlines(txt: Optional[str]) -> str:
	if txt is None:
	return ""
	return txt.replace("\n", "")
	# Currently SQLite dialect in SQLA does not generate comments; see
	# https://github.com/sqlalchemy/sqlalchemy/issues/1546#issuecomment-1067389172
	# As a workaround we add these as "--" comments via direct string manipulation
	include_comments = self.dialect == "sqlite"
	sv = SchemaView(schema)
	for cn, c in schema.classes.items():
	if include_comments:
	ddl_str += f'-- # Class: "{cn}" Description: "{strip_newlines(c.description)}"\n'
	pk_slot = sv.get_identifier_slot(cn)
	if c.attributes:
	cols = []
	for sn, s in c.attributes.items():
	is_pk = "primary_key" in s.annotations
	if pk_slot:
	is_pk = sn == pk_slot.name
	# else:
	# is_pk = True ## TODO: use unique key
	args = []
	if s.range in schema.classes and self.use_foreign_keys:
	fk = sql_name(self.get_foreign_key(s.range, sv))
	args = [ForeignKey(fk)]
	field_type = self.get_sql_range(s, schema)
	col = Column(
	sql_name(sn),
	field_type,
	*args,
	primary_key=is_pk,
	nullable=not s.required,
	)
	if include_comments:
	ddl_str += f"-- * Slot: {sn} Description: {strip_newlines(s.description)}\n"
	if s.description:
	col.comment = s.description
	cols.append(col)
	for uc_name, uc in c.unique_keys.items():
	def _sql_name(sn: str):
	if sn in c.attributes:
	return sql_name(sn)
	else:
	# for candidate in c.attributes.values():
	# if "original_slot" in candidate.annotations:
	# original = candidate.annotations["original_slot"]
	# if original.value == sn:
	# return sql_name(candidate.name)
	return None
	sql_names = [_sql_name(sn) for sn in uc.unique_key_slots]
	if any(sn is None for sn in sql_names):
	continue
	sql_uc = UniqueConstraint(*sql_names)
	cols.append(sql_uc)
	Table(sql_name(cn), schema_metadata, *cols, comment=str(c.description))
	schema_metadata.create_all(engine)
	return ddl_str

[hrshdhgd]

Interesting! That's awesome but can this be called via CLI ?

gen-sqla schema.yaml

produces the python script by default. I checked --help and I don't see anything that suggests generating DDL output specifically.

[cmungall]

gen-sqltables uses that code!

[hrshdhgd]

Aah! I got it the other way around, sorry! Perfect!!

[cmungall]

you were mostly right - the same underlying logic is shared by both generating the SQLA classes and the DDL!