Misc dynamodb integration test improvements:

 - DELTA_DYNAMO_TABLE_OVERWRITE env variable
 - add run_id column
 - cleanup

Author: mrk-its

storage-dynamodb/integration_tests/dynamodb_logstore.py

@@ -16,26 +16,17 @@
 import os
 import sys
 import threading
+import random
+import uuid
 
 from pyspark import SparkContext
-from pyspark.sql import Column, DataFrame, SparkSession, SQLContext, functions
-from pyspark.sql.functions import *
+from pyspark.sql import Column, DataFrame, SparkSession, SQLContext, functions as F
 from py4j.java_collections import MapConverter
 from delta.tables import *
 from multiprocessing.pool import ThreadPool
 import time
 
 """
-create required dynamodb table with:
-
-$ aws --region us-west-2 dynamodb create-table \
-    --table-name delta_log_test \
-    --attribute-definitions AttributeName=tablePath,AttributeType=S \
-                            AttributeName=fileName,AttributeType=S \
-    --key-schema AttributeName=tablePath,KeyType=HASH \
-                AttributeName=fileName,KeyType=RANGE \
-    --provisioned-throughput ReadCapacityUnits=5,WriteCapacityUnits=5
-
 run this script in root dir of repository:
 
 export VERSION=$(cat version.sbt|cut -d '"' -f 2)
@@ -44,7 +35,7 @@
 export DELTA_TABLE_PATH=s3a://test-bucket/delta-test/
 export DELTA_DYNAMO_TABLE=delta_log_test
 export DELTA_DYNAMO_REGION=us-west-2
-export DELTA_STORAGE=io.delta.storage.DynamoDBLogStoreScala # TODO: remove `Scala` when Java version finished
+export DELTA_STORAGE=io.delta.storage.DynamoDBLogStore
 export DELTA_NUM_ROWS=16
 
 ./run-integration-tests.py --run-storage-dynamodb-integration-tests \
@@ -59,11 +50,11 @@
 concurrent_readers = int(os.environ.get("DELTA_CONCURRENT_READERS", 2))
 num_rows = int(os.environ.get("DELTA_NUM_ROWS", 16))
 
-# TODO change back to default io.delta.storage.DynamoDBLogStore
-delta_storage = os.environ.get("DELTA_STORAGE", "io.delta.storage.DynamoDBLogStoreScala")
+delta_storage = os.environ.get("DELTA_STORAGE", "io.delta.storage.DynamoDBLogStore")
 dynamo_table_name = os.environ.get("DELTA_DYNAMO_TABLE", "delta_log_test")
 dynamo_region = os.environ.get("DELTA_DYNAMO_REGION", "us-west-2")
 dynamo_error_rates = os.environ.get("DELTA_DYNAMO_ERROR_RATES", "")
+table_overwrite = os.environ.get("DELTA_DYNAMO_TABLE_OVERWRITE", "true").lower() == "true"
 
 if delta_table_path is None:
     print(f"\nSkipping Python test {os.path.basename(__file__)} due to the missing env variable "
@@ -90,24 +81,41 @@
     .master("local[*]") \
     .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") \
     .config("spark.delta.logStore.class", delta_storage) \
-    .config("spark.delta.DynamoDBLogStoreScala.tableName", dynamo_table_name) \
-    .config("spark.delta.DynamoDBLogStoreScala.region", dynamo_region) \
-    .config("spark.delta.DynamoDBLogStoreScala.errorRates", dynamo_error_rates) \
+    .config("spark.delta.DynamoDBLogStore.tableName", dynamo_table_name) \
+    .config("spark.delta.DynamoDBLogStore.region", dynamo_region) \
+    .config("spark.delta.DynamoDBLogStore.errorRates", dynamo_error_rates) \
     .getOrCreate()
 
-data = spark.createDataFrame([], "id: int, a: int")
-data.write.format("delta").mode("overwrite").partitionBy("id").save(delta_table_path)
+SCHEMA = "run_id: string, id: int, a: int"
+
+RUN_ID = str(uuid.uuid4())
+
+data = spark.createDataFrame([], SCHEMA)
+
+if table_overwrite:
+    data.write.format("delta").mode("overwrite").partitionBy("run_id", "id").save(delta_table_path)
+
 
 def write_tx(n):
-    data = spark.createDataFrame([[n, n]], "id: int, a: int")
-    data.write.format("delta").mode("append").partitionBy("id").save(delta_table_path)
+    data = spark.createDataFrame([[RUN_ID, random.randrange(2**16), n]], SCHEMA)
+    data.write.format("delta").mode("append").partitionBy("run_id", "id").save(delta_table_path)
+
+
+def count():
+    return (
+        spark.read.format("delta")
+        .load(delta_table_path)
+        .filter(F.col("run_id") == RUN_ID)
+        .count()
+    )
 
 
 stop_reading = threading.Event()
 
 def read_data():
     while not stop_reading.is_set():
-        print("Reading {:d} rows ...".format(spark.read.format("delta").load(delta_table_path).distinct().count()))
+        cnt = count()
+        print(f"Reading {cnt} rows ...")
         time.sleep(1)
 
 
@@ -127,7 +135,7 @@ def start_read_thread():
 for thread in read_threads:
     thread.join()
 
-actual = spark.read.format("delta").load(delta_table_path).distinct().count()
+actual = count()
 print("Number of written rows:", actual)
 assert actual == num_rows