[Data] Add incremental join to Ray data

[owenowenisme]

Description

Currently ray data materialize both side of the table to do join, but we can actually build one part and incrementally build output, this can reduce the memory usage and prevent the whole output table in our worker.

We can do that after pyarrow expose the reusability of the build table to prevent rebuild cost every time a new record batch(for us it's probe side), this way we can control the memory usage when still being efficient.

See arrow PR apache/arrow#50802

Current way:

ray/python/ray/data/_internal/execution/operators/join.py

Lines 169 to 234 in 7445e1f

	def join_tables(
	left_table: "pa.Table",
	right_table: "pa.Table",
	*,
	join_type: JoinType,
	left_key_col_names: Tuple[str, ...],
	right_key_col_names: Tuple[str, ...],
	left_columns_suffix: Optional[str] = None,
	right_columns_suffix: Optional[str] = None,
	) -> "pa.Table":
	"""Apply preprocess -> ``pa.Table.join`` -> postprocess to two input tables.
	Shared between the physical executor (``JoiningAggregation.finalize``)
	and plan-time schema inference (``Join.infer_schema``), which calls
	this with empty tables built from the input schemas. Plan-time and
	runtime schemas therefore agree by construction.
	"""
	left_on = list(left_key_col_names)
	right_on = list(right_key_col_names)
	# Eagerly validate suffix conflicts so callers get a clear error instead
	# of the opaque PyArrow schema-merge error ('Field X exists 2 times').
	# Skip for semi/anti joins: only one side's columns appear in the result,
	# so overlapping non-key names between left and right are harmless.
	if join_type not in (
	JoinType.LEFT_SEMI,
	JoinType.LEFT_ANTI,
	JoinType.RIGHT_SEMI,
	JoinType.RIGHT_ANTI,
	):
	left_cols = set(left_table.schema.names)
	# PyArrow drops right key columns from output (coalescing them into
	# the left keys), so only right non-key columns can collide with
	# left columns. Subtracting only right_on (not left_on) correctly
	# handles asymmetric key names (left_on != right_on).
	right_output_cols = set(right_table.schema.names) - set(right_on)
	collisions = left_cols & right_output_cols
	if left_columns_suffix is None and right_columns_suffix is None and collisions:
	raise ValueError(
	"Left and right columns suffixes cannot be both None "
	f"(overlapping columns: {sorted(collisions)})"
	)
	# Preprocess: split unsupported columns and add index columns if needed
	preprocess_result_l, preprocess_result_r = _preprocess(
	left_table, right_table, left_on, right_on, join_type
	)
	# Perform the join on supported columns
	arrow_join_type = _JOIN_TYPE_TO_ARROW_JOIN_VERB_MAP[join_type]
	supported = preprocess_result_l.supported_projection.join(
	preprocess_result_r.supported_projection,
	join_type=arrow_join_type,
	keys=left_on,
	right_keys=right_on,
	left_suffix=left_columns_suffix,
	right_suffix=right_columns_suffix,
	)
	# Add back unsupported columns
	return _postprocess(
	supported,
	preprocess_result_l.unsupported_projection,
	preprocess_result_r.unsupported_projection,
	)

Use case

No response

[odncode]

Hi @owenowenisme, Am I okay to take this on, or will you be handling the implementation?

It looks like you've basically done most of the heavy lifting on the logic, and I can implement the changes if you prefer.

[hoamgh]

Hi @owenowenisme ,

This sounds like a great task! Would it be alright if I pick this up and work on the implementation for the incremental join? Thank you!