only repartition big dfs
@@ -138,6 +138,11 @@ class SparkLinker(Linker):
if self.break_lineage_after_blocking:
regex_to_persist.append(r"__splink__df_blocked")
+ names_to_repartition = (
+ "__splink__df_comparison_vectors",
+ "__splink__df_blocked",
+ )
+
if re.match(r"|".join(regex_to_persist), templated_name):
if self.break_lineage_method == "persist":
@@ -147,14 +152,14 @@ class SparkLinker(Linker):
spark_df = spark_df.persist(self.persist_level)
logger.debug(f"persisted {templated_name}")
elif self.break_lineage_method == "checkpoint":
- if templated_name == "__splink__df_comparison_vectors":
+ if templated_name in names_to_repartition:
spark_df = spark_df.repartition(self.num_partitions_on_repartition)
spark_df = spark_df.checkpoint()
logger.debug(f"Checkpointed {templated_name}")
elif self.break_lineage_method == "parquet":
checkpoint_dir = self.spark.sparkContext.getCheckpointDir()
write_path = os.path.join(checkpoint_dir, physical_name)
-
+ if templated_name in names_to_repartition:
spark_df = spark_df.repartition(self.num_partitions_on_repartition)
spark_df.write.mode("overwrite").parquet(write_path)
