fix cross-shard ORDER BY on non-projected columns

[murex971]

• Temporarily project missing ORDER BY columns during query rewriting.
• Use hidden columns for correct global cross-shard sorting.
• Remove hidden columns before sending rows and RowDescription to clients.
• Support direct-shard results.
• closes [Query Engine] ORDER BY on non-projected column can produce incorrect cross-shard ordering #1135

[codecov]

Codecov Report

❌ Patch coverage is 98.51485% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ontend/router/parser/rewrite/statement/order_by.rs	97.12%	4 Missing ⚠️
pgdog/src/frontend/router/parser/query/select.rs	96.96%	1 Missing ⚠️
pgdog/src/net/messages/bind.rs	90.90%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[levkk]

I would probably push this concern into the connection binding (e.g.. read method). The architecture is as such that piece of code handles the cross-shard/direct-to-shard message manipulation. The query engine doesn't actually know whether it's talking to one or more servers, so this leaks it here.

[sgrif]

It might be out of scope for this PR, but a refactor I've been wanting to make is to pass through the original target list of the select statement, and changing the code to keep those columns rather than keep track of which ones to drop. This could probably even just be a length instead of a list since we're always going to append any helpers to the end of the list

[sgrif]

Suggested change

	let mut order_by = Self::select_sort(stmt, context.router_context.bind);
	for helper in cached_ast.rewrite_plan.projection.order_by_helpers() {
	let Some((_, column)) = order_by
	.iter_mut()
	.find(|(position, _)| *position == helper.sort_position)
	else {
	continue;
	};
	*column = if column.asc() {
	OrderBy::Asc(helper.projected_column + 1)
	} else {
	OrderBy::Desc(helper.projected_column + 1)
	};
	}
	let order_by = order_by
	.into_iter()
	.map(|(_, order_by)| order_by)
	.collect::<Vec<_>>();
	let helpers = cached_ast.rewrite_plan.projection.order_by_helpers();
	let order_by = Self::select_sort(stmt, router_context.bind)
	.into_iter()
	.map(|(position, column)| {
	if let Some(helper) = order_by.iter().find(|helper| position == helper.sort_position) {
	if column.asc() {
	OrderBy::Asc(helper.projected_column + 1)
	} else {
	OrderBy::Desc(helper.projected_column + 1)
	}
	} else {
	column
	}
	})
	.collect::<Vec<_>>();

[sgrif]

While we're touching this signature, what do you think about changing this to return impl Iterator instead?

[sgrif]

We can skip an intermediate vec here.

Suggested change

	})
	.collect::<Vec<_>>();
	if helpers.is_empty() {
	return;
	}
	let (targets, mappings): (Vec<_>, Vec<_>) = helpers.into_iter().unzip();
	});
	let (targets, mappings): (Vec<_>, Vec<_>) = helpers.unzip();
	if targets.is_empty() {
	return;
	}

[sgrif]

What do you think about this instead? The way qualifiers is constructed feels noisy to me, and we can skip the intermediate vec.

Suggested change

	let fields = column_ref.fields();
	if !matches!(fields.iter().next_back(), Some(Node::A_Star(_))) {
	return false;
	}
	let qualifiers = fields
	.iter()
	.take(fields.len().saturating_sub(1))
	.filter_map(Node::as_str)
	.collect::<Vec<_>>();
	match qualifiers.as_slice() {
	[] => true,
	[table] => order_by.table == Some(*table),
	[schema, table] => order_by.schema == Some(*schema) && order_by.table == Some(*table),
	_ => false,
	}
	let mut fields = column_ref.fields().iter();
	if !matches!(fields.next_back(), Some(Node::A_Star(_))) {
	return false;
	}
	let table = fields.next_back();
	let schema = fields.next_back();
	match (schema, table) {
	(None, None) => true,
	(Some(table), None) => order_by.table == Some(*table),
	(schema, table) => order_by.schema == schema && order_by.table == table,
	}

[murex971]

@sgrif I am closing this and handling it in #1545
will handle suggested fixes there