feat: additional pg catalog fix from dbeaver queries

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::<Int32Type>();
+
+        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::<Int32Type>();
+
+        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::<Int32Type>();
+
+        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::<Int32Type>();
+
+        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<P>(
@@ -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(())
 }

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,20 +241,43 @@ impl ResolveUnqualifiedIdentifer {
         }
     }
 
-    fn rewrite_expr(expr: &mut Expr, wildcard_alias: &str, table_aliases: &HashSet<String>) {
+    fn get_projection_aliases(projection: &[SelectItem]) -> HashSet<String> {
+        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<String>,
+        projection_aliases: &HashSet<String>,
+    ) {
         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(),
                     ]);
                 }
             }
             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]

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]