[IMP] Decrease Merge Peak Memory Usage of ConnectedComponents

[VibhuJawa]

Describe the bug

On smaller GPU skews we are running into memory issues in the broadcast merge in Connected Components. We have to decrease that memory footprint without hurting performance too much.

NeMo-Curator/nemo_curator/modules/fuzzy_dedup.py

Lines 1605 to 1655 in f7441ea

	def _write_encoded_jaccard_pair(self, dedup_parsed_id_path):
	output_path = f"{self.cache_dir}/encoded_jaccard_pair/"
	t0 = time.time()
	with performance_report_if_with_ts_suffix(
	self.profile_dir, "connected-components-encoded-jaccard-pair"
	):
	ddf_id = dask_cudf.read_parquet(
	dedup_parsed_id_path, blocksize="2GB", aggregate_files=True
	)
	ddf = dask_cudf.read_parquet(
	self.jaccard_pairs_path,
	blocksize="1GB",
	aggregate_files=True,
	)
	self._merge_and_write(
	ddf=ddf,
	ddf_id=ddf_id,
	output_path=output_path,
	id_column=self.id_column,
	)
	self._logger.info(
	f"Time taken for Encoding Jaccard Pairs = {time.time() - t0}s and output written at {output_path}"
	)
	return output_path
	def _merge_and_write(
	self,
	ddf: dask_cudf.DataFrame,
	ddf_id: dask_cudf.DataFrame,
	output_path: str,
	id_column: str,
	) -> None:
	st = time.time()
	# Ensure 'id_columns' is a list
	ddf_id = ddf_id.set_index(id_column)
	for tag in ["x", "y"]:
	pair_id = f"{id_column}_{tag}"
	# Merge 'ddf' with 'ddf_id' to map ids to uids
	ddf = ddf.merge(
	ddf_id,
	left_on=pair_id,
	right_index=True,
	how="inner",
	broadcast=True,
	)
	ddf = ddf.drop(columns=pair_id)
	ddf = ddf.rename(columns={"uid": f"{self.id_column}_{tag}"})
	ddf = ddf[[self.left_id, self.right_id, "jaccard"]]
	ddf.to_parquet(output_path, write_index=False, overwrite=True)

CC: @ayushdg