Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 24 additions & 23 deletions rdsa_utils/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,7 @@ def print_full_table_and_raise_error(
ValueError
Raises error and stops pipeline if switch applied.
"""
if show_records & stop_pipeline:
if show_records and stop_pipeline:
logger.error(df.to_string())
logger.error(message)
raise ValueError
Expand Down Expand Up @@ -490,34 +490,35 @@ def wrapper_decorator(*args, **kwargs):
),
)

try:
# Run the decorated function in its normal way, but catch its output
# so it can be counted.
df_return = func(*args, **kwargs)
finally:
# Ensure input df cache is always released, even if func raises
if kwargs.get("df") and isinstance(kwargs["df"], SparkDF):
kwargs["df"].unpersist()

# Check to ensure that the function returns a single value that is a
# spark dataframe, as otherwise the count operation will fail.
if isinstance(df_return, SparkDF):
# Persist the dataframe to be returned prior to counting to allow
# more efficient processing downstream. We persist here as we will
# not be removing this dataframe from the in-memory cache in this
# decorator. Therefore, we don't want it to get pushed onto disk
# (and incur an expensive swap operation).
df_return.persist(StorageLevel.MEMORY_ONLY)
logger.info(
f"Rows in dataframe after {func_name} : {df_return.count()}",
) # noqa: E501

else:
logger.warning(
f"{func_name} should return a spark dataframe for decorator, "
f"but returned {type(df_return)}",
)
# Check to ensure that the function returns a single value that is a
# spark dataframe, as otherwise the count operation will fail.
if isinstance(df_return, SparkDF):
# Persist the dataframe to be returned prior to counting to allow
# more efficient processing downstream. We persist here as we will
# not be removing this dataframe from the in-memory cache in this
# decorator. Therefore, we don't want it to get pushed onto disk
# (and incur an expensive swap operation).
df_return.persist(StorageLevel.MEMORY_ONLY)
logger.info(
f"Rows in dataframe after {func_name} : {df_return.count()}",
) # noqa: E501

if kwargs.get("df"):
# Unpersist the cached input df to manage memory.
kwargs["df"].unpersist()
else:
logger.warning(
f"{func_name} should return a spark dataframe for decorator, "
f"but returned {type(df_return)}",
)

return df_return
return df_return

return wrapper_decorator

Expand Down