1use super::{DOC_ID_COLUMN, TABLE_OID_COLUMN};
10use crate::ast::ReturningAliases;
11use crate::plan::ProjectionPlan;
12use crate::{ColumnIdentity, RowSchema, SQLError, ScalarExpr};
13
14pub fn projection_columns(projections: &[ProjectionPlan]) -> Vec<String> {
15 projections.iter().map(projection_label_at).collect()
16}
17
18pub fn projection_label_at(proj: &ProjectionPlan) -> String {
20 if let Some(a) = &proj.alias {
21 return a.clone();
22 }
23 match &proj.expr {
24 ScalarExpr::Column(c) => c.clone(),
25 ScalarExpr::QualifiedColumn { column, .. } => column.clone(),
26 ScalarExpr::Star | ScalarExpr::QualifiedStar(_) => "*".into(),
27 ScalarExpr::Func {
28 name,
29 binding,
30 args,
31 ..
32 } => field_selection_label(binding.as_ref(), args.get(1))
33 .unwrap_or_else(|| function_projection_label(name, binding.as_ref())),
34 ScalarExpr::WindowCall { name, .. } => function_projection_label(name, None),
35 _ => "?column?".into(),
36 }
37}
38
39pub(crate) fn field_selection_label<E: FieldNameLiteral>(
41 binding: Option<&crate::ast::FunctionBinding>,
42 field: Option<&E>,
43) -> Option<String> {
44 (binding.and_then(|binding| binding.dispatch)
45 == Some(crate::ast::FunctionDispatch::FieldSelect))
46 .then(|| field.and_then(FieldNameLiteral::field_name))
47 .flatten()
48 .map(str::to_owned)
49}
50
51pub(crate) trait FieldNameLiteral {
53 fn field_name(&self) -> Option<&str>;
54}
55
56impl FieldNameLiteral for ScalarExpr {
57 fn field_name(&self) -> Option<&str> {
58 match self {
59 ScalarExpr::Literal(uqa_core::Value::Str(field)) => Some(field),
60 _ => None,
61 }
62 }
63}
64
65impl FieldNameLiteral for crate::ast::Expr {
66 fn field_name(&self) -> Option<&str> {
67 match self {
68 crate::ast::Expr::Literal(uqa_core::Value::Str(field)) => Some(field),
69 _ => None,
70 }
71 }
72}
73
74pub fn function_projection_label(
75 name: &str,
76 binding: Option<&crate::ast::FunctionBinding>,
77) -> String {
78 if matches!(
80 binding.and_then(|binding| binding.dispatch),
81 Some(
82 crate::ast::FunctionDispatch::NumericOperator(_)
83 | crate::ast::FunctionDispatch::ArrayConcat
84 | crate::ast::FunctionDispatch::ArrayAppend
85 | crate::ast::FunctionDispatch::ArrayPrepend
86 | crate::ast::FunctionDispatch::IsDistinct
87 | crate::ast::FunctionDispatch::AnyOperator
88 | crate::ast::FunctionDispatch::AllOperator
89 | crate::ast::FunctionDispatch::BetweenSymmetric
90 | crate::ast::FunctionDispatch::JsonExtract { .. }
91 )
92 ) {
93 return "?column?".into();
94 }
95 let name = crate::parse_regobject_name(name)
96 .and_then(|mut names| names.pop())
97 .unwrap_or_else(|| name.to_string());
98 if operator_call(&name) {
99 return "?column?".into();
100 }
101 name
102}
103
104fn operator_call(name: &str) -> bool {
106 matches!(
107 name,
108 "concat_op"
109 | "fts_match"
110 | "jsonpath_exists"
111 | "contains_op"
112 | "contained_by_op"
113 | "array_overlap"
114 | "json_has_key"
115 | "json_has_any_key"
116 | "json_has_all_keys"
117 | "json_delete_path"
118 | "like"
119 | "ilike"
120 | "similar_to"
121 | "regex_match_op"
122 | "regex_imatch_op"
123 )
124}
125
126pub fn expand_from_star_columns(
127 columns: Vec<String>,
128 projections: &[ProjectionPlan],
129 source_schema: &RowSchema,
130) -> Result<Vec<String>, SQLError> {
131 let mut output = Vec::new();
132 for (position, projection) in projections.iter().enumerate() {
133 match &projection.expr {
134 ScalarExpr::Star => {
135 output.extend(
136 source_schema
137 .columns()
138 .iter()
139 .enumerate()
140 .filter(|(position, _)| {
141 visible_projection_source_position(source_schema, *position)
142 })
143 .map(|(source_position, column)| {
144 source_schema
145 .public_name(source_position)
146 .unwrap_or(column)
147 .to_string()
148 }),
149 );
150 }
151 ScalarExpr::QualifiedStar(qualifier) => {
152 let qualified_columns = source_schema
153 .qualified_star_position_layout(qualifier)
154 .into_iter()
155 .filter(|(_, logical, _, _)| {
156 logical.is_none_or(|position| {
157 visible_projection_source_position(source_schema, position)
158 })
159 })
160 .map(|(column, _, _, _)| column)
161 .collect::<Vec<_>>();
162 if qualified_columns.is_empty() {
163 return Err(SQLError::UnknownTable(qualifier.clone()));
164 }
165 output.extend(qualified_columns);
166 }
167 _ => output.push(columns[position].clone()),
168 }
169 }
170 Ok(output)
171}
172
173pub fn bound_projection_expression(schema: &RowSchema, position: usize) -> ScalarExpr {
174 let Some(identity) = schema.identity(position) else {
175 return ScalarExpr::Position(position);
176 };
177 if let Some(qualifier) = identity.qualifier() {
178 if schema.qualified_position(qualifier, identity.column()) == Some(position) {
179 return ScalarExpr::qualified_column(qualifier, identity.column());
180 }
181 } else if schema.unqualified_position(identity.column()) == Some(position) {
182 return ScalarExpr::Column(identity.column().to_string());
183 }
184 ScalarExpr::Position(position)
185}
186
187pub fn expand_bound_projection_stars(
188 projections: &[ProjectionPlan],
189 schema: &RowSchema,
190) -> Result<Vec<ProjectionPlan>, SQLError> {
191 let mut expanded = Vec::new();
192 for projection in projections {
193 match &projection.expr {
194 ScalarExpr::Star => {
195 for (position, column) in schema.columns().iter().enumerate() {
196 if !visible_projection_source_position(schema, position) {
197 continue;
198 }
199 expanded.push(ProjectionPlan {
200 expr: bound_projection_expression(schema, position),
201 alias: Some(schema.public_name(position).unwrap_or(column).to_string()),
202 });
203 }
204 }
205 ScalarExpr::QualifiedStar(qualifier) => {
206 let layout = schema.qualified_star_position_layout(qualifier);
207 if layout.is_empty() {
208 return Err(SQLError::UnknownTable(qualifier.clone()));
209 }
210 for (column, logical, _, _) in layout {
211 if logical.is_some_and(|position| {
212 !visible_projection_source_position(schema, position)
213 }) {
214 continue;
215 }
216 expanded.push(ProjectionPlan {
217 expr: logical.map_or_else(
218 || ScalarExpr::qualified_column(qualifier, &column),
219 |position| bound_projection_expression(schema, position),
220 ),
221 alias: Some(column),
222 });
223 }
224 }
225 _ => expanded.push(projection.clone()),
226 }
227 }
228 Ok(expanded)
229}
230
231pub fn visible_projection_source_position(schema: &RowSchema, position: usize) -> bool {
232 schema.wildcard_position_visible(position)
233}
234
235pub fn returning_context_schema(
236 columns: &[String],
237 types: &[Option<crate::ast::ColumnType>],
238 composite_width: usize,
239 target_qualifier: &str,
240 aliases: &ReturningAliases,
241) -> RowSchema {
242 let target =
243 RowSchema::with_qualified_types(target_qualifier, columns.to_vec(), types.to_vec());
244 let target = RowSchema::with_wildcard_hidden_positions(&target, composite_width..columns.len());
245 let hidden_types = types
246 .iter()
247 .cloned()
248 .chain(types.iter().cloned())
249 .collect::<Vec<_>>();
250 let schema = RowSchema::append_hidden_typed(&target, &hidden_types);
251 let width = columns.len();
252 let identity_aliases = columns
253 .iter()
254 .enumerate()
255 .flat_map(|(position, column)| {
256 [
257 (
258 ColumnIdentity::qualified(&aliases.old, column),
259 width + position,
260 types[position].clone(),
261 ),
262 (
263 ColumnIdentity::qualified(&aliases.new, column),
264 width * 2 + position,
265 types[position].clone(),
266 ),
267 ]
268 })
269 .collect::<Vec<_>>();
270 RowSchema::with_physical_identity_aliases(&schema, &identity_aliases)
271}
272
273pub fn returning_expression_schema(
274 target: &RowSchema,
275 target_qualifier: &str,
276 aliases: &ReturningAliases,
277 supplemental: Option<&RowSchema>,
278) -> RowSchema {
279 let composite_width = target.len();
280 let (columns, types) = target_with_system_columns(target);
281 let target =
282 returning_context_schema(&columns, &types, composite_width, target_qualifier, aliases);
283 supplemental.map_or(target.clone(), |source| {
284 RowSchema::join(&target, source, std::iter::empty())
285 })
286}
287
288#[must_use]
290pub fn mutation_clause_schema(
291 target: &RowSchema,
292 target_qualifier: &str,
293 supplemental: Option<&RowSchema>,
294) -> RowSchema {
295 let composite_width = target.len();
296 let (columns, types) = target_with_system_columns(target);
297 let width = columns.len();
298 let target = RowSchema::with_wildcard_hidden_positions(
299 &RowSchema::with_qualified_types(target_qualifier, columns, types),
300 composite_width..width,
301 );
302 supplemental.map_or(target.clone(), |source| {
303 RowSchema::join(&target, source, std::iter::empty())
304 })
305}
306
307fn target_with_system_columns(
309 target: &RowSchema,
310) -> (Vec<String>, Vec<Option<crate::ast::ColumnType>>) {
311 let mut columns = target.columns().to_vec();
312 let mut types = target.column_types().to_vec();
313 if !columns.iter().any(|column| column == DOC_ID_COLUMN) {
314 columns.push(DOC_ID_COLUMN.into());
315 types.push(Some(crate::ast::ColumnType::BigInteger));
316 }
317 columns.push(TABLE_OID_COLUMN.into());
318 types.push(Some(crate::ast::ColumnType::Oid));
319 columns.push(crate::semantics::XMIN_COLUMN.into());
320 types.push(Some(crate::ast::ColumnType::Xid));
321 (columns, types)
322}
323
324pub fn query_plan_output_columns(plan: &crate::plan::QueryPlan) -> Option<Vec<String>> {
325 match &plan.root {
326 crate::plan::RelationalPlan::QueryBlock(block) => {
327 Some(projection_columns(&block.projections))
328 }
329 crate::plan::RelationalPlan::SetOp { left, .. } => query_plan_output_columns(left),
330 crate::plan::RelationalPlan::Values { rows, .. } => rows.first().map(|row| {
331 (1..=row.len())
332 .map(|index| format!("column{index}"))
333 .collect()
334 }),
335 }
336}
337
338pub fn should_defer_distinct_limit(stmt: &crate::plan::QueryBlockPlan) -> bool {
339 stmt.distinct && (stmt.limit.is_some() || stmt.offset.is_some())
340}
341
342pub fn select_execution_stmt(
343 stmt: &crate::plan::QueryBlockPlan,
344 defer_distinct_limit: bool,
345) -> crate::plan::QueryBlockPlan {
346 if !defer_distinct_limit {
347 return stmt.clone();
348 }
349 let mut exec_stmt = stmt.clone();
350 exec_stmt.limit = None;
351 exec_stmt.offset = None;
352 exec_stmt
353}