diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attrdef.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attrdef.feather index 4470857..0c5d9e2 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attrdef.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attrdef.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attribute.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attribute.feather index ee0af97..118e122 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attribute.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_attribute.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_auth_members.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_auth_members.feather index 5e5dde1..c968261 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_auth_members.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_auth_members.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_authid.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_authid.feather index 87d11d1..460f285 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_authid.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_authid.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_class.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_class.feather index ee79ccf..41faca4 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_class.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_class.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_constraint.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_constraint.feather index a618337..5572bd9 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_constraint.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_constraint.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_database.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_database.feather index a43468f..5d4c2da 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_database.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_database.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_db_role_setting.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_db_role_setting.feather index 0d2577d..7e2460c 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_db_role_setting.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_db_role_setting.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_default_acl.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_default_acl.feather index 097878e..adaf398 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_default_acl.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_default_acl.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_depend.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_depend.feather index 88178cc..72f9f1a 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_depend.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_depend.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_description.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_description.feather index b0d71b9..6e0c91e 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_description.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_description.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_enum.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_enum.feather index dbad2b7..0844d53 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_enum.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_enum.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_event_trigger.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_event_trigger.feather index e63feda..7ac630b 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_event_trigger.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_event_trigger.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_extension.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_extension.feather index c3a438b..0801e5d 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_extension.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_extension.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_data_wrapper.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_data_wrapper.feather index 61899a5..22d7d31 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_data_wrapper.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_data_wrapper.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_server.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_server.feather index 8d22dc6..c04a3c3 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_server.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_server.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_table.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_table.feather index 47079d4..30e9080 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_table.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_foreign_table.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_index.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_index.feather index 1e9b990..a4f6585 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_index.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_index.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_inherits.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_inherits.feather index ea4e2ea..d6649cb 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_inherits.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_inherits.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_init_privs.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_init_privs.feather index 892210a..f17a229 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_init_privs.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_init_privs.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject.feather index 296772f..1c90d18 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject_metadata.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject_metadata.feather index 193c4c8..fad0b2d 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject_metadata.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_largeobject_metadata.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_namespace.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_namespace.feather index ebe0609..302d58c 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_namespace.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_namespace.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_parameter_acl.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_parameter_acl.feather index 365964c..49dff9f 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_parameter_acl.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_parameter_acl.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_partitioned_table.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_partitioned_table.feather index da6e42d..71e2338 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_partitioned_table.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_partitioned_table.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_policy.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_policy.feather index ac03a61..1084a11 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_policy.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_policy.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication.feather index f659f1c..20f57b8 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_namespace.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_namespace.feather index 8bb2eeb..7afbe07 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_namespace.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_namespace.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_rel.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_rel.feather index ea8ee46..cb0e864 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_rel.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_publication_rel.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_replication_origin.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_replication_origin.feather index 7b0298a..0c76fc7 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_replication_origin.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_replication_origin.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_rewrite.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_rewrite.feather index dc08887..64a3889 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_rewrite.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_rewrite.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_seclabel.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_seclabel.feather index a327f32..ed03fb8 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_seclabel.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_seclabel.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_sequence.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_sequence.feather index ca46efd..c0e1b33 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_sequence.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_sequence.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdepend.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdepend.feather index cbdb81e..526e971 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdepend.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdepend.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdescription.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdescription.feather index d7f039c..591a167 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdescription.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shdescription.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shseclabel.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shseclabel.feather index 2907e5d..4bbb85b 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shseclabel.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_shseclabel.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic.feather index 3c1ea37..2a3aea9 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext.feather index b1a5a70..8109dfb 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext_data.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext_data.feather index 3eaee81..386c8f6 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext_data.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_statistic_ext_data.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription.feather index 2bf7478..0a5af23 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription_rel.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription_rel.feather index 4b7c393..d5197f8 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription_rel.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_subscription_rel.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_tablespace.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_tablespace.feather index f2035f8..1c59e2f 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_tablespace.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_tablespace.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_transform.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_transform.feather index 8cf564a..0a0ec3b 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_transform.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_transform.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_trigger.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_trigger.feather index b8524c9..339fec8 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_trigger.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_trigger.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_ts_config_map.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_ts_config_map.feather index d4688f2..29d904c 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_ts_config_map.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_ts_config_map.feather differ diff --git a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_user_mapping.feather b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_user_mapping.feather index cde913b..1609631 100644 Binary files a/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_user_mapping.feather and b/datafusion-pg-catalog/pg_catalog_arrow_exports/pg_user_mapping.feather differ diff --git a/datafusion-pg-catalog/src/pg_catalog.rs b/datafusion-pg-catalog/src/pg_catalog.rs index 1cc731b..ecc3bad 100644 --- a/datafusion-pg-catalog/src/pg_catalog.rs +++ b/datafusion-pg-catalog/src/pg_catalog.rs @@ -1224,7 +1224,7 @@ pub fn create_pg_get_statisticsobjdef_columns_udf() -> ScalarUDF { create_udf( "pg_get_statisticsobjdef_columns", - vec![DataType::UInt32], + vec![DataType::Int32], DataType::Utf8, Volatility::Stable, Arc::new(func), @@ -1259,6 +1259,8 @@ pub fn create_pg_encoding_to_char_udf() -> ScalarUDF { ) } +const BACKEND_PID: i32 = 1; + pub fn create_pg_backend_pid_udf() -> ScalarUDF { let func = move |_args: &[ColumnarValue]| { let mut builder = Int32Builder::new(); @@ -1276,7 +1278,97 @@ pub fn create_pg_backend_pid_udf() -> ScalarUDF { ) } -const BACKEND_PID: i32 = 1; +pub fn create_pg_relation_size_udf() -> ScalarUDF { + let func = move |args: &[ColumnarValue]| { + let args = ColumnarValue::values_to_arrays(args)?; + let oids = &args[0].as_primitive::(); + + let mut builder = Int32Builder::new(); + for _ in 0..oids.len() { + builder.append_value(0); + } + + let array: ArrayRef = Arc::new(builder.finish()); + Ok(ColumnarValue::Array(array)) + }; + + create_udf( + "pg_relation_size", + vec![DataType::Int32], + DataType::Int32, + Volatility::Stable, + Arc::new(func), + ) +} + +pub fn create_pg_total_relation_size_udf() -> ScalarUDF { + let func = move |args: &[ColumnarValue]| { + let args = ColumnarValue::values_to_arrays(args)?; + let oids = &args[0].as_primitive::(); + + let mut builder = Int32Builder::new(); + for _ in 0..oids.len() { + builder.append_value(0); + } + + let array: ArrayRef = Arc::new(builder.finish()); + Ok(ColumnarValue::Array(array)) + }; + + create_udf( + "pg_total_relation_size", + vec![DataType::Int32], + DataType::Int32, + Volatility::Stable, + Arc::new(func), + ) +} + +pub fn create_pg_stat_get_numscans() -> ScalarUDF { + let func = move |args: &[ColumnarValue]| { + let args = ColumnarValue::values_to_arrays(args)?; + let index_rel_id = &args[0].as_primitive::(); + + let mut builder = Int32Builder::new(); + for _ in 0..index_rel_id.len() { + builder.append_value(0); + } + + let array: ArrayRef = Arc::new(builder.finish()); + Ok(ColumnarValue::Array(array)) + }; + + create_udf( + "pg_stat_get_numscans", + vec![DataType::Int32], + DataType::Int32, + Volatility::Stable, + Arc::new(func), + ) +} + +pub fn create_pg_get_constraintdef() -> ScalarUDF { + let func = move |args: &[ColumnarValue]| { + let args = ColumnarValue::values_to_arrays(args)?; + let oids = &args[0].as_primitive::(); + + let mut builder = StringBuilder::new(); + for _ in 0..oids.len() { + builder.append_value(""); + } + + let array: ArrayRef = Arc::new(builder.finish()); + Ok(ColumnarValue::Array(array)) + }; + + create_udf( + "pg_get_constraintdef", + vec![DataType::Int32], + DataType::Utf8, + Volatility::Stable, + Arc::new(func), + ) +} /// Install pg_catalog and postgres UDFs to current `SessionContext` pub fn setup_pg_catalog

( @@ -1329,6 +1421,10 @@ where session_context.register_udf(create_pg_get_statisticsobjdef_columns_udf()); session_context.register_udf(create_pg_encoding_to_char_udf()); session_context.register_udf(create_pg_backend_pid_udf()); + session_context.register_udf(create_pg_relation_size_udf()); + session_context.register_udf(create_pg_total_relation_size_udf()); + session_context.register_udf(create_pg_stat_get_numscans()); + session_context.register_udf(create_pg_get_constraintdef()); Ok(()) } diff --git a/datafusion-pg-catalog/src/sql/rules.rs b/datafusion-pg-catalog/src/sql/rules.rs index b6ddc65..85e5bce 100644 --- a/datafusion-pg-catalog/src/sql/rules.rs +++ b/datafusion-pg-catalog/src/sql/rules.rs @@ -162,16 +162,29 @@ impl ResolveUnqualifiedIdentifer { let wildcard_alias = qualified_wildcard_alias.unwrap(); + // Step 2.5: Collect all projection aliases to avoid rewriting them + let projection_aliases = Self::get_projection_aliases(&select.projection); + // Step 3: Rewrite expressions in the WHERE and ORDER BY clauses. if let Some(selection) = &mut select.selection { - Self::rewrite_expr(selection, &wildcard_alias, &table_aliases); + Self::rewrite_expr( + selection, + &wildcard_alias, + &table_aliases, + &projection_aliases, + ); } if let Some(OrderByKind::Expressions(order_by_exprs)) = query.order_by.as_mut().map(|o| &mut o.kind) { for order_by_expr in order_by_exprs { - Self::rewrite_expr(&mut order_by_expr.expr, &wildcard_alias, &table_aliases); + Self::rewrite_expr( + &mut order_by_expr.expr, + &wildcard_alias, + &table_aliases, + &projection_aliases, + ); } } } @@ -228,11 +241,34 @@ impl ResolveUnqualifiedIdentifer { } } - fn rewrite_expr(expr: &mut Expr, wildcard_alias: &str, table_aliases: &HashSet) { + fn get_projection_aliases(projection: &[SelectItem]) -> HashSet { + let mut aliases = HashSet::new(); + for item in projection { + match item { + SelectItem::ExprWithAlias { alias, .. } => { + aliases.insert(alias.value.clone()); + } + SelectItem::UnnamedExpr(Expr::Identifier(ident)) => { + aliases.insert(ident.value.clone()); + } + _ => {} + } + } + aliases + } + + fn rewrite_expr( + expr: &mut Expr, + wildcard_alias: &str, + table_aliases: &HashSet, + projection_aliases: &HashSet, + ) { match expr { Expr::Identifier(ident) => { - // If the identifier is not a table alias itself, rewrite it. - if !table_aliases.contains(&ident.value) { + // If the identifier is not a table alias itself and not already aliased in projection, rewrite it. + if !table_aliases.contains(&ident.value) + && !projection_aliases.contains(&ident.value) + { *expr = Expr::CompoundIdentifier(vec![ Ident::new(wildcard_alias.to_string()), ident.clone(), @@ -240,8 +276,8 @@ impl ResolveUnqualifiedIdentifer { } } Expr::BinaryOp { left, right, .. } => { - Self::rewrite_expr(left, wildcard_alias, table_aliases); - Self::rewrite_expr(right, wildcard_alias, table_aliases); + Self::rewrite_expr(left, wildcard_alias, table_aliases, projection_aliases); + Self::rewrite_expr(right, wildcard_alias, table_aliases, projection_aliases); } // Add more cases for other expression types as needed (e.g., `InList`, `Between`, etc.) _ => {} @@ -907,6 +943,11 @@ mod tests { "SELECT n.oid,n.*,d.description FROM pg_catalog.pg_namespace n LEFT OUTER JOIN pg_catalog.pg_description d ON d.objoid=n.oid AND d.objsubid=0 AND d.classoid='pg_namespace' ORDER BY nspsname", "SELECT n.oid, n.*, d.description FROM pg_catalog.pg_namespace AS n LEFT OUTER JOIN pg_catalog.pg_description AS d ON d.objoid = n.oid AND d.objsubid = 0 AND d.classoid = 'pg_namespace' ORDER BY n.nspsname" ); + + assert_rewrite!(&rules, + "SELECT i.*,i.indkey as keys,c.relname,c.relnamespace,c.relam,c.reltablespace,tc.relname as tabrelname,dsc.description FROM pg_catalog.pg_index i INNER JOIN pg_catalog.pg_class c ON c.oid=i.indexrelid INNER JOIN pg_catalog.pg_class tc ON tc.oid=i.indrelid LEFT OUTER JOIN pg_catalog.pg_description dsc ON i.indexrelid=dsc.objoid WHERE i.indrelid=1 ORDER BY tabrelname, c.relname", + "SELECT i.*, i.indkey AS keys, c.relname, c.relnamespace, c.relam, c.reltablespace, tc.relname AS tabrelname, dsc.description FROM pg_catalog.pg_index AS i INNER JOIN pg_catalog.pg_class AS c ON c.oid = i.indexrelid INNER JOIN pg_catalog.pg_class AS tc ON tc.oid = i.indrelid LEFT OUTER JOIN pg_catalog.pg_description AS dsc ON i.indexrelid = dsc.objoid WHERE i.indrelid = 1 ORDER BY tabrelname, c.relname" + ); } #[test] diff --git a/datafusion-postgres/tests/dbeaver.rs b/datafusion-postgres/tests/dbeaver.rs index 6602b6d..c69c96f 100644 --- a/datafusion-postgres/tests/dbeaver.rs +++ b/datafusion-postgres/tests/dbeaver.rs @@ -21,7 +21,25 @@ const DBEAVER_QUERIES: &[&str] = &[ "SELECT c.oid,c.*,d.description,pg_catalog.pg_get_expr(c.relpartbound, c.oid) as partition_expr, pg_catalog.pg_get_partkeydef(c.oid) as partition_key FROM pg_catalog.pg_class c LEFT OUTER JOIN pg_catalog.pg_description d ON d.objoid=c.oid AND d.objsubid=0 AND d.classoid='pg_class'::regclass - WHERE c.relnamespace=11 AND c.relkind not in ('i','I','c')" + WHERE c.relnamespace=11 AND c.relkind not in ('i','I','c')", + "select c.oid,pg_catalog.pg_total_relation_size(c.oid) as total_rel_size,pg_catalog.pg_relation_size(c.oid) as rel_size + FROM pg_class c + WHERE c.relnamespace='public'", + + "SELECT i.*,i.indkey as keys,c.relname,c.relnamespace,c.relam,c.reltablespace,tc.relname as tabrelname,dsc.description,pg_catalog.pg_get_expr(i.indpred, i.indrelid) as pred_expr,pg_catalog.pg_get_expr(i.indexprs, i.indrelid, true) as expr,pg_catalog.pg_relation_size(i.indexrelid) as index_rel_size,pg_catalog.pg_stat_get_numscans(i.indexrelid) as index_num_scans FROM pg_catalog.pg_index i + INNER JOIN pg_catalog.pg_class c ON c.oid=i.indexrelid + INNER JOIN pg_catalog.pg_class tc ON tc.oid=i.indrelid + LEFT OUTER JOIN pg_catalog.pg_description dsc ON i.indexrelid=dsc.objoid + WHERE i.indrelid=1 ORDER BY tabrelname, c.relname", + + "SELECT c.oid,c.*,t.relname as tabrelname,rt.relnamespace as refnamespace,d.description, case when c.contype='c' then \"substring\"(pg_get_constraintdef(c.oid), 7) else null end consrc_copy + FROM pg_catalog.pg_constraint c + INNER JOIN pg_catalog.pg_class t ON t.oid=c.conrelid + LEFT OUTER JOIN pg_catalog.pg_class rt ON rt.oid=c.confrelid + LEFT OUTER JOIN pg_catalog.pg_description d ON d.objoid=c.oid AND d.objsubid=0 AND d.classoid='pg_constraint'::regclass + WHERE c.conrelid=1 + ORDER BY c.oid", + ]; #[tokio::test]