feat(expr-ir): Implement Acero order_by, hashjoin for over + DataFrame.filter

[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]

TODO: Investigate using non-table_source nodes

• AFAICT, hashjoin should be able to accept things like project
• Defining the handling for cross joins using acero directly seems very achievable
  

  narwhals/narwhals/_arrow/dataframe.py

  Lines 400 to 422 in f4787d3

  	if how == "cross":
  	plx = self.__narwhals_namespace__()
  	key_token = generate_temporary_column_name(
  	n_bytes=8, columns=[*self.columns, *other.columns]
  	)
  	return self._with_native(
  	self.with_columns(
  	plx.lit(0, None).alias(key_token).broadcast(ExprKind.LITERAL)
  	)
  	.native.join(
  	other.with_columns(
  	plx.lit(0, None).alias(key_token).broadcast(ExprKind.LITERAL)
  	).native,
  	keys=key_token,
  	right_keys=key_token,
  	join_type="inner",
  	right_suffix=suffix,
  	)
  	.drop([key_token])
  	)

• There does need to be a new layer for tracking Schema changes
  • Which is needed for with_columns also
  • Generally, the responsibility for a future LogicalPlan

[dangotbanned]

It's interesting I got this far without DataFrame.filter 😂

I do have an acero version, which operates over pc.Expression though

narwhals/narwhals/_plan/arrow/acero.py

Lines 172 to 174 in b0c2a4d

	def filter(*predicates: Expr, **constraints: IntoExpr) -> Decl:
	expr = _parse_all_horizontal(predicates, constraints)
	return Decl("filter", options=pac.FilterNodeOptions(expr))

So the missing link between those two is approximately this:

https://github.com/apache/datafusion-python/blob/e75addfa64a91c0d91ef059917a451e04495b83a/src/pyarrow_filter_expression.rs#L96-L181

But with a fallback to an eager path like main:

narwhals/narwhals/_arrow/dataframe.py

Lines 521 to 529 in 8ac061c

	def filter(self, predicate: ArrowExpr | list[bool | None]) -> Self:
	if isinstance(predicate, list):
	mask_native: Mask | ChunkedArrayAny = predicate
	else:
	# `[0]` is safe as the predicate's expression only returns a single column
	mask_native = self._evaluate_into_exprs(predicate)[0].native
	return self._with_native(
	self.native.filter(mask_native), validate_column_names=False
	)

[dangotbanned]

If anyone see's this - don't copy this logic to solve (#3182)

I haven't handled it here either

Update

Should be clearer now with these failing tests (test: Add test_filter_mask_mixed)

[dangotbanned]

@FBruzzesi #3183 (comment)

This is very nice now 😄