feat: handle metadata and bind packets

darthunix · darthunix · commit bcf1cac5eda0 · 2025-04-02T22:38:50.000+07:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -22,6 +22,7 @@ datafusion = "44.0"
 datafusion-sql = "44.0"
 fasthash = "0.4"
 libc = "0.2"
+once_cell = "1.21"
 pgrx = "0.12"
 rmp = "0.8"
 rust-fsm = { version = "0.7", features = ["diagram"] }
diff --git a/src/backend.rs b/src/backend.rs
@@ -113,18 +113,16 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
                     error!("Failed to send the table metadata: {}", err);
                 }
             }
-            Packet::Ack => break,
-            _ => error!("Unexpected packet in backend: {:?}", header.packet),
+            Packet::Bind => break,
+            Packet::Parse => error!("Unexpected packet in backend: {:?}", header.packet),
         }
     }
     let param_list = (*list_nth(list, 1)).ptr_value as ParamListInfo;
     let num_params = unsafe { (*param_list).numParams } as usize;
-    if num_params > 0 {
-        let params = unsafe { (*param_list).params.as_slice(num_params) };
-        let stream = wait_stream();
-        if let Err(err) = send_params(my_slot(), stream, params) {
-            error!("Failed to send the parameter list: {}", err);
-        }
+    let params = unsafe { (*param_list).params.as_slice(num_params) };
+    let stream = wait_stream();
+    if let Err(err) = send_params(my_slot(), stream, params) {
+        error!("Failed to send the parameter list: {}", err);
     }
     let css = CustomScanState {
         methods: exec_methods(),
diff --git a/src/data_type.rs b/src/data_type.rs
@@ -1,7 +1,7 @@
 use crate::error::FusionError;
 use crate::ipc::SlotStream;
 use anyhow::{bail, Result};
-use datafusion::arrow::datatypes::DataType;
+use datafusion::arrow::datatypes::{DataType, IntervalUnit, TimeUnit};
 use datafusion::common::arrow::array::types::IntervalMonthDayNano;
 use datafusion::common::arrow::datatypes::Field;
 use datafusion::common::ScalarValue;
@@ -107,6 +107,24 @@ pub(crate) enum EncodedType {
     Interval = 10,
 }
 
+impl EncodedType {
+    pub(crate) fn to_arrow(&self) -> DataType {
+        match self {
+            EncodedType::Boolean => DataType::Boolean,
+            EncodedType::Utf8 => DataType::Utf8,
+            EncodedType::Int16 => DataType::Int16,
+            EncodedType::Int32 => DataType::Int32,
+            EncodedType::Int64 => DataType::Int64,
+            EncodedType::Float32 => DataType::Float32,
+            EncodedType::Float64 => DataType::Float64,
+            EncodedType::Date32 => DataType::Date32,
+            EncodedType::Time64 => DataType::Time64(TimeUnit::Microsecond),
+            EncodedType::Timestamp => DataType::Timestamp(TimeUnit::Microsecond, None),
+            EncodedType::Interval => DataType::Interval(IntervalUnit::MonthDayNano),
+        }
+    }
+}
+
 impl TryFrom<u8> for EncodedType {
     type Error = FusionError;
 
diff --git a/src/protocol.rs b/src/protocol.rs
@@ -42,11 +42,10 @@ impl TryFrom<u8> for Direction {
 #[derive(Clone, Debug, Default, PartialEq)]
 pub enum Packet {
     #[default]
-    Ack = 0,
-    Bind = 1,
-    Failure = 2,
-    Metadata = 3,
-    Parse = 4,
+    Bind = 0,
+    Failure = 1,
+    Metadata = 2,
+    Parse = 3,
 }
 
 impl TryFrom<u8> for Packet {
@@ -55,11 +54,10 @@ impl TryFrom<u8> for Packet {
     fn try_from(value: u8) -> Result<Self, Self::Error> {
         assert!(value < 128);
         match value {
-            0 => Ok(Packet::Ack),
-            1 => Ok(Packet::Bind),
-            2 => Ok(Packet::Failure),
-            3 => Ok(Packet::Metadata),
-            4 => Ok(Packet::Parse),
+            0 => Ok(Packet::Bind),
+            1 => Ok(Packet::Failure),
+            2 => Ok(Packet::Metadata),
+            3 => Ok(Packet::Parse),
             _ => Err(FusionError::Deserialize("packet".to_string(), value.into())),
         }
     }
@@ -230,6 +228,21 @@ pub(crate) fn send_params(
     Ok(())
 }
 
+pub(crate) fn request_params(slot_id: SlotNumber, mut stream: SlotStream) -> Result<()> {
+    stream.reset();
+    let header = Header {
+        direction: Direction::ToBackend,
+        packet: Packet::Bind,
+        length: 0,
+        flag: Flag::Last,
+    };
+    write_header(&mut stream, &header)?;
+    // Unlock the slot after writing the parameters.
+    let _guard = Slot::from(stream);
+    signal(slot_id, Direction::ToBackend);
+    Ok(())
+}
+
 // FAILURE
 
 pub(crate) fn read_error(stream: &mut SlotStream) -> Result<String> {
@@ -334,21 +347,23 @@ pub(crate) fn send_table_refs(
 fn serialize_table(rel_oid: Oid, stream: &mut SlotStream) -> Result<()> {
     // The destructor will release the lock.
     let rel = unsafe { PgRelation::with_lock(rel_oid, pg_sys::AccessShareLock as i32) };
+    write_u32(stream, rel_oid.as_u32())?;
+    write_str(stream, rel.namespace())?;
+    write_str(stream, rel.name())?;
     let tuple_desc = rel.tuple_desc();
     let attr_num = u32::try_from(tuple_desc.iter().filter(|a| !a.is_dropped()).count())?;
-    write_u32(stream, rel_oid.as_u32())?;
     write_array_len(stream, attr_num)?;
     for attr in tuple_desc.iter() {
         if attr.is_dropped() {
             continue;
         }
+        write_array_len(stream, 3)?;
         let etype = EncodedType::try_from(attr.type_oid().value())?;
+        write_u8(stream, etype as u8)?;
         let is_nullable = !attr.attnotnull;
+        write_bool(stream, is_nullable)?;
         let name = attr.name();
-        write_array_len(stream, 3)?;
         write_str(stream, name)?;
-        write_u8(stream, etype as u8)?;
-        write_bool(stream, is_nullable)?;
     }
     Ok(())
 }
@@ -403,7 +418,7 @@ mod tests {
     fn test_header() {
         let header = Header {
             direction: Direction::ToWorker,
-            packet: Packet::Ack,
+            packet: Packet::Parse,
             length: 42,
             flag: Flag::Last,
         };
@@ -556,31 +571,39 @@ mod tests {
         // t1
         let oid = read_u32(&mut stream).unwrap();
         assert_eq!(oid, t1_oid.as_u32());
+        let ns_len = read_str_len(&mut stream).unwrap();
+        let ns = stream.look_ahead(ns_len as usize).unwrap();
+        assert_eq!(ns, b"public");
+        stream.rewind(ns_len as usize).unwrap();
+        let name_len = read_str_len(&mut stream).unwrap();
+        let name = stream.look_ahead(name_len as usize).unwrap();
+        assert_eq!(name, b"t1");
+        stream.rewind(name_len as usize).unwrap();
         let attr_num = read_array_len(&mut stream).unwrap();
         assert_eq!(attr_num, 2);
         // a
         let elem_num = read_array_len(&mut stream).unwrap();
         assert_eq!(elem_num, 3);
+        let etype = read_u8(&mut stream).unwrap();
+        assert_eq!(etype, EncodedType::Int32 as u8);
+        let is_nullable = read_bool(&mut stream).unwrap();
+        assert!(!is_nullable);
         let name_len = read_str_len(&mut stream).unwrap();
         assert_eq!(name_len, "a".len() as u32);
         let name = stream.look_ahead(name_len as usize).unwrap();
         assert_eq!(name, b"a");
         stream.rewind(name_len as usize).unwrap();
-        let etype = read_u8(&mut stream).unwrap();
-        assert_eq!(etype, EncodedType::Int32 as u8);
-        let is_nullable = read_bool(&mut stream).unwrap();
-        assert!(!is_nullable);
         // b
         let elem_num = read_array_len(&mut stream).unwrap();
         assert_eq!(elem_num, 3);
+        let etype = read_u8(&mut stream).unwrap();
+        assert_eq!(etype, EncodedType::Utf8 as u8);
+        let is_nullable = read_bool(&mut stream).unwrap();
+        assert!(is_nullable);
         let name_len = read_str_len(&mut stream).unwrap();
         assert_eq!(name_len, "b".len() as u32);
         let name = stream.look_ahead(name_len as usize).unwrap();
         assert_eq!(name, b"b");
         stream.rewind(name_len as usize).unwrap();
-        let etype = read_u8(&mut stream).unwrap();
-        assert_eq!(etype, EncodedType::Utf8 as u8);
-        let is_nullable = read_bool(&mut stream).unwrap();
-        assert!(is_nullable);
     }
 }
diff --git a/src/sql.rs b/src/sql.rs
@@ -1,12 +1,8 @@
-use std::cell::OnceCell;
-use std::collections::HashMap;
-use std::sync::Arc;
-
+use crate::data_type::EncodedType;
+use crate::ipc::SlotStream;
 use ahash::AHashMap;
-use datafusion::arrow::datatypes::DataType;
-use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::common::DFSchemaRef;
-use datafusion::common::ScalarValue;
+use anyhow::Result;
+use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion::config::ConfigOptions;
 use datafusion::error::DataFusionError;
 use datafusion::error::Result as DataFusionResult;
@@ -16,34 +12,22 @@ use datafusion::functions_aggregate::all_default_aggregate_functions;
 use datafusion::functions_window::all_default_window_functions;
 use datafusion::logical_expr::planner::ContextProvider;
 use datafusion::logical_expr::planner::ExprPlanner;
-use datafusion::logical_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF};
-use datafusion_sql::planner::SqlToRel;
-use datafusion_sql::sqlparser::dialect::PostgreSqlDialect;
-use datafusion_sql::sqlparser::parser::Parser;
+use datafusion::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
 use datafusion_sql::TableReference;
+use once_cell::sync::Lazy;
 use pgrx::pg_sys::Oid;
+use rmp::decode::read_array_len;
+use rmp::decode::read_bool;
+use rmp::decode::read_str_len;
+use rmp::decode::{read_u32, read_u8};
 use smol_str::SmolStr;
+use std::collections::HashMap;
+use std::str::from_utf8;
+use std::sync::Arc;
 
-// fn sql_to_logical_plan(
-//     sql: &str,
-//     params: Vec<ScalarValue>,
-// ) -> Result<LogicalPlan, DataFusionError> {
-//     let dialect = PostgreSqlDialect {};
-//     let ast = Parser::parse_sql(&dialect, sql).map_err(|e| DataFusionError::SQL(e, None))?;
-//     assert_eq!(ast.len(), 1);
-//     let statement = ast.into_iter().next().expect("ast is not empty");
-//
-//     // Cash metadata provider in a static variable to avoid re-allocation on each query.
-//     let base_plan = CATALOG.with(|catalog| {
-//         let catalog = catalog.get_or_init(Builtin::new);
-//         let sql_to_rel = SqlToRel::new(catalog);
-//         sql_to_rel.sql_statement_to_plan(statement)
-//     })?;
-//     let plan = base_plan.with_param_values(params)?;
-//
-//     Ok(plan)
-// }
+static BUILDIN: Lazy<Arc<Builtin>> = Lazy::new(|| Arc::new(Builtin::new()));
 
+#[derive(PartialEq, Eq, Hash)]
 pub(crate) struct Table {
     oid: Oid,
     schema: SchemaRef,
@@ -92,36 +76,86 @@ impl Builtin {
     }
 }
 
-struct Catalog {
+pub(crate) struct Catalog {
     builtin: Arc<Builtin>,
-    tables: AHashMap<SmolStr, Arc<dyn TableSource>>,
+    tables: AHashMap<TableReference, Arc<dyn TableSource>>,
+}
+
+impl Catalog {
+    pub(crate) fn from_stream(stream: &mut SlotStream) -> Result<Self> {
+        // The header should be consumed before calling this function.
+        let table_num = read_array_len(stream)?;
+        let mut tables = AHashMap::with_capacity(table_num as usize);
+        for _ in 0..table_num {
+            let oid = read_u32(stream)?;
+            let ns_len = read_str_len(stream)?;
+            let ns_bytes = stream.look_ahead(ns_len as usize)?;
+            let schema = SmolStr::from(from_utf8(ns_bytes)?);
+            stream.rewind(ns_len as usize)?;
+            let name_len = read_str_len(stream)?;
+            let name_bytes = stream.look_ahead(name_len as usize)?;
+            let name = from_utf8(name_bytes)?;
+            let table_ref = TableReference::partial(schema.as_str(), name);
+            stream.rewind(name_len as usize)?;
+            let column_num = read_array_len(stream)?;
+            let mut fields = Vec::with_capacity(column_num as usize);
+            for _ in 0..column_num {
+                let elem_num = read_array_len(stream)?;
+                assert_eq!(elem_num, 3);
+                let etype = read_u8(stream)?;
+                let df_type = EncodedType::try_from(etype)?.to_arrow();
+                let is_nullable = read_bool(stream)?;
+                let name_len = read_str_len(stream)?;
+                let name_bytes = stream.look_ahead(name_len as usize)?;
+                let name = from_utf8(name_bytes)?;
+                let field = Field::new(name, df_type, is_nullable);
+                fields.push(field);
+            }
+            let schema = Schema::new(fields);
+            let table = Table {
+                oid: Oid::from(oid),
+                schema: Arc::new(schema),
+            };
+            tables.insert(table_ref, Arc::new(table) as Arc<dyn TableSource>);
+        }
+        Ok(Self {
+            builtin: Arc::clone(&*BUILDIN),
+            tables,
+        })
+    }
 }
 
 impl ContextProvider for Catalog {
-    fn get_table_source(&self, name: TableReference) -> DataFusionResult<Arc<dyn TableSource>> {
-        match self.tables.get(name.table()) {
+    fn get_table_source(&self, table: TableReference) -> DataFusionResult<Arc<dyn TableSource>> {
+        match self.tables.get(&table) {
             Some(table) => Ok(Arc::clone(table)),
-            _ => Err(DataFusionError::Plan(format!(
-                "Table not found: {}",
-                name.table()
-            ))),
+            _ => {
+                let schema = match table.schema() {
+                    Some(schema) => &format!("{schema}."),
+                    _ => table.table(),
+                };
+                Err(DataFusionError::Plan(format!(
+                    "Table not found: {schema}{}",
+                    table.table()
+                )))
+            }
         }
     }
 
     fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>> {
-        self.builtin.scalar_udf.get(name).map(|f| Arc::clone(f))
+        self.builtin.scalar_udf.get(name).map(Arc::clone)
     }
 
     fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
-        self.builtin.agg_udf.get(name).map(|f| Arc::clone(f))
+        self.builtin.agg_udf.get(name).map(Arc::clone)
     }
 
     fn get_variable_type(&self, _variable_names: &[String]) -> Option<DataType> {
         None
     }
 
     fn get_window_meta(&self, name: &str) -> Option<Arc<WindowUDF>> {
-        self.builtin.window_udf.get(name).map(|f| Arc::clone(f))
+        self.builtin.window_udf.get(name).map(Arc::clone)
     }
 
     fn options(&self) -> &ConfigOptions {
diff --git a/src/worker.rs b/src/worker.rs

Original file line number	Diff line number	Diff line change
`@@ -113,18 +113,16 @@ unsafe extern "C" fn create_df_scan_state(cscan: mut CustomScan) -> mut Node {`
`113`	`113`	`error!("Failed to send the table metadata: {}", err);`
`114`	`114`	`}`
`115`	`115`	`}`
`116`		`- Packet::Ack => break,`
`117`		`- _ => error!("Unexpected packet in backend: {:?}", header.packet),`
	`116`	`+ Packet::Bind => break,`
	`117`	`+ Packet::Parse => error!("Unexpected packet in backend: {:?}", header.packet),`
`118`	`118`	`}`
`119`	`119`	`}`
`120`	`120`	`let param_list = (*list_nth(list, 1)).ptr_value as ParamListInfo;`
`121`	`121`	`let num_params = unsafe { (*param_list).numParams } as usize;`
`122`		`- if num_params > 0 {`
`123`		`- let params = unsafe { (*param_list).params.as_slice(num_params) };`
`124`		`- let stream = wait_stream();`
`125`		`- if let Err(err) = send_params(my_slot(), stream, params) {`
`126`		`- error!("Failed to send the parameter list: {}", err);`
`127`		`- }`
	`122`	`+ let params = unsafe { (*param_list).params.as_slice(num_params) };`
	`123`	`+ let stream = wait_stream();`
	`124`	`+ if let Err(err) = send_params(my_slot(), stream, params) {`
	`125`	`+ error!("Failed to send the parameter list: {}", err);`
`128`	`126`	`}`
`129`	`127`	`let css = CustomScanState {`
`130`	`128`	`methods: exec_methods(),`