Update pandas_redis_store.py

Author: Saul Frank

src/dataplane/pipelinerun/data_persist/pandas_redis_store.py

@@ -16,8 +16,9 @@ def RedisCheck(r):
 DataFrame: Pandas dataframe to pass
 Expire: Expires the data if true.
 ExpireDuration: If expires is true, how much time to expire. Default 15 mins
+Compression: "gzip", "bz2", 'zstd', 'tar', "zip",  None
 """
-def pipeline_pandas_redis_store(StoreKey, DataFrame, Redis, Expire=True, ExpireDuration=timedelta(minutes=15)):
+def pipeline_pandas_redis_store(StoreKey, DataFrame, Redis, Expire=True, ExpireDuration=timedelta(minutes=15), Compression="gzip"):
 
     import os
     import io
@@ -33,7 +34,7 @@ def pipeline_pandas_redis_store(StoreKey, DataFrame, Redis, Expire=True, ExpireD
         raise Exception("Redis connection failed.")
 
     buffer = io.BytesIO()
-    DataFrame.to_pickle(buffer, compression='gzip')
+    DataFrame.to_pickle(buffer, compression=Compression)
     buffer.seek(0) # re-set the pointer to the beginning after reading
 
     if Expire:
@@ -50,7 +51,7 @@ def pipeline_pandas_redis_store(StoreKey, DataFrame, Redis, Expire=True, ExpireD
 StoreKey: is the key to look up for retrieval (set with RedisStore). 
 Redis: e.g. Redis = redis.Redis(host='redis-service', port=6379, db=0)
 """
-def pipeline_pandas_redis_get(StoreKey, Redis):
+def pipeline_pandas_redis_get(StoreKey, Redis, Compression="gzip"):
 
     import os
     import io
@@ -69,7 +70,7 @@ def pipeline_pandas_redis_get(StoreKey, Redis):
     buffer = io.BytesIO(Redis.get(InsertKey))
     buffer.seek(0)
     import pandas as pd
-    df = pd.read_pickle(buffer,compression='gzip')
+    df = pd.read_pickle(buffer,compression=Compression)
 
     duration = datetime.now() - start