feat: Deletion vectors (#19)

Author: simonpokorny

README.md

@@ -226,9 +226,10 @@ def spark():
 
 At the beginning of your test method call `reinit_local_metastore` function 
 from the testing package to initialize the metastore with the tables from 
-your json folder (`JSON_TABLES_DIR`). If the method is called while the 
-metastore  already exists, it will delete all the existing tables before 
-initializing the new ones.
+your json folder (`JSON_TABLES_DIR`). You can also choose to enable or disable
+ deletion vectors for Delta tables (default: enabled). If the method is called
+  while the metastore already exists, it will delete all the existing tables
+   before initializing the new ones.
 
 *Alternatively, you can call this method only once per testing module, 
 but then individual testing methods might affect each other by modifying 
@@ -242,7 +243,7 @@ from pyspark.testing import assertDataFrameEqual
 def test_process_data(
     spark: SparkSession,
 ):
-    reinit_local_metastore(spark, JSON_TABLES_DIR)
+    reinit_local_metastore(spark, JSON_TABLES_DIR, deletion_vectors=True)
 
     process_data(
         spark=spark,

pysparkdt/metastore.py

@@ -8,6 +8,7 @@
 def reinit_local_metastore(
     spark: SparkSession,
     json_tables_dir: str,
+    deletion_vectors: bool = True,
 ) -> None:
     """Re-initializes dynamic metastore acting as Databricks data catalog
     using provided input delta table data in json format.
@@ -51,6 +52,9 @@ def reinit_local_metastore(
         Local Spark session.
     json_tables_dir
         Directory where the delta tables and their schemas are located.
+    deletion_vectors
+        Whether to enable deletion vectors for the delta tables.
+        Defaults to True.
     """
     # Clear all existing tables (must be done through SQL, not by clearing the
     # folder)
@@ -87,4 +91,14 @@ def reinit_local_metastore(
             query = query.option('inferSchema', True)
         df = query.load(data_path)
 
-        df.write.format('delta').saveAsTable(table_name)
+        write_query = df.write.format('delta')
+
+        if deletion_vectors:
+            write_query = write_query.option(
+                'delta.enableDeletionVectors', 'true'
+            )
+        else:
+            write_query = write_query.option(
+                'delta.enableDeletionVectors', 'false'
+            )
+        write_query.saveAsTable(table_name)

tests/test_deletion_vectors.py

@@ -0,0 +1,38 @@
+import os
+
+from pyspark.sql import SparkSession
+from pytest import fixture
+
+from pysparkdt import reinit_local_metastore, spark_base
+
+DATA_DIR = f'{os.path.dirname(__file__)}/data'
+JSON_TABLES_DIR = f'{DATA_DIR}/tables'
+TMP_DIR = f'{DATA_DIR}/tmp'
+METASTORE_DIR = f'{TMP_DIR}/metastore'
+
+
+@fixture(scope='module')
+def spark():
+    yield from spark_base(METASTORE_DIR)
+
+
+def test_deletion_vectors_disabled(spark: SparkSession):
+    """Test that deletion vectors are disabled when deletion_vectors=False"""
+    reinit_local_metastore(spark, JSON_TABLES_DIR, deletion_vectors=False)
+
+    # Check if deletion vectors are disabled for the table
+    table_properties = spark.sql('DESCRIBE DETAIL example_input').collect()[0]
+    properties = table_properties.properties
+
+    assert properties.get('delta.enableDeletionVectors') == 'false'
+
+
+def test_deletion_vectors_enabled(spark: SparkSession):
+    """Test that deletion vectors are enabled when deletion_vectors=True"""
+    reinit_local_metastore(spark, JSON_TABLES_DIR, deletion_vectors=True)
+
+    # Check if deletion vectors are enabled for the table
+    table_properties = spark.sql('DESCRIBE DETAIL example_input').collect()[0]
+    properties = table_properties.properties
+
+    assert properties.get('delta.enableDeletionVectors') == 'true'