bill/unchecked-io
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

initial commit

Browse Source
This commit is contained in:
Bill
2025-11-17 14:29:06 -07:00
parent ef8f0d8b96
commit 6e5c0acb58
4 changed files with 487 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
config.yaml
View File

@@ -0,0 +1,28 @@
connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres
# The query to execute. This MUST be a COPY...TO STDOUT (FORMAT binary) command.
# We cast all columns to match the binary types our parser expects.
query: "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN, last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8 FROM benchmark_table) TO STDOUT (FORMAT binary)"
# 'schema' is the ordered list of columns *exactly* as they appear in the query's SELECT statement.
schema:
- column_name: id
arrow_type: Int64
- column_name: uuid
arrow_type: Utf8
- column_name: username
arrow_type: Utf8
- column_name: score
arrow_type: Float32
- column_name: is_active
arrow_type: Boolean
- column_name: last_login
arrow_type: Timestamp(Nanosecond, None)
- column_name: notes
arrow_type: Utf8
- column_name: course_id
arrow_type: Int32
- column_name: start_date
arrow_type: Date32
- column_name: rating
arrow_type: Float64

127
setup_db.sql
View File

@@ -0,0 +1,127 @@
-- 1. Drop the old table if it exists
DROP TABLE IF EXISTS benchmark_table;
-- 2. Create the new 10-column table
CREATE TABLE benchmark_table (
id BIGINT,
uuid TEXT, -- Storing as TEXT for simplicity, can be native UUID
username TEXT,
score REAL,
is_active BOOLEAN,
last_login TIMESTAMP,
notes TEXT,
course_id INT,
start_date DATE,
rating FLOAT8
);
-- 3. Use COPY FROM STDIN to bulk-load 100 rows of data
-- This is the fastest way to insert data into Postgres.
COPY benchmark_table (id, uuid, username, score, is_active, last_login, notes, course_id, start_date, rating) FROM STDIN WITH (FORMAT csv);
1,a1b2c3d4-e5f6-7890-1234-567890abcdef,user_1,95.5,true,2025-01-10 10:30:00,First note,101,2024-01-01,4.5
2,b2c3d4e5-f6a7-8901-2345-678901bcdef0,user_2,88.0,true,2025-02-15 11:45:10,Second note,102,2024-01-02,4.2
3,c3d4e5f6-a7b8-9012-3456-789012cdef01,user_3,72.1,false,2025-03-20 12:00:05,,101,2024-01-03,3.8
4,d4e5f6a7-b8c9-0123-4567-890123def012,user_4,99.9,true,2025-04-25 14:15:30,High score,103,2024-01-04,4.9
5,e5f6a7b8-c9d0-1234-5678-901234ef0123,user_5,65.0,false,2025-05-30 09:10:20,Needs review,102,2024-01-05,3.1
6,f6a7b8c9-d0e1-2345-6789-012345f01234,user_6,78.5,true,2025-06-01 18:05:00,Passed,101,2024-01-06,4.0
7,a7b8c9d0-e1f2-3456-7890-123456012345,user_7,82.3,true,2025-07-07 13:20:45,,104,2024-01-07,4.1
8,b8c9d0e1-f2a3-4567-8901-234567123456,user_8,91.0,true,2025-08-12 10:00:00,Excellent,103,2024-01-08,4.7
9,c9d0e1f2-a3b4-5678-9012-345678234567,user_9,55.0,false,2025-09-18 20:30:15,Failed,102,2024-01-09,2.5
10,d0e1f2a3-b4c5-6789-0123-456789345678,user_10,89.5,true,2025-10-22 11:11:11,Good work,101,2024-01-10,4.6
11,e1f2a3b4-c5d6-7890-1234-567890456789,user_11,92.0,true,2025-11-28 14:55:00,Solid,103,2024-01-11,4.8
12,f2a3b4c5-d6e7-8901-2345-678901567890,user_12,75.0,true,2025-12-01 08:30:00,,104,2024-01-12,3.9
13,a3b4c5d6-e7f8-9012-3456-789012678901,user_13,68.2,false,2025-01-15 17:45:30,Re-test,102,2024-01-13,3.3
14,b4c5d6e7-f8a9-0123-4567-890123789012,user_14,81.0,true,2025-02-20 12:01:00,Passed,101,2024-01-14,4.3
15,c5d6e7f8-a9b0-1234-5678-901234890123,user_15,77.7,true,2025-03-25 10:10:10,Average,103,2024-01-15,3.7
16,d6e7f8a9-b0c1-2345-6789-012345901234,user_16,94.5,true,2025-04-30 16:20:25,,104,2024-01-16,4.8
17,e7f8a9b0-c1d2-3456-7890-123456012345,user_17,83.0,true,2025-05-05 11:30:00,Good,101,2024-01-17,4.2
18,f8a9b0c1-d2e3-4567-8901-234567123456,user_18,60.5,false,2025-06-10 19:00:50,Needs help,102,2024-01-18,2.9
19,a9b0c1d2-e3f4-5678-9012-345678234567,user_19,88.5,true,2025-07-15 14:45:00,Great,103,2024-01-19,4.6
20,b0c1d2e3-f4a5-6789-0123-456789345678,user_20,90.0,true,2025-08-20 09:00:00,,104,2024-01-20,4.7
21,c1d2e3f4-a5b6-7890-1234-567890456789,user_21,93.2,true,2025-09-25 10:30:00,Top score,101,2024-01-21,4.9
22,d2e3f4a5-b6c7-8901-2345-678901567890,user_22,76.8,true,2025-10-30 11:45:10,Passed,102,2024-01-22,3.9
23,e3f4a5b6-c7d8-9012-3456-789012678901,user_23,69.9,false,2025-11-04 12:00:05,,101,2024-01-23,3.5
24,f4a5b6c7-d8e9-0123-4567-890123789012,user_24,85.0,true,2025-12-10 14:15:30,Good,103,2024-01-24,4.4
25,a5b6c7d8-e9f0-1234-5678-901234890123,user_25,70.0,true,2025-01-15 09:10:20,Review,102,2024-01-25,3.6
26,b6c7d8e9-f0a1-2345-6789-012345901234,user_26,81.5,true,2025-02-20 18:05:00,Passed,101,2024-01-26,4.1
27,c7d8e9f0-a1b2-3456-7890-123456012345,user_27,87.3,true,2025-03-27 13:20:45,,104,2024-01-27,4.5
28,d8e9f0a1-b2c3-4567-8901-234567123456,user_28,94.0,true,2025-04-01 10:00:00,Excellent,103,2024-01-28,4.8
29,e9f0a1b2-c3d4-5678-9012-345678234567,user_29,58.0,false,2025-05-07 20:30:15,Failed,102,2024-01-29,2.8
30,f0a1b2c3-d4e5-6789-0123-456789345678,user_30,90.5,true,2025-06-12 11:11:11,Great work,101,2024-01-30,4.7
31,0a1b2c3d-4e5f-6789-0123-456789abcdef,user_31,96.0,true,2025-07-18 14:55:00,Top,103,2024-01-31,4.9
32,1b2c3d4e-5f6a-7890-1234-567890bcdef0,user_32,74.0,true,2025-08-23 08:30:00,,104,2024-02-01,3.8
33,2c3d4e5f-6a7b-8901-2345-678901cdef01,user_33,67.0,false,2025-09-28 17:45:30,Re-test,102,2024-02-02,3.2
34,3d4e5f6a-7b8c-9012-3456-789012def012,user_34,80.0,true,2025-10-03 12:01:00,Passed,101,2024-02-03,4.2
35,4e5f6a7b-8c9d-0123-4567-890123ef0123,user_35,79.0,true,2025-11-08 10:10:10,Average,103,2024-02-04,3.8
36,5f6a7b8c-9d0e-1234-5678-901234f01234,user_36,95.0,true,2025-12-14 16:20:25,,104,2024-02-05,4.9
37,6a7b8c9d-0e1f-2345-6789-012345012345,user_37,84.5,true,2025-01-20 11:30:00,Good,101,2024-02-06,4.3
38,7b8c9d0e-1f2a-3456-7890-123456123456,user_38,62.0,false,2025-02-25 19:00:50,Needs help,102,2024-02-07,3.0
39,8c9d0e1f-2a3b-4567-8901-234567234567,user_39,87.0,true,2025-03-02 14:45:00,Great,103,2024-02-08,4.5
40,9d0e1f2a-3b4c-5678-9012-345678345678,user_40,91.8,true,2025-04-07 09:00:00,,104,2024-02-09,4.7
41,0e1f2a3b-4c5d-6789-0123-456789456789,user_41,94.2,true,2025-05-13 10:30:00,Top score,101,2024-02-10,4.8
42,1f2a3b4c-5d6e-7890-1234-567890567890,user_42,78.0,true,2025-06-18 11:45:10,Passed,102,2024-02-11,4.0
43,2a3b4c5d-6e7f-8901-2345-678901678901,user_43,71.5,false,2025-07-24 12:00:05,,101,2024-02-12,3.6
44,3b4c5d6e-7f8a-9012-3456-789012789012,user_44,86.3,true,2025-08-29 14:15:30,Good,103,2024-02-13,4.5
45,4c5d6e7f-8a9b-0123-4567-890123890123,user_45,72.0,true,2025-09-03 09:10:20,Review,102,2024-02-14,3.7
46,5d6e7f8a-9b0c-1234-5678-901234901234,user_46,83.0,true,2025-10-09 18:05:00,Passed,101,2024-02-15,4.2
47,6e7f8a9b-0c1d-2345-6789-012345012345,user_47,89.0,true,2025-11-14 13:20:45,,104,2024-02-16,4.6
48,7f8a9b0c-1d2e-3456-7890-123456123456,user_48,96.5,true,2025-12-20 10:00:00,Excellent,103,2024-02-17,4.9
49,8a9b0c1d-2e3f-4567-8901-234567234567,user_49,60.0,false,2025-01-25 20:30:15,Failed,102,2024-02-18,2.9
50,9b0c1d2e-3f4a-5678-9012-345678345678,user_50,92.5,true,2025-02-01 11:11:11,Great work,101,2024-02-19,4.8
51,a1b2c3d4-e5f6-7890-1234-567890abcdef,user_51,95.5,true,2025-01-10 10:30:00,First note,101,2024-01-01,4.5
52,b2c3d4e5-f6a7-8901-2345-678901bcdef0,user_52,88.0,true,2025-02-15 11:45:10,Second note,102,2024-01-02,4.2
53,c3d4e5f6-a7b8-9012-3456-789012cdef01,user_53,72.1,false,2025-03-20 12:00:05,,101,2024-01-03,3.8
54,d4e5f6a7-b8c9-0123-4567-890123def012,user_54,99.9,true,2025-04-25 14:15:30,High score,103,2024-01-04,4.9
55,e5f6a7b8-c9d0-1234-5678-901234ef0123,user_55,65.0,false,2025-05-30 09:10:20,Needs review,102,2024-01-05,3.1
56,f6a7b8c9-d0e1-2345-6789-012345f01234,user_56,78.5,true,2025-06-01 18:05:00,Passed,101,2024-01-06,4.0
57,a7b8c9d0-e1f2-3456-7890-123456012345,user_57,82.3,true,2025-07-07 13:20:45,,104,2024-01-07,4.1
58,b8c9d0e1-f2a3-4567-8901-234567123456,user_58,91.0,true,2025-08-12 10:00:00,Excellent,103,2024-01-08,4.7
59,c9d0e1f2-a3b4-5678-9012-345678234567,user_59,55.0,false,2025-09-18 20:30:15,Failed,102,2024-01-09,2.5
60,d0e1f2a3-b4c5-6789-0123-456789345678,user_60,89.5,true,2025-10-22 11:11:11,Good work,101,2024-01-10,4.6
61,e1f2a3b4-c5d6-7890-1234-567890456789,user_61,92.0,true,2025-11-28 14:55:00,Solid,103,2024-01-11,4.8
62,f2a3b4c5-d6e7-8901-2345-678901567890,user_62,75.0,true,2025-12-01 08:30:00,,104,2024-01-12,3.9
63,a3b4c5d6-e7f8-9012-3456-789012678901,user_63,68.2,false,2025-01-15 17:45:30,Re-test,102,2024-01-13,3.3
64,b4c5d6e7-f8a9-0123-4567-890123789012,user_64,81.0,true,2025-02-20 12:01:00,Passed,101,2024-01-14,4.3
65,c5d6e7f8-a9b0-1234-5678-901234890123,user_65,77.7,true,2025-03-25 10:10:10,Average,103,2024-01-15,3.7
66,d6e7f8a9-b0c1-2345-6789-012345901234,user_66,94.5,true,2025-04-30 16:20:25,,104,2024-01-16,4.8
67,e7f8a9b0-c1d2-3456-7890-123456012345,user_67,83.0,true,2025-05-05 11:30:00,Good,101,2024-01-12,4.2
68,f8a9b0c1-d2e3-4567-8901-234567123456,user_68,60.5,false,2025-06-10 19:00:50,Needs help,102,2024-01-18,2.9
69,a9b0c1d2-e3f4-5678-9012-345678234567,user_69,88.5,true,2025-07-15 14:45:00,Great,103,2024-01-19,4.6
70,b0c1d2e3-f4a5-6789-0123-456789345678,user_70,90.0,true,2025-08-20 09:00:00,,104,2024-01-20,4.7
71,c1d2e3f4-a5b6-7890-1234-567890456789,user_71,93.2,true,2025-09-25 10:30:00,Top score,101,2024-01-21,4.9
72,d2e3f4a5-b6c7-8901-2345-678901567890,user_72,76.8,true,2025-10-30 11:45:10,Passed,102,2024-01-22,3.9
73,e3f4a5b6-c7d8-9012-3456-789012678901,user_73,69.9,false,2025-11-04 12:00:05,,101,2024-01-23,3.5
74,f4a5b6c7-d8e9-0123-4567-890123789012,user_74,85.0,true,2025-12-10 14:15:30,Good,103,2024-01-24,4.4
75,a5b6c7d8-e9f0-1234-5678-901234890123,user_75,70.0,true,2025-01-15 09:10:20,Review,102,2024-01-25,3.6
76,b6c7d8e9-f0a1-2345-6789-012345901234,user_76,81.5,true,2025-02-20 18:05:00,Passed,101,2024-01-26,4.1
77,c7d8e9f0-a1b2-3456-7890-123456012345,user_77,87.3,true,2025-03-27 13:20:45,,104,2024-01-27,4.5
78,d8e9f0a1-b2c3-4567-8901-234567123456,user_78,94.0,true,2025-04-01 10:00:00,Excellent,103,2024-01-28,4.8
79,e9f0a1b2-c3d4-5678-9012-345678234567,user_79,58.0,false,2025-05-07 20:30:15,Failed,102,2024-01-29,2.8
80,f0a1b2c3-d4e5-6789-0123-456789345678,user_80,90.5,true,2025-06-12 11:11:11,Great work,101,2024-01-30,4.7
81,0a1b2c3d-4e5f-6789-0123-456789abcdef,user_81,96.0,true,2025-07-18 14:55:00,Top,103,2024-01-31,4.9
82,1b2c3d4e-5f6a-7890-1234-567890bcdef0,user_82,74.0,true,2025-08-23 08:30:00,,104,2024-02-01,3.8
83,2c3d4e5f-6a7b-8901-2345-678901cdef01,user_83,67.0,false,2025-09-28 17:45:30,Re-test,102,2024-02-02,3.2
84,3d4e5f6a-7b8c-9012-3456-789012def012,user_84,80.0,true,2025-10-03 12:01:00,Passed,101,2024-02-03,4.2
85,4e5f6a7b-8c9d-0123-4567-890123ef0123,user_85,79.0,true,2025-11-08 10:10:10,Average,103,2024-02-04,3.8
86,5f6a7b8c-9d0e-1234-5678-901234f01234,user_86,95.0,true,2025-12-14 16:20:25,,104,2024-02-05,4.9
87,6a7b8c9d-0e1f-2345-6789-012345012345,user_87,84.5,true,2025-01-20 11:30:00,Good,101,2024-02-06,4.3
88,7b8c9d0e-1f2a-3456-7890-123456123456,user_88,62.0,false,2025-02-25 19:00:50,Needs help,102,2024-02-07,3.0
89,8c9d0e1f-2a3b-4567-8901-234567234567,user_89,87.0,true,2025-03-02 14:45:00,Great,103,2024-02-08,4.5
90,9d0e1f2a-3b4c-5678-9012-345678345678,user_90,91.8,true,2025-04-07 09:00:00,,104,2024-02-09,4.7
91,0e1f2a3b-4c5d-6789-0123-456789456789,user_91,94.2,true,2025-05-13 10:30:00,Top score,101,2024-02-10,4.8
92,1f2a3b4c-5d6e-7890-1234-567890567890,user_92,78.0,true,2025-06-18 11:45:10,Passed,102,2024-02-11,4.0
93,2a3b4c5d-6e7f-8901-2345-678901678901,user_93,71.5,false,2025-07-24 12:00:05,,101,2024-02-12,3.6
94,3b4c5d6e-7f8a-9012-3456-789012789012,user_94,86.3,true,2025-08-29 14:15:30,Good,103,2024-02-13,4.5
95,4c5d6e7f-8a9b-0123-4567-890123890123,user_95,72.0,true,2025-09-03 09:10:20,Review,102,2024-02-14,3.7
96,5d6e7f8a-9b0c-1234-5678-901234901234,user_96,83.0,true,2025-10-09 18:05:00,Passed,101,2024-02-15,4.2
97,6e7f8a9b-0c1d-2345-6789-012345012345,user_97,89.0,true,2025-11-14 13:20:45,,104,2024-02-16,4.6
98,7f8a9b0c-1d2e-3456-7890-123456123456,user_98,96.5,true,2025-12-20 10:00:00,Excellent,103,2024-02-17,4.9
99,8a9b0c1d-2e3f-4567-8901-234567234567,user_99,60.0,false,2025-01-25 20:30:15,Failed,102,2024-02-18,2.9
100,9b0c1d2e-3f4a-5678-9012-345678345678,user_100,92.5,true,2025-02-01 11:11:11,Great work,101,2024-02-19,4.8
\.
-- ```
-- **To run this script:**
-- 1. Save it as `setup_db.sql`.
-- 2. Use the `psql` command-line tool (which comes with Postgres) to execute it against your database (the one running in Docker):
-- ```bash
-- psql "postgresql://postgres:mysecretpassword@localhost:5433/postgres" -f setup_db.sql

61
src/config.rs
View File

@@ -0,0 +1,61 @@
// --- External Crates ---
use serde::Deserialize;
use std::collections::HashMap;
use anyhow::{Context, Result, anyhow};
use std::fmt::{self, Display};
// --- 1. CUSTOM ERROR DEFINITION ---
#[derive(Debug)]
pub struct ConfigError(String);
impl Display for ConfigError {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "Configuration Error: {}", self.0)
}
}
impl std::error::Error for ConfigError {}
// --- 2. CONFIGURATION STRUCTS (The "Configured Opinion") ---
// We make these 'pub' (public) so src/lib.rs can use them.
#[derive(Debug, Deserialize, Clone)]
pub struct ColumnConfig {
pub column_name: String,
pub arrow_type: String,
}
#[derive(Debug, Deserialize)]
pub struct ConnectorConfig {
pub connection_string: String,
pub query: String, // <--- ADD THIS FIELD
pub schema: Vec<ColumnConfig>,
}
// --- 3. CONFIG LOADING FUNCTION ---
// This is also 'pub' so src/lib.rs can call it.
pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> {
let file_content = std::fs::read_to_string(path)
.context(format!("Failed to read config file at path: {}", path))?;
let config: ConnectorConfig = serde_yaml::from_str(&file_content)
.context("Failed to deserialize YAML configuration")?;
if config.connection_string.is_empty() {
return Err(ConfigError("Connection string cannot be empty.".to_string())).map_err(anyhow::Error::from)?;
}
if config.query.is_empty() {
return Err(ConfigError("Query cannot be empty.".to_string())).map_err(anyhow::Error::from)?;
}
// FIX: Make the COPY check more flexible (allows newlines)
let uppercase_query = config.query.trim().to_uppercase();
if !uppercase_query.starts_with("COPY") || !uppercase_query.contains("TO STDOUT") || !uppercase_query.contains("FORMAT BINARY") {
return Err(ConfigError("Query must be a 'COPY ... TO STDOUT (FORMAT binary)' command.".to_string())).map_err(anyhow::Error::from)?;
}
if config.schema.is_empty() {
return Err(ConfigError("Schema cannot be empty.".to_string())).map_err(anyhow::Error::from)?;
}
Ok(config)
}

271
src/parser.rs
View File

@@ -0,0 +1,271 @@
// --- External Crates ---
use std::pin::Pin;
use std::sync::Arc;
use tokio_postgres::{NoTls, CopyOutStream};
use anyhow::{Context, Result, anyhow};
// FIX: Import new builders
use arrow::array::{
ArrayBuilder, ArrayRef,
Int64Builder, Float64Builder, Float32Builder, StringBuilder, BooleanBuilder,
TimestampNanosecondBuilder, Date32Builder, Int32Builder
};
// FIX: Import new types
use arrow::datatypes::{
DataType, Field, Schema,
Float64Type, Float32Type, Int64Type, Int32Type, Utf8Type, BooleanType, TimestampNanosecondType,
Date32Type
};
use arrow::record_batch::RecordBatch;
use futures_util::stream::StreamExt;
use bytes::Bytes;
use byteorder::{BigEndian, ReadBytesExt};
use std::io::{Cursor, Read};
use std::str;
use chrono::{NaiveDateTime, NaiveDate};
// --- Internal Crates ---
use crate::config::ConnectorConfig;
// --- 1. CORE DATABASE LOGIC ---
pub async fn run_db_logic(config: ConnectorConfig) -> Result<()> {
// 1. Establish the connection
println!("UncheckedIO: Attempting connection...");
let (client, connection) = tokio_postgres::connect(&config.connection_string, NoTls).await
.context("Failed to connect to PostgreSQL database")?;
tokio::spawn(async move {
if let Err(e) = connection.await {
eprintln!("Postgres connection error: {}", e);
}
});
// 2. Execute the COPY TO STDOUT command
// FIX: We are now using the query from the config file!
let copy_query = &config.query;
println!("UncheckedIO: Executing user-defined query...");
let copy_stream = client.copy_out(copy_query.as_str()).await
.context("Failed to execute COPY TO STDOUT protocol. Check your query syntax and permissions.")?;
// 3. Handle the Binary Stream
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
// Build the Arrow Schema from the config
let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| {
// We now allow "nullable" to be controlled by the column name, a temporary "hack"
let nullable = col_cfg.column_name == "notes";
let arrow_type = match col_cfg.arrow_type.as_str() {
"Int64" => DataType::Int64,
"Int32" => DataType::Int32, // NEW
"Float64" => DataType::Float64, // NEW
"Float32" => DataType::Float32,
"Utf8" | "String" => DataType::Utf8,
"Boolean" => DataType::Boolean,
"Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
"Date32" => DataType::Date32, // NEW
_ => panic!("Unsupported type in config: {}", col_cfg.arrow_type),
};
Field::new(&col_cfg.column_name, arrow_type, nullable)
}).collect();
let arrow_schema = Arc::new(Schema::new(schema_fields));
// Call the parser
let (rows_processed, record_batch) = handle_binary_copy(pinned_stream, arrow_schema.clone()).await?;
println!("UncheckedIO: Successfully parsed {} rows via binary stream.", rows_processed);
// 4. Final Output Confirmation
println!("UncheckedIO: Built RecordBatch with {} rows and {} columns.",
record_batch.num_rows(), record_batch.num_columns());
println!("UncheckedIO: Data transfer complete. We lived.");
Ok(())
}
// --- 2. BINARY STREAM HANDLER (THE DYNAMIC PARSER) ---
// This enum will hold our different builder types
enum DynamicBuilder {
Int64(Box<Int64Builder>),
Int32(Box<Int32Builder>), // NEW
Float64(Box<Float64Builder>), // NEW
Float32(Box<Float32Builder>),
String(Box<StringBuilder>),
Boolean(Box<BooleanBuilder>),
Timestamp(Box<TimestampNanosecondBuilder>),
Date32(Box<Date32Builder>), // NEW
}
// Postgres Epoch for timestamps
const POSTGRES_EPOCH_NAIVE: NaiveDateTime = NaiveDate::from_ymd_opt(2000, 1, 1).unwrap().and_hms_opt(0, 0, 0).unwrap();
// Unix Epoch for dates
const UNIX_EPOCH_NAIVE_DATE: NaiveDate = NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
// This function is now private to this module (it's not 'pub')
async fn handle_binary_copy(
mut stream: Pin<Box<CopyOutStream>>,
arrow_schema: Arc<Schema>
) -> Result<(usize, RecordBatch)> {
let mut binary_buffer: Vec<u8> = Vec::new();
// 1. Aggregate all bytes from the stream into our buffer
while let Some(segment_result) = stream.next().await {
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;
binary_buffer.extend_from_slice(&segment);
}
let bytes_received = binary_buffer.len();
if bytes_received == 0 {
return Err(anyhow!("Received zero bytes from COPY stream. Check query and permissions."));
}
println!("UncheckedIO: Received {} total bytes from stream.", bytes_received);
// --- PARSE POSTGRES BINARY HEADER ---
let mut cursor = Cursor::new(&binary_buffer[..]);
let mut magic_signature = [0u8; 11];
cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?;
if &magic_signature != b"PGCOPY\n\xff\r\n\0" {
return Err(anyhow!("Invalid Postgres COPY binary signature."));
}
let _flags = cursor.read_u32::<BigEndian>().context("Failed to read flags")?;
let _header_ext_len = cursor.read_u32::<BigEndian>().context("Failed to read header extension length")?;
println!("UncheckedIO: Postgres binary header validated.");
// --- DYNAMIC FIELD DESERIALIZATION ---
// 1. Create a dynamic list of builders based on the schema
let mut builders: Vec<DynamicBuilder> = arrow_schema.fields().iter().map(|field| {
match field.data_type() {
DataType::Int64 => DynamicBuilder::Int64(Box::new(Int64Builder::new())),
DataType::Int32 => DynamicBuilder::Int32(Box::new(Int32Builder::new())),
DataType::Float64 => DynamicBuilder::Float64(Box::new(Float64Builder::new())),
DataType::Float32 => DynamicBuilder::Float32(Box::new(Float32Builder::new())),
DataType::Utf8 => DynamicBuilder::String(Box::new(StringBuilder::new())),
DataType::Boolean => DynamicBuilder::Boolean(Box::new(BooleanBuilder::new())),
DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None) => {
DynamicBuilder::Timestamp(Box::new(TimestampNanosecondBuilder::new()))
},
DataType::Date32 => DynamicBuilder::Date32(Box::new(Date32Builder::new())),
_ => panic!("Unsupported type in builder creation!"),
}
}).collect();
let mut rows_processed = 0;
// 2. Loop through the rest of the buffer until we hit the 2-byte trailer
loop {
let col_count = match cursor.read_i16::<BigEndian>() {
Ok(count) => count,
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
break; // Reached end of buffer *before* trailer, assume it's the end
}
Err(e) => return Err(e.into()),
};
// Check for the 2-byte trailer (-1) which signals the end of the data
if col_count == -1 {
println!("UncheckedIO: Reached end-of-stream trailer.");
break; // End of stream
}
rows_processed += 1;
// 3. Dynamic Row-Parsing Loop
for (i, builder) in builders.iter_mut().enumerate() {
let field_len_i32 = cursor.read_i32::<BigEndian>()
.context(format!("Failed to read field length for col {}", i))?;
if field_len_i32 == -1 {
// Handle NULLs
match builder {
DynamicBuilder::Int64(b) => b.append_null(),
DynamicBuilder::Float64(b) => b.append_null(),
DynamicBuilder::Int32(b) => b.append_null(),
DynamicBuilder::Float32(b) => b.append_null(),
DynamicBuilder::String(b) => b.append_null(),
DynamicBuilder::Boolean(b) => b.append_null(),
DynamicBuilder::Timestamp(b) => b.append_null(),
DynamicBuilder::Date32(b) => b.append_null(),
}
continue; // Go to the next field
}
let field_len_usize = field_len_i32 as usize;
match builder {
DynamicBuilder::Int64(b) => {
let val = cursor.read_i64::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::Int32(b) => {
let val = cursor.read_i32::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::Float64(b) => {
let val = cursor.read_f64::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::Float32(b) => {
let val = cursor.read_f32::<BigEndian>()?;
b.append_value(val);
}
DynamicBuilder::String(b) => {
let mut str_buf = vec![0; field_len_usize];
cursor.read_exact(&mut str_buf)?;
let val_str = str::from_utf8(&str_buf)
.context(format!("Failed to parse UTF-8 string for col {}", i))?;
b.append_value(val_str);
}
DynamicBuilder::Boolean(b) => {
let val_bool = cursor.read_u8()?;
b.append_value(val_bool != 0);
}
DynamicBuilder::Timestamp(b) => {
let pg_micros = cursor.read_i64::<BigEndian>()?;
let unix_epoch = NaiveDateTime::from_timestamp_opt(0, 0).unwrap();
let pg_epoch = POSTGRES_EPOCH_NAIVE;
let epoch_delta_micros = (pg_epoch - unix_epoch).num_microseconds().unwrap();
let unix_micros = epoch_delta_micros + pg_micros;
let unix_nanos = unix_micros * 1000;
b.append_value(unix_nanos);
}
DynamicBuilder::Date32(b) => {
// Postgres binary DATE is i32 days since 2000-01-01
let pg_days = cursor.read_i32::<BigEndian>()?;
// Arrow Date32 is i32 days since 1970-01-01 (Unix Epoch)
let epoch_delta_days = (POSTGRES_EPOCH_NAIVE.date() - UNIX_EPOCH_NAIVE_DATE).num_days() as i32;
let unix_days = epoch_delta_days + pg_days;
b.append_value(unix_days);
}
}
}
}
// 4. Finalize the Arrays
let final_columns: Vec<ArrayRef> = builders.into_iter().map(|builder| {
match builder {
DynamicBuilder::Int64(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Int32(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Float64(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Float32(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::String(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Boolean(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Timestamp(mut b) => Arc::new(b.finish()) as ArrayRef,
DynamicBuilder::Date32(mut b) => Arc::new(b.finish()) as ArrayRef,
}
}).collect();
// 5. Build the Final RecordBatch
let record_batch = RecordBatch::try_new(
arrow_schema.clone(),
final_columns,
).context("Failed to create final Arrow RecordBatch")?;
Ok((rows_processed, record_batch))
}