wip: return columns from the worker and fix errors

Author: darthunix

src/backend.rs

@@ -13,8 +13,9 @@ use smallvec::SmallVec;
 use std::ffi::c_char;
 use std::time::Duration;
 
+use crate::data_type::unpack_target_entry;
 use crate::error::FusionError;
-use crate::ipc::{my_slot, Bus, SlotStream};
+use crate::ipc::{my_slot, worker_id, Bus, SlotStream, INVALID_PROC_NUMBER};
 use crate::protocol::{
     consume_header, read_error, request_explain, send_metadata, send_params, send_query, Direction,
     NeedSchema, Packet,
@@ -88,6 +89,9 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
     }
     let mut skip_wait = true;
     loop {
+        if worker_id() == INVALID_PROC_NUMBER {
+            error!("Worker ID is invalid");
+        }
         if !skip_wait {
             wait_latch(Some(BACKEND_WAIT_TIMEOUT));
             skip_wait = false;
@@ -96,7 +100,11 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
             continue;
         };
         let mut stream = SlotStream::from(slot);
-        let header = consume_header(&mut stream).expect("Failed to consume header");
+        let header = match consume_header(&mut stream) {
+            Ok(header) => header,
+            // TODO: before panic we should send a Failure message to the worker.
+            Err(err) => error!("Failed to consume header: {}", err),
+        };
         if header.direction != Direction::ToBackend {
             continue;
         }
@@ -105,37 +113,44 @@ unsafe extern "C" fn create_df_scan_state(cscan: *mut CustomScan) -> *mut Node {
                 // No data, just continue waiting.
                 continue;
             }
-            Packet::Failure => {
-                let msg = read_error(&mut stream).expect("Failed to read the error message");
-                error!("Failed to compile the query: {}", msg);
-            }
+            Packet::Failure => match read_error(&mut stream) {
+                Ok(msg) => error!("Failed to compile the query: {}", msg),
+                Err(err) => error!("Double error: {}", err),
+            },
             Packet::Metadata => {
-                let oids = table_oids(&mut stream).expect("Failed to read table OIDs");
+                let oids = match table_oids(&mut stream) {
+                    Ok(oids) => oids,
+                    Err(err) => error!("Failed to read the table OIDs: {}", err),
+                };
                 if let Err(err) = send_metadata(my_slot(), stream, &oids) {
                     error!("Failed to send the table metadata: {}", err);
                 }
+            }
+            Packet::Bind => {
+                let mut params: &[ParamExternData] = &[];
+                let param_list = (*list_nth(list, 1)).ptr_value as ParamListInfo;
+                if !param_list.is_null() {
+                    let num_params = unsafe { (*param_list).numParams } as usize;
+                    params = unsafe { (*param_list).params.as_slice(num_params) };
+                }
+                if let Err(err) = send_params(my_slot(), stream, params) {
+                    error!("Failed to send the parameter list: {}", err);
+                }
+            }
+            Packet::Columns => {
+                if let Err(err) = unpack_target_entry(&mut stream, (*cscan).custom_scan_tlist) {
+                    error!("Failed to unpack target entry: {}", err);
+                }
                 break;
             }
-            Packet::Bind | Packet::Parse | Packet::Explain => {
+            Packet::Parse | Packet::Explain => {
                 error!(
-                    "Unexpected packet for create custom plan: {:?}",
+                    "Unexpected packet while creating a custom plan: {:?}",
                     header.packet
                 )
             }
         }
     }
-    let mut params: &[ParamExternData] = &[];
-    let param_list = (*list_nth(list, 1)).ptr_value as ParamListInfo;
-    if !param_list.is_null() {
-        let num_params = unsafe { (*param_list).numParams } as usize;
-        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);
-    }
-    // TODO: request plan fields from the worker to build custom_scan_tlist
-    // with repack_output().
     let css = CustomScanState {
         methods: exec_methods(),
         ..Default::default()
@@ -200,7 +215,7 @@ unsafe extern "C" fn explain_df_scan(
                 let msg = read_error(&mut stream).expect("Failed to read the error message");
                 error!("Failed to compile the query: {}", msg);
             }
-            Packet::Metadata | Packet::Bind | Packet::Parse => {
+            Packet::Columns | Packet::Metadata | Packet::Bind | Packet::Parse => {
                 error!("Unexpected packet for explain: {:?}", header.packet)
             }
             Packet::Explain => {

src/data_type.rs

@@ -13,7 +13,8 @@ use pgrx::pg_sys::{
     TargetEntry, GETSTRUCT,
 };
 use rmp::decode::{
-    read_bool, read_f32, read_f64, read_i16, read_i32, read_i64, read_str_len, RmpRead,
+    read_array_len, read_bin_len, read_bool, read_f32, read_f64, read_i16, read_i32, read_i64,
+    read_str_len, read_u8, RmpRead,
 };
 use rmp::encode::{
     write_bool, write_f32, write_f64, write_i16, write_i32, write_i64, write_pfix, write_str,
@@ -173,6 +174,27 @@ impl TryFrom<pg_sys::Oid> for EncodedType {
     }
 }
 
+impl TryFrom<&DataType> for EncodedType {
+    type Error = FusionError;
+
+    fn try_from(value: &DataType) -> Result<Self, Self::Error> {
+        match value {
+            DataType::Boolean => Ok(EncodedType::Boolean),
+            DataType::Utf8 => Ok(EncodedType::Utf8),
+            DataType::Int16 => Ok(EncodedType::Int16),
+            DataType::Int32 => Ok(EncodedType::Int32),
+            DataType::Int64 => Ok(EncodedType::Int64),
+            DataType::Float32 => Ok(EncodedType::Float32),
+            DataType::Float64 => Ok(EncodedType::Float64),
+            DataType::Date32 => Ok(EncodedType::Date32),
+            DataType::Time64(_) => Ok(EncodedType::Time64),
+            DataType::Timestamp(_, _) => Ok(EncodedType::Timestamp),
+            DataType::Interval(_) => Ok(EncodedType::Interval),
+            _ => Err(FusionError::UnsupportedType(format!("{:?}", value))),
+        }
+    }
+}
+
 #[inline]
 pub(crate) fn write_scalar_value(stream: &mut SlotStream, value: &ScalarValue) -> Result<()> {
     let write_null = |stream: &mut SlotStream| -> Result<()> {
@@ -416,6 +438,41 @@ pub(crate) fn repack_output(columns: &[Field]) -> *mut List {
     list
 }
 
+// The header of the stream must already be consumed.
+pub(crate) fn unpack_target_entry(stream: &mut SlotStream, list: *mut List) -> Result<()> {
+    let column_len = read_array_len(stream)?;
+    assert!(column_len < i16::MAX as u32);
+    for position in 0..column_len {
+        let pos = position as i16 + 1;
+        let etype = read_u8(stream)?;
+        let oid = type_to_oid(&EncodedType::try_from(etype)?.to_arrow());
+        let tuple =
+            unsafe { pg_sys::SearchSysCache1(TYPEOID as i32, pg_sys::ObjectIdGetDatum(oid)) };
+        if tuple.is_null() {
+            bail!(FusionError::UnsupportedType(format!("{:?}", oid)));
+        }
+        let name_len = read_bin_len(stream)? as usize;
+        let name = stream.look_ahead(name_len)?;
+        unsafe {
+            let typtup = GETSTRUCT(tuple) as pg_sys::Form_pg_type;
+            let expr = makeVar(
+                pg_sys::INDEX_VAR,
+                pos,
+                oid,
+                (*typtup).typtypmod,
+                (*typtup).typcollation,
+                0,
+            );
+            let col_name = palloc0(name_len) as *mut u8;
+            std::ptr::copy_nonoverlapping(name.as_ptr(), col_name, name_len);
+            let entry = makeTargetEntry(expr as *mut Expr, pos, col_name as *mut i8, false);
+            pg_sys::ReleaseSysCache(tuple);
+            list_append_unique_ptr(list, entry as *mut c_void);
+        }
+    }
+    Ok(())
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {

src/protocol.rs

@@ -7,7 +7,7 @@ use crate::ipc::{worker_id, Bus, Slot, SlotNumber, SlotStream, DATA_SIZE};
 use crate::sql::Table;
 use ahash::AHashMap;
 use anyhow::Result;
-use datafusion::arrow::datatypes::{Field, Schema};
+use datafusion::arrow::datatypes::{Field, Fields, Schema};
 use datafusion::logical_expr::TableSource;
 use datafusion::scalar::ScalarValue;
 use datafusion_sql::TableReference;
@@ -57,6 +57,7 @@ pub enum Packet {
     Metadata = 3,
     Parse = 4,
     Explain = 5,
+    Columns = 6,
 }
 
 impl TryFrom<u8> for Packet {
@@ -71,6 +72,7 @@ impl TryFrom<u8> for Packet {
             3 => Ok(Packet::Metadata),
             4 => Ok(Packet::Parse),
             5 => Ok(Packet::Explain),
+            6 => Ok(Packet::Columns),
             _ => Err(FusionError::Deserialize("packet".to_string(), value.into())),
         }
     }
@@ -512,6 +514,34 @@ pub(crate) fn request_explain(slot_id: SlotNumber, mut stream: SlotStream) -> Re
     Ok(())
 }
 
+// COLUMNS
+
+pub(crate) fn prepare_columns(stream: &mut SlotStream, columns: &Fields) -> Result<()> {
+    stream.reset();
+    write_header(stream, &Header::default())?;
+    let pos_init = stream.position();
+    write_array_len(stream, u32::try_from(columns.len())?)?;
+    for column in columns {
+        write_u8(stream, EncodedType::try_from(column.data_type())? as u8)?;
+        let len = u32::try_from(column.name().len() + 1)?;
+        write_bin_len(stream, len)?;
+        stream.write_bytes(column.name().as_bytes())?;
+        write_pfix(stream, 0)?;
+    }
+    let pos_final = stream.position();
+    let length = u16::try_from(pos_final - pos_init)?;
+    let header = Header {
+        direction: Direction::ToBackend,
+        packet: Packet::Columns,
+        length,
+        flag: Flag::Last,
+    };
+    stream.reset();
+    write_header(stream, &header)?;
+    stream.rewind(length as usize)?;
+    Ok(())
+}
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pg_schema]
 mod tests {