feat(expr-ir): Acero order_by, hashjoin , DataFrame.{filter,join}, Expr.is_{first,last}_distinct

[dangotbanned]

Tracking

• Bug: DataFrame.filter silently ignores **constraints when using list[bool] #3182

Related issues

• Child of feat(RFC): A richer Expr IR #2572
• Follow-up to feat(expr-ir): Support group_by, utilize pyarrow.acero #3143

Description

Note

I've used the name sort_by for our wrapped of order_by.
The node corresponds to pa.Table.sort_by, whereas the name order_by doesn't appear anywhere else in pyarrow

Building out more acero parts to be able to support .over(*partition_by)

[dangotbanned]

It's an improvement over the previous version, but far from ideal.

Still doesn't resolve this case, and I'm not entirely sure why yet

narwhals/narwhals/_plan/compliant/column.py

Lines 49 to 60 in f77bb4c

	@classmethod
	def align(
	cls, *exprs: OneOrIterable[SupportsBroadcast[SeriesT, LengthT]]
	) -> Iterator[SeriesT]:
	exprs = tuple[SupportsBroadcast[SeriesT, LengthT], ...](flatten_hash_safe(exprs))
	length = cls._length_required(exprs)
	if length is None:
	for e in exprs:
	yield e.to_series()
	else:
	for e in exprs:
	yield e.broadcast(length)

[dangotbanned]

As of feat(expr-ir): Impl acero.sort_by, I still need to make use of this in a plan.

A good candidate might be in either/both of

over(order_by=...)

narwhals/narwhals/_plan/arrow/expr.py

Lines 328 to 350 in f77bb4c

	def over_ordered(
	self, node: ir.OrderedWindowExpr, frame: Frame, name: str
	) -> Self | Scalar:
	if node.partition_by:
	msg = f"Need to implement `group_by`, `join` for:\n{node!r}"
	raise NotImplementedError(msg)
	# NOTE: Converting `over(order_by=..., options=...)` into the right shape for `DataFrame.sort`
	sort_by = tuple(NamedIR.from_ir(e) for e in node.order_by)
	options = node.sort_options.to_multiple(len(node.order_by))
	idx_name = temp.column_name(frame)
	sorted_context = frame.with_row_index(idx_name).sort(sort_by, options)
	evaluated = node.expr.dispatch(self, sorted_context.drop([idx_name]), name)
	if isinstance(evaluated, ArrowScalar):
	# NOTE: We're already sorted, defer broadcasting to the outer context
	# Wouldn't be suitable for partitions, but will be fine here
	# - https://github.com/narwhals-dev/narwhals/pull/2528/commits/2ae42458cae91f4473e01270919815fcd7cb9667
	# - https://github.com/narwhals-dev/narwhals/pull/2528/commits/b8066c4c57d4b0b6c38d58a0f5de05eefc2cae70
	return self._with_native(evaluated.native, name)
	indices = pc.sort_indices(sorted_context.get_column(idx_name).native)
	height = len(sorted_context)
	result = evaluated.broadcast(height).native.take(indices)
	return self._with_native(result, name)

is_{first,last}_distinct

• test: Port over is_first_distinct tests

narwhals/narwhals/_arrow/series.py

Lines 719 to 747 in 715be22

	def is_first_distinct(self) -> Self:
	import numpy as np # ignore-banned-import
	row_number = pa.array(np.arange(len(self)))
	col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name])
	first_distinct_index = (
	pa.Table.from_arrays([self.native], names=[self.name])
	.append_column(col_token, row_number)
	.group_by(self.name)
	.aggregate([(col_token, "min")])
	.column(f"{col_token}_min")
	)
	return self._with_native(pc.is_in(row_number, first_distinct_index))
	def is_last_distinct(self) -> Self:
	import numpy as np # ignore-banned-import
	row_number = pa.array(np.arange(len(self)))
	col_token = generate_temporary_column_name(n_bytes=8, columns=[self.name])
	last_distinct_index = (
	pa.Table.from_arrays([self.native], names=[self.name])
	.append_column(col_token, row_number)
	.group_by(self.name)
	.aggregate([(col_token, "max")])
	.column(f"{col_token}_max")
	)
	return self._with_native(pc.is_in(row_number, last_distinct_index))

[dangotbanned]

@FBruzzesi #3183 (comment)

This is very nice now 😄