add pg_collation to pg_catalog

Giangblackk · Giangblackk · commit 06260dd302ec · 2025-09-13T16:29:19.000+07:00
diff --git a/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/mod.rs b/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/mod.rs
@@ -58,6 +58,7 @@ mod pg_type;
 mod pg_user;
 mod pg_user_mapping;
 mod pg_views;
+mod pg_collation;
 mod role_column_grants;
 mod role_table_grants;
 mod testing_blocking;
@@ -109,6 +110,7 @@ pub use pg_type::*;
 pub use pg_user::*;
 pub use pg_user_mapping::*;
 pub use pg_views::*;
+pub use pg_collation::*;
 pub use role_column_grants::*;
 pub use role_table_grants::*;
 pub use sql_implementation_info::*;
diff --git a/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_collation.rs b/rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_collation.rs
@@ -0,0 +1,249 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use datafusion::{
+    arrow::{
+        array::{Array, ArrayRef, BooleanBuilder, Int32Builder, StringBuilder, UInt32Builder},
+        datatypes::{DataType, Field, Schema},
+        record_batch::RecordBatch,
+    },
+    datasource::{datasource::TableProviderFilterPushDown, TableProvider},
+    error::Result,
+    logical_plan::Expr,
+    physical_plan::{memory::MemoryExec, ExecutionPlan},
+};
+
+use crate::compile::engine::information_schema::postgres::PG_NAMESPACE_CATALOG_OID;
+
+struct PgCollation {
+    oid: u32,
+    collname: &'static str,
+    collnamespace: u32,
+    collowner: u32,
+    collprovider: String,
+    collisdeterministic: bool,
+    collencoding: i32,
+    collcollate: Option<String>,
+    collctype: Option<String>,
+    // Column `colliculocale` renamed to `colllocale` since PostgreSQL 17.
+    // Support both columns for backward-compatibility.
+    // Reference: https://pgpedia.info/p/pg_collation.html
+    colliculocale: Option<String>,
+    colllocale: Option<String>,
+    collicurules: Option<String>,
+    collversion: Option<String>,
+}
+
+struct PgCatalogCollationBuilder {
+    oid: UInt32Builder,
+    collname: StringBuilder,
+    collnamespace: UInt32Builder,
+    collowner: UInt32Builder,
+    collprovider: StringBuilder,
+    collisdeterministic: BooleanBuilder,
+    collencoding: Int32Builder,
+    colliculocale: StringBuilder,
+    collcollate: StringBuilder,
+    collctype: StringBuilder,
+    colllocale: StringBuilder,
+    collicurules: StringBuilder,
+    collversion: StringBuilder,
+}
+
+impl PgCatalogCollationBuilder {
+    fn new() -> Self {
+        let capacity = 5;
+        Self {
+            oid: UInt32Builder::new(capacity),
+            collname: StringBuilder::new(capacity),
+            collnamespace: UInt32Builder::new(capacity),
+            collowner: UInt32Builder::new(capacity),
+            collprovider: StringBuilder::new(capacity),
+            collisdeterministic: BooleanBuilder::new(capacity),
+            collencoding: Int32Builder::new(capacity),
+            colliculocale: StringBuilder::new(capacity),
+            collcollate: StringBuilder::new(capacity),
+            collctype: StringBuilder::new(capacity),
+            colllocale: StringBuilder::new(capacity),
+            collicurules: StringBuilder::new(capacity),
+            collversion: StringBuilder::new(capacity),
+        }
+    }
+    fn add_collation(&mut self, coll: &PgCollation) {
+        self.oid.append_value(coll.oid).unwrap();
+        self.collname.append_value(coll.collname).unwrap();
+        self.collnamespace.append_value(coll.collnamespace).unwrap();
+        self.collowner.append_value(coll.collowner).unwrap();
+        self.collprovider
+            .append_value(coll.collprovider.clone())
+            .unwrap();
+        self.collisdeterministic
+            .append_value(coll.collisdeterministic)
+            .unwrap();
+        self.collencoding.append_value(coll.collencoding).unwrap();
+        self.colliculocale
+            .append_option(coll.colliculocale.clone())
+            .unwrap();
+        self.collcollate
+            .append_option(coll.collcollate.clone())
+            .unwrap();
+        self.collctype
+            .append_option(coll.collctype.clone())
+            .unwrap();
+        self.collicurules
+            .append_option(coll.collicurules.clone())
+            .unwrap();
+        self.colllocale
+            .append_option(coll.colllocale.clone())
+            .unwrap();
+        self.collversion
+            .append_option(coll.collversion.clone())
+            .unwrap();
+    }
+
+    fn finish(mut self) -> Vec<Arc<dyn Array>> {
+        let columns: Vec<Arc<dyn Array>> = vec![
+            Arc::new(self.oid.finish()),
+            Arc::new(self.collname.finish()),
+            Arc::new(self.collnamespace.finish()),
+            Arc::new(self.collowner.finish()),
+            Arc::new(self.collprovider.finish()),
+            Arc::new(self.collisdeterministic.finish()),
+            Arc::new(self.collencoding.finish()),
+            Arc::new(self.colliculocale.finish()),
+            Arc::new(self.collcollate.finish()),
+            Arc::new(self.collctype.finish()),
+            Arc::new(self.collicurules.finish()),
+            Arc::new(self.colllocale.finish()),
+            Arc::new(self.collversion.finish()),
+        ];
+        columns
+    }
+}
+
+pub struct PgCatalogCollationProvider {
+    data: Arc<Vec<ArrayRef>>,
+}
+
+impl PgCatalogCollationProvider {
+    pub fn new() -> Self {
+        // See https://github.com/postgres/postgres/blob/REL_16_4/src/include/catalog/pg_collation.h
+        let mut builder = PgCatalogCollationBuilder::new();
+
+        // Initial contents of the pg_collation system catalog.
+        // See https://github.com/postgres/postgres/blob/REL_16_4/src/include/catalog/pg_collation.dat
+        builder.add_collation(&PgCollation {
+            oid: 100,
+            collname: "default",
+            collnamespace: PG_NAMESPACE_CATALOG_OID,
+            collowner: 10,
+            collprovider: "d".to_string(),
+            collisdeterministic: true,
+            collencoding: -1,
+            collcollate: None,
+            colliculocale: None,
+            collctype: None,
+            colllocale: None,
+            collicurules: None,
+            collversion: None,
+        });
+        builder.add_collation(&PgCollation {
+            oid: 950,
+            collname: "C",
+            collnamespace: PG_NAMESPACE_CATALOG_OID,
+            collowner: 10,
+            collprovider: "c".to_string(),
+            collisdeterministic: true,
+            collencoding: -1,
+            collcollate: Some("C".to_string()),
+            colliculocale: Some("C".to_string()),
+            collctype: Some("C".to_string()),
+            colllocale: None,
+            collicurules: None,
+            collversion: None,
+        });
+        builder.add_collation(&PgCollation {
+            oid: 951,
+            collname: "POSIX",
+            collnamespace: PG_NAMESPACE_CATALOG_OID,
+            collowner: 10,
+            collprovider: "c".to_string(),
+            collisdeterministic: true,
+            collencoding: -1,
+            collcollate: Some("POSIX".to_string()),
+            colliculocale: Some("POSIX".to_string()),
+            collctype: Some("POSIX".to_string()),
+            colllocale: None,
+            collicurules: None,
+            collversion: None,
+        });
+        builder.add_collation(&PgCollation {
+            oid: 962,
+            collname: "usc_basic",
+            collnamespace: PG_NAMESPACE_CATALOG_OID,
+            collowner: 10,
+            collprovider: "c".to_string(),
+            collisdeterministic: true,
+            collencoding: 6,
+            collcollate: Some("C".to_string()),
+            colliculocale: Some("C".to_string()),
+            collctype: Some("C".to_string()),
+            colllocale: None,
+            collicurules: None,
+            collversion: None,
+        });
+        builder.add_collation(&PgCollation {
+            oid: 963,
+            collname: "unicode",
+            collnamespace: PG_NAMESPACE_CATALOG_OID,
+            collowner: 10,
+            collprovider: "i".to_string(),
+            collisdeterministic: true,
+            collencoding: -1,
+            collcollate: None,
+            colliculocale: None,
+            collctype: None,
+            colllocale: Some("und".to_string()),
+            collicurules: None,
+            collversion: Some("153.121".to_string()),
+        });
+        Self {
+            data: Arc::new(builder.finish()),
+        }
+    }
+}
+
+#[async_trait]
+impl TableProvider for PgCatalogCollationProvider {
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+    fn schema(&self) -> datafusion::arrow::datatypes::SchemaRef {
+        Arc::new(Schema::new(vec![
+            Field::new("oid", DataType::UInt32, false),
+            Field::new("collname", DataType::Utf8, false),
+            Field::new("collnamespace", DataType::UInt32, false),
+        ]))
+    }
+    async fn scan(
+        &self,
+        projection: &Option<Vec<usize>>,
+        _filters: &[Expr],
+        // limit can be used to reduce the amount scanned
+        // from the datasource as a performance optimization.
+        // If set, it contains the amount of rows needed by the `LogicalPlan`,
+        // The datasource should return *at least* this number of rows if available.
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let batch = RecordBatch::try_new(self.schema(), self.data.to_vec())?;
+        Ok(Arc::new(MemoryExec::try_new(
+            &[vec![batch]],
+            self.schema(),
+            projection.clone(),
+        )?))
+    }
+    fn supports_filter_pushdown(&self, _filter: &Expr) -> Result<TableProviderFilterPushDown> {
+        Ok(TableProviderFilterPushDown::Unsupported)
+    }
+}
