diff --git a/pyproject.toml b/pyproject.toml index 78f2f4a..77e7e75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "lakebench" -version = "1.1.0" +version = "1.2.0" authors = [ { name="Miles Cole" }, ] diff --git a/src/lakebench/engines/fabric_spark.py b/src/lakebench/engines/fabric_spark.py index 807c847..3a94ccf 100644 --- a/src/lakebench/engines/fabric_spark.py +++ b/src/lakebench/engines/fabric_spark.py @@ -26,6 +26,7 @@ def __init__( cost_per_vcore_hour: Optional[float] = None, collect_stats_on_write: bool = True, compute_stats_all_cols: Optional[bool] = None, + tblproperties: Optional[dict] = None, ): """ Parameters @@ -43,6 +44,8 @@ def __init__( Whether Fabric Delta extended statistics should be collected during write operations. compute_stats_all_cols : bool, optional Deprecated alias for ``collect_stats_on_write``. When provided, it takes precedence. + tblproperties : dict, optional + Delta table properties to inject into CREATE TABLE statements. """ collect_stats_on_write = self._resolve_collect_stats_on_write( collect_stats_on_write=collect_stats_on_write, @@ -55,6 +58,7 @@ def __init__( spark_measure_telemetry=spark_measure_telemetry, cost_per_vcore_hour=cost_per_vcore_hour, compute_stats_all_cols=False, + tblproperties=tblproperties, ) self.collect_stats_on_write = collect_stats_on_write diff --git a/src/lakebench/engines/hdi_spark.py b/src/lakebench/engines/hdi_spark.py index 210e5c2..bf2a8d3 100644 --- a/src/lakebench/engines/hdi_spark.py +++ b/src/lakebench/engines/hdi_spark.py @@ -9,7 +9,11 @@ class HDISpark(Spark): """ def __init__( - self, schema_name: str, spark_measure_telemetry: bool = False, cost_per_vcore_hour: Optional[float] = None + self, + schema_name: str, + spark_measure_telemetry: bool = False, + cost_per_vcore_hour: Optional[float] = None, + tblproperties: Optional[dict] = None, ): """ Parameters @@ -21,6 +25,8 @@ def __init__( cost_per_vcore_hour : float, optional The cost per vCore hour for the Spark cluster. If None, cost calculations are auto calculated where possible. + tblproperties : dict, optional + Delta table properties to inject into CREATE TABLE statements. """ super().__init__( @@ -29,4 +35,5 @@ def __init__( spark_measure_telemetry=spark_measure_telemetry, cost_per_vcore_hour=cost_per_vcore_hour, compute_stats_all_cols=False, + tblproperties=tblproperties, ) diff --git a/src/lakebench/engines/spark.py b/src/lakebench/engines/spark.py index b52a3e8..b5c5d53 100644 --- a/src/lakebench/engines/spark.py +++ b/src/lakebench/engines/spark.py @@ -46,6 +46,7 @@ def __init__( spark_measure_telemetry: bool = False, cost_per_vcore_hour: Optional[float] = None, compute_stats_all_cols: bool = False, + tblproperties: Optional[dict] = None, ): """ Parameters @@ -153,6 +154,7 @@ def __init__( self.compute_stats_all_cols = compute_stats_all_cols self.run_analyze_after_load = self.compute_stats_all_cols + self.tblproperties = tblproperties if tblproperties is not None else {} def __get_spark_session_configs(self) -> dict: """ @@ -240,6 +242,32 @@ def _create_empty_table(self, table_name: Optional[str], ddl: str): ) ddl = create_node.sql(dialect="spark", pretty=True) + if self.tblproperties: + import sqlglot + + expression = sqlglot.parse_one(ddl, dialect="spark") + create_node = expression.find(sqlglot.exp.Create) + if create_node is not None: + existing_props = create_node.args.get("properties") + existing_exprs = existing_props.expressions if existing_props else [] + # Prepend so TBLPROPERTIES appears after USING + create_node.set( + "properties", + sqlglot.exp.Properties( + expressions=[ + *existing_exprs, + *[ + sqlglot.exp.Property( + this=sqlglot.exp.Literal.string(k), + value=sqlglot.exp.Literal.string(v), + ) + for k, v in self.tblproperties.items() + ], + ] + ), + ) + ddl = create_node.sql(dialect="spark", pretty=True) + self.execute_sql_statement(ddl) def _convert_generic_to_specific_schema(self, generic_schema: list): diff --git a/src/lakebench/engines/synapse_spark.py b/src/lakebench/engines/synapse_spark.py index 8c10d50..3d6d8b6 100644 --- a/src/lakebench/engines/synapse_spark.py +++ b/src/lakebench/engines/synapse_spark.py @@ -15,6 +15,7 @@ def __init__( schema_uri: Optional[str] = None, spark_measure_telemetry: bool = False, cost_per_vcore_hour: Optional[float] = None, + tblproperties: Optional[dict] = None, ): """ Parameters @@ -28,6 +29,8 @@ def __init__( cost_per_vcore_hour : float, optional The cost per vCore hour for the Spark cluster. If None, cost calculations are auto calculated where possible. + tblproperties : dict, optional + Delta table properties to inject into CREATE TABLE statements. """ super().__init__( @@ -37,6 +40,7 @@ def __init__( spark_measure_telemetry=spark_measure_telemetry, cost_per_vcore_hour=cost_per_vcore_hour, compute_stats_all_cols=False, + tblproperties=tblproperties, ) if self.runtime != "synapse": diff --git a/tests/test_engine.py b/tests/test_engine.py index 5ea6c53..cc44b33 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -120,6 +120,7 @@ def _make_spark_engine(): engine.fs = None engine.runtime = "local_unknown" engine.operating_system = engine._detect_os() + engine.tblproperties = {} original_execute = ( engine.execute_sql_statement.__func__ if hasattr(engine.execute_sql_statement, "__func__") else None @@ -176,6 +177,22 @@ def test_existing_using_parquet_preserved(self): assert "using parquet" in result assert "using delta" not in result + def test_tblproperties_are_injected_after_using_delta(self): + engine = _make_spark_engine() + engine.tblproperties = { + "delta.enableDeletionVectors": "false", + "delta.targetFileSize": "134217728", + } + + engine._create_empty_table(table_name="t", ddl="CREATE TABLE t (id INT)") + + result = engine.executed_statements[0] + assert "USING DELTA" in result + assert "TBLPROPERTIES" in result + assert "'delta.enableDeletionVectors'='false'" in result.replace(" ", "") + assert "'delta.targetFileSize'='134217728'" in result.replace(" ", "") + assert result.index("USING DELTA") < result.index("TBLPROPERTIES") + class TestSparkAnalyzeTable: def test_full_analysis_uses_all_columns(self): diff --git a/tests/test_spark_subclasses.py b/tests/test_spark_subclasses.py new file mode 100644 index 0000000..1a5ff61 --- /dev/null +++ b/tests/test_spark_subclasses.py @@ -0,0 +1,40 @@ +import pytest + +from lakebench.engines.fabric_spark import FabricSpark +from lakebench.engines.hdi_spark import HDISpark +from lakebench.engines.spark import Spark +from lakebench.engines.synapse_spark import SynapseSpark + + +class _ParentInitCalled(Exception): + pass + + +@pytest.mark.parametrize( + ("engine_class", "constructor_kwargs"), + [ + ( + FabricSpark, + { + "lakehouse_name": "lakehouse", + "lakehouse_schema_name": "schema", + }, + ), + (HDISpark, {"schema_name": "schema"}), + (SynapseSpark, {"schema_name": "schema"}), + ], +) +def test_spark_subclasses_forward_tblproperties(monkeypatch, engine_class, constructor_kwargs): + tblproperties = {"delta.enableDeletionVectors": "false"} + captured_kwargs = {} + + def capture_parent_init(self, **kwargs): + captured_kwargs.update(kwargs) + raise _ParentInitCalled + + monkeypatch.setattr(Spark, "__init__", capture_parent_init) + + with pytest.raises(_ParentInitCalled): + engine_class(**constructor_kwargs, tblproperties=tblproperties) + + assert captured_kwargs["tblproperties"] is tblproperties diff --git a/uv.lock b/uv.lock index ee44a5c..46be560 100644 --- a/uv.lock +++ b/uv.lock @@ -908,7 +908,7 @@ wheels = [ [[package]] name = "lakebench" -version = "1.1.0" +version = "1.2.0" source = { editable = "." } dependencies = [ { name = "fsspec" },