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 { name, .. } => crate::parse_regobject_name(name)
28 .and_then(|mut names| names.pop())
29 .unwrap_or_else(|| name.clone()),
30 _ => "?column?".into(),
31 }
32}
33
34pub fn expand_from_star_columns(
35 columns: Vec<String>,
36 projections: &[ProjectionPlan],
37 source_schema: &RowSchema,
38) -> Result<Vec<String>, SQLError> {
39 let mut output = Vec::new();
40 for (position, projection) in projections.iter().enumerate() {
41 match &projection.expr {
42 ScalarExpr::Star => {
43 output.extend(
44 source_schema
45 .columns()
46 .iter()
47 .enumerate()
48 .filter(|(position, _)| {
49 visible_projection_source_position(source_schema, *position)
50 })
51 .map(|(source_position, column)| {
52 source_schema
53 .public_name(source_position)
54 .unwrap_or(column)
55 .to_string()
56 }),
57 );
58 }
59 ScalarExpr::QualifiedStar(qualifier) => {
60 let qualified_columns = source_schema
61 .qualified_star_position_layout(qualifier)
62 .into_iter()
63 .filter(|(_, logical, _, _)| {
64 logical.is_none_or(|position| {
65 visible_projection_source_position(source_schema, position)
66 })
67 })
68 .map(|(column, _, _, _)| column)
69 .collect::<Vec<_>>();
70 if qualified_columns.is_empty() {
71 return Err(SQLError::UnknownTable(qualifier.clone()));
72 }
73 output.extend(qualified_columns);
74 }
75 _ => output.push(columns[position].clone()),
76 }
77 }
78 Ok(output)
79}
80
81pub fn bound_projection_expression(schema: &RowSchema, position: usize) -> ScalarExpr {
82 let Some(identity) = schema.identity(position) else {
83 return ScalarExpr::Position(position);
84 };
85 if let Some(qualifier) = identity.qualifier() {
86 if schema.qualified_position(qualifier, identity.column()) == Some(position) {
87 return ScalarExpr::qualified_column(qualifier, identity.column());
88 }
89 } else if schema.unqualified_position(identity.column()) == Some(position) {
90 return ScalarExpr::Column(identity.column().to_string());
91 }
92 ScalarExpr::Position(position)
93}
94
95pub fn expand_bound_projection_stars(
96 projections: &[ProjectionPlan],
97 schema: &RowSchema,
98) -> Result<Vec<ProjectionPlan>, SQLError> {
99 let mut expanded = Vec::new();
100 for projection in projections {
101 match &projection.expr {
102 ScalarExpr::Star => {
103 for (position, column) in schema.columns().iter().enumerate() {
104 if !visible_projection_source_position(schema, position) {
105 continue;
106 }
107 expanded.push(ProjectionPlan {
108 expr: bound_projection_expression(schema, position),
109 alias: Some(schema.public_name(position).unwrap_or(column).to_string()),
110 });
111 }
112 }
113 ScalarExpr::QualifiedStar(qualifier) => {
114 let layout = schema.qualified_star_position_layout(qualifier);
115 if layout.is_empty() {
116 return Err(SQLError::UnknownTable(qualifier.clone()));
117 }
118 for (column, logical, _, _) in layout {
119 if logical.is_some_and(|position| {
120 !visible_projection_source_position(schema, position)
121 }) {
122 continue;
123 }
124 expanded.push(ProjectionPlan {
125 expr: logical.map_or_else(
126 || ScalarExpr::qualified_column(qualifier, &column),
127 |position| bound_projection_expression(schema, position),
128 ),
129 alias: Some(column),
130 });
131 }
132 }
133 _ => expanded.push(projection.clone()),
134 }
135 }
136 Ok(expanded)
137}
138
139pub fn visible_projection_source_position(schema: &RowSchema, position: usize) -> bool {
140 schema.wildcard_position_visible(position)
141}
142
143pub fn returning_context_schema(
144 columns: &[String],
145 types: &[Option<crate::ast::ColumnType>],
146 composite_width: usize,
147 target_qualifier: &str,
148 aliases: &ReturningAliases,
149) -> RowSchema {
150 let target =
151 RowSchema::with_qualified_types(target_qualifier, columns.to_vec(), types.to_vec());
152 let target = RowSchema::with_wildcard_hidden_positions(&target, composite_width..columns.len());
153 let hidden_types = types
154 .iter()
155 .cloned()
156 .chain(types.iter().cloned())
157 .collect::<Vec<_>>();
158 let schema = RowSchema::append_hidden_typed(&target, &hidden_types);
159 let width = columns.len();
160 let identity_aliases = columns
161 .iter()
162 .enumerate()
163 .flat_map(|(position, column)| {
164 [
165 (
166 ColumnIdentity::qualified(&aliases.old, column),
167 width + position,
168 types[position].clone(),
169 ),
170 (
171 ColumnIdentity::qualified(&aliases.new, column),
172 width * 2 + position,
173 types[position].clone(),
174 ),
175 ]
176 })
177 .collect::<Vec<_>>();
178 RowSchema::with_physical_identity_aliases(&schema, &identity_aliases)
179}
180
181pub fn returning_expression_schema(
182 target: &RowSchema,
183 target_qualifier: &str,
184 aliases: &ReturningAliases,
185 supplemental: Option<&RowSchema>,
186) -> RowSchema {
187 let composite_width = target.len();
188 let mut columns = target.columns().to_vec();
189 let mut types = target.column_types().to_vec();
190 if !columns.iter().any(|column| column == DOC_ID_COLUMN) {
191 columns.push(DOC_ID_COLUMN.into());
192 types.push(Some(crate::ast::ColumnType::BigInteger));
193 }
194 columns.push(TABLE_OID_COLUMN.into());
195 types.push(Some(crate::ast::ColumnType::Oid));
196 columns.push(crate::semantics::XMIN_COLUMN.into());
197 types.push(Some(crate::ast::ColumnType::Xid));
198 let target =
199 returning_context_schema(&columns, &types, composite_width, target_qualifier, aliases);
200 supplemental.map_or(target.clone(), |source| {
201 RowSchema::join(&target, source, std::iter::empty())
202 })
203}
204
205pub fn query_plan_output_columns(plan: &crate::plan::QueryPlan) -> Option<Vec<String>> {
206 match &plan.root {
207 crate::plan::RelationalPlan::QueryBlock(block) => {
208 Some(projection_columns(&block.projections))
209 }
210 crate::plan::RelationalPlan::SetOp { left, .. } => query_plan_output_columns(left),
211 crate::plan::RelationalPlan::Values { rows, .. } => rows.first().map(|row| {
212 (1..=row.len())
213 .map(|index| format!("column{index}"))
214 .collect()
215 }),
216 }
217}
218
219pub fn should_defer_distinct_limit(stmt: &crate::plan::QueryBlockPlan) -> bool {
220 stmt.distinct && (stmt.limit.is_some() || stmt.offset.is_some())
221}
222
223pub fn select_execution_stmt(
224 stmt: &crate::plan::QueryBlockPlan,
225 defer_distinct_limit: bool,
226) -> crate::plan::QueryBlockPlan {
227 if !defer_distinct_limit {
228 return stmt.clone();
229 }
230 let mut exec_stmt = stmt.clone();
231 exec_stmt.limit = None;
232 exec_stmt.offset = None;
233 exec_stmt
234}