diff --git a/src/dataplane/pipelinerun/data_persist/pandas_redis_store.py b/src/dataplane/pipelinerun/data_persist/pandas_redis_store.py index b954313..a5ed961 100644 --- a/src/dataplane/pipelinerun/data_persist/pandas_redis_store.py +++ b/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