1use core::ptr;
5
6use reifydb_abi::{
7 callbacks::builder::{ColumnBufferHandle, EmitDiffKind},
8 context::context::ContextFFI,
9 data::column::ColumnTypeCode,
10};
11use reifydb_value::value::row_number::RowNumber;
12
13use crate::{error::SdkError, operator::context::ffi::FFIOperatorContext};
14
15pub struct ColumnBuilder<'a> {
16 ctx: *mut ContextFFI,
17 handle: *mut ColumnBufferHandle,
18 type_code: ColumnTypeCode,
19 committed: bool,
20 _phantom: core::marker::PhantomData<&'a ()>,
21}
22
23#[derive(Clone, Copy)]
24pub struct CommittedColumn {
25 handle: *mut ColumnBufferHandle,
26 row_count: usize,
27}
28
29impl<'a> ColumnBuilder<'a> {
30 pub fn data_ptr(&self) -> *mut u8 {
31 unsafe {
32 let cb = (*self.ctx).callbacks.builder;
33 (cb.data_ptr)(self.handle)
34 }
35 }
36
37 pub fn offsets_ptr(&self) -> *mut u64 {
38 unsafe {
39 let cb = (*self.ctx).callbacks.builder;
40 (cb.offsets_ptr)(self.handle)
41 }
42 }
43
44 pub fn bitvec_ptr(&self) -> *mut u8 {
45 unsafe {
46 let cb = (*self.ctx).callbacks.builder;
47 (cb.bitvec_ptr)(self.handle)
48 }
49 }
50
51 pub fn grow(&self, additional: usize) -> Result<(), SdkError> {
52 let code = unsafe {
53 let cb = (*self.ctx).callbacks.builder;
54 (cb.grow)(self.handle, additional)
55 };
56 if code != 0 {
57 return Err(SdkError::Other(format!("ColumnBuilder::grow failed: {}", code)));
58 }
59 Ok(())
60 }
61
62 pub fn commit(mut self, written_count: usize) -> Result<CommittedColumn, SdkError> {
63 let code = unsafe {
64 let cb = (*self.ctx).callbacks.builder;
65 (cb.commit)(self.handle, written_count)
66 };
67 self.committed = true;
68 if code != 0 {
69 return Err(SdkError::Other(format!("ColumnBuilder::commit failed: {}", code)));
70 }
71 Ok(CommittedColumn {
72 handle: self.handle,
73 row_count: written_count,
74 })
75 }
76
77 pub fn type_code(&self) -> ColumnTypeCode {
78 self.type_code
79 }
80
81 pub fn write_bool(self, values: &[bool]) -> Result<CommittedColumn, SdkError> {
82 debug_assert_eq!(self.type_code, ColumnTypeCode::Bool, "write_bool requires a Bool ColumnBuilder");
83
84 let byte_count = values.len().div_ceil(8);
85 let buffer_byte_len = values.len();
86 let mut packed = vec![0u8; buffer_byte_len.max(byte_count)];
87 for (i, &b) in values.iter().enumerate() {
88 if b {
89 packed[i / 8] |= 1 << (i % 8);
90 }
91 }
92 if !packed.is_empty() {
93 unsafe {
94 core::ptr::copy_nonoverlapping(packed.as_ptr(), self.data_ptr(), packed.len());
95 }
96 }
97 self.commit(values.len())
98 }
99
100 pub fn write_f32(self, values: &[f32]) -> Result<CommittedColumn, SdkError> {
101 debug_assert_eq!(self.type_code, ColumnTypeCode::Float4);
102 unsafe { write_scalar(self, values) }
103 }
104
105 pub fn write_f64(self, values: &[f64]) -> Result<CommittedColumn, SdkError> {
106 debug_assert_eq!(self.type_code, ColumnTypeCode::Float8);
107 unsafe { write_scalar(self, values) }
108 }
109
110 pub fn write_i8(self, values: &[i8]) -> Result<CommittedColumn, SdkError> {
111 debug_assert_eq!(self.type_code, ColumnTypeCode::Int1);
112 unsafe { write_scalar(self, values) }
113 }
114
115 pub fn write_i16(self, values: &[i16]) -> Result<CommittedColumn, SdkError> {
116 debug_assert_eq!(self.type_code, ColumnTypeCode::Int2);
117 unsafe { write_scalar(self, values) }
118 }
119
120 pub fn write_i32(self, values: &[i32]) -> Result<CommittedColumn, SdkError> {
121 debug_assert_eq!(self.type_code, ColumnTypeCode::Int4);
122 unsafe { write_scalar(self, values) }
123 }
124
125 pub fn write_i64(self, values: &[i64]) -> Result<CommittedColumn, SdkError> {
126 debug_assert_eq!(self.type_code, ColumnTypeCode::Int8);
127 unsafe { write_scalar(self, values) }
128 }
129
130 pub fn write_i128(self, values: &[i128]) -> Result<CommittedColumn, SdkError> {
131 debug_assert_eq!(self.type_code, ColumnTypeCode::Int16);
132 unsafe { write_scalar(self, values) }
133 }
134
135 pub fn write_u8(self, values: &[u8]) -> Result<CommittedColumn, SdkError> {
136 debug_assert_eq!(self.type_code, ColumnTypeCode::Uint1);
137 unsafe { write_scalar(self, values) }
138 }
139
140 pub fn write_u16(self, values: &[u16]) -> Result<CommittedColumn, SdkError> {
141 debug_assert_eq!(self.type_code, ColumnTypeCode::Uint2);
142 unsafe { write_scalar(self, values) }
143 }
144
145 pub fn write_u32(self, values: &[u32]) -> Result<CommittedColumn, SdkError> {
146 debug_assert_eq!(self.type_code, ColumnTypeCode::Uint4);
147 unsafe { write_scalar(self, values) }
148 }
149
150 pub fn write_u64(self, values: &[u64]) -> Result<CommittedColumn, SdkError> {
151 debug_assert_eq!(self.type_code, ColumnTypeCode::Uint8);
152 unsafe { write_scalar(self, values) }
153 }
154
155 pub fn write_u128(self, values: &[u128]) -> Result<CommittedColumn, SdkError> {
156 debug_assert_eq!(self.type_code, ColumnTypeCode::Uint16);
157 unsafe { write_scalar(self, values) }
158 }
159
160 pub fn write_utf8<S: AsRef<str>>(self, values: &[S]) -> Result<CommittedColumn, SdkError> {
161 debug_assert_eq!(self.type_code, ColumnTypeCode::Utf8, "write_utf8 requires a Utf8 ColumnBuilder");
162 write_var_len(self, values.iter().map(|s| s.as_ref().as_bytes()))
163 }
164
165 pub fn write_blob<B: AsRef<[u8]>>(self, values: &[B]) -> Result<CommittedColumn, SdkError> {
166 debug_assert_eq!(self.type_code, ColumnTypeCode::Blob, "write_blob requires a Blob ColumnBuilder");
167 write_var_len(self, values.iter().map(|b| b.as_ref()))
168 }
169
170 pub fn set_defined(&self, defined: &[bool]) {
171 let bytes = defined.len().div_ceil(8);
172 if bytes == 0 {
173 return;
174 }
175 let mut packed = vec![0u8; bytes];
176 for (i, &b) in defined.iter().enumerate() {
177 if b {
178 packed[i / 8] |= 1 << (i % 8);
179 }
180 }
181 unsafe {
182 core::ptr::copy_nonoverlapping(packed.as_ptr(), self.bitvec_ptr(), bytes);
183 }
184 }
185}
186
187unsafe fn write_scalar<T: Copy>(col: ColumnBuilder<'_>, values: &[T]) -> Result<CommittedColumn, SdkError> {
188 let bytes = core::mem::size_of_val(values);
189 if bytes > 0 {
190 unsafe {
191 core::ptr::copy_nonoverlapping(values.as_ptr() as *const u8, col.data_ptr(), bytes);
192 }
193 }
194 col.commit(values.len())
195}
196
197fn write_var_len<'b, I>(col: ColumnBuilder<'_>, items: I) -> Result<CommittedColumn, SdkError>
198where
199 I: IntoIterator<Item = &'b [u8]>,
200{
201 let items: Vec<&[u8]> = items.into_iter().collect();
202 let total: usize = items.iter().map(|b| b.len()).sum();
203 let needed = total.max(items.len());
204 if needed > 0 {
205 col.grow(needed)?;
206 }
207 let mut cursor = 0usize;
208 unsafe {
209 let data = col.data_ptr();
210 let offsets = col.offsets_ptr();
211 core::ptr::write(offsets, 0u64);
212 for (i, bytes) in items.iter().enumerate() {
213 if !bytes.is_empty() {
214 core::ptr::copy_nonoverlapping(bytes.as_ptr(), data.add(cursor), bytes.len());
215 }
216 cursor += bytes.len();
217 core::ptr::write(offsets.add(i + 1), cursor as u64);
218 }
219 }
220 col.commit(items.len())
221}
222
223impl<'a> Drop for ColumnBuilder<'a> {
224 fn drop(&mut self) {
225 if !self.committed {
226 unsafe {
227 let cb = (*self.ctx).callbacks.builder;
228 (cb.release)(self.handle);
229 }
230 }
231 }
232}
233
234pub struct ColumnsBuilder<'a> {
235 ctx: *mut ContextFFI,
236 _phantom: core::marker::PhantomData<&'a mut ()>,
237}
238
239impl<'a> ColumnsBuilder<'a> {
240 pub fn new(ctx: &'a mut FFIOperatorContext) -> Self {
241 Self {
242 ctx: ctx.ctx,
243 _phantom: core::marker::PhantomData,
244 }
245 }
246
247 pub fn from_raw_ctx(ctx: *mut ContextFFI) -> Self {
248 Self {
249 ctx,
250 _phantom: core::marker::PhantomData,
251 }
252 }
253
254 pub fn acquire(&mut self, type_code: ColumnTypeCode, capacity: usize) -> Result<ColumnBuilder<'_>, SdkError> {
255 let handle = unsafe {
256 let cb = (*self.ctx).callbacks.builder;
257 (cb.acquire)(self.ctx, type_code, capacity)
258 };
259 if handle.is_null() {
260 return Err(SdkError::Other(format!(
261 "ColumnsBuilder::acquire failed for type {:?}",
262 type_code
263 )));
264 }
265 Ok(ColumnBuilder {
266 ctx: self.ctx,
267 handle,
268 type_code,
269 committed: false,
270 _phantom: core::marker::PhantomData,
271 })
272 }
273
274 pub fn emit_insert(
275 &mut self,
276 post: &[CommittedColumn],
277 names: &[&str],
278 row_numbers: &[RowNumber],
279 ) -> Result<(), SdkError> {
280 assert_eq!(post.len(), names.len(), "emit_insert: post columns and names must have matching length");
281 let row_count = post.first().map(|c| c.row_count).unwrap_or(0);
282 assert_eq!(row_numbers.len(), row_count, "emit_insert: row_numbers length must equal post row count");
283 self.emit_internal(EmitDiffKind::Insert, &[], &[], 0, &[], post, names, row_count, row_numbers)
284 }
285
286 #[allow(clippy::too_many_arguments)]
287 pub fn emit_update(
288 &mut self,
289 pre: &[CommittedColumn],
290 pre_names: &[&str],
291 pre_row_count: usize,
292 pre_row_numbers: &[RowNumber],
293 post: &[CommittedColumn],
294 post_names: &[&str],
295 post_row_count: usize,
296 post_row_numbers: &[RowNumber],
297 ) -> Result<(), SdkError> {
298 assert_eq!(pre.len(), pre_names.len(), "emit_update: pre columns/names mismatch");
299 assert_eq!(post.len(), post_names.len(), "emit_update: post columns/names mismatch");
300 assert_eq!(pre_row_numbers.len(), pre_row_count, "emit_update: pre_row_numbers length mismatch");
301 assert_eq!(post_row_numbers.len(), post_row_count, "emit_update: post_row_numbers length mismatch");
302 self.emit_internal(
303 EmitDiffKind::Update,
304 pre,
305 pre_names,
306 pre_row_count,
307 pre_row_numbers,
308 post,
309 post_names,
310 post_row_count,
311 post_row_numbers,
312 )
313 }
314
315 pub fn emit_remove(
316 &mut self,
317 pre: &[CommittedColumn],
318 names: &[&str],
319 row_numbers: &[RowNumber],
320 ) -> Result<(), SdkError> {
321 assert_eq!(pre.len(), names.len(), "emit_remove: pre columns and names must have matching length");
322 let row_count = pre.first().map(|c| c.row_count).unwrap_or(0);
323 assert_eq!(row_numbers.len(), row_count, "emit_remove: row_numbers length must equal pre row count");
324 self.emit_internal(EmitDiffKind::Remove, pre, names, row_count, row_numbers, &[], &[], 0, &[])
325 }
326
327 #[allow(clippy::too_many_arguments)]
328 fn emit_internal(
329 &mut self,
330 kind: EmitDiffKind,
331 pre: &[CommittedColumn],
332 pre_names: &[&str],
333 pre_row_count: usize,
334 pre_row_numbers: &[RowNumber],
335 post: &[CommittedColumn],
336 post_names: &[&str],
337 post_row_count: usize,
338 post_row_numbers: &[RowNumber],
339 ) -> Result<(), SdkError> {
340 let pre_handles: Vec<*mut ColumnBufferHandle> = pre.iter().map(|c| c.handle).collect();
341 let pre_name_ptrs: Vec<*const u8> = pre_names.iter().map(|n| n.as_ptr()).collect();
342 let pre_name_lens: Vec<usize> = pre_names.iter().map(|n| n.len()).collect();
343 let pre_row_nums: Vec<u64> = pre_row_numbers.iter().map(|r| r.0).collect();
344 let post_handles: Vec<*mut ColumnBufferHandle> = post.iter().map(|c| c.handle).collect();
345 let post_name_ptrs: Vec<*const u8> = post_names.iter().map(|n| n.as_ptr()).collect();
346 let post_name_lens: Vec<usize> = post_names.iter().map(|n| n.len()).collect();
347 let post_row_nums: Vec<u64> = post_row_numbers.iter().map(|r| r.0).collect();
348
349 let code = unsafe {
350 let cb = (*self.ctx).callbacks.builder;
351 (cb.emit_diff)(
352 self.ctx,
353 kind,
354 if pre_handles.is_empty() {
355 ptr::null()
356 } else {
357 pre_handles.as_ptr()
358 },
359 if pre_name_ptrs.is_empty() {
360 ptr::null()
361 } else {
362 pre_name_ptrs.as_ptr()
363 },
364 if pre_name_lens.is_empty() {
365 ptr::null()
366 } else {
367 pre_name_lens.as_ptr()
368 },
369 pre_handles.len(),
370 pre_row_count,
371 if pre_row_nums.is_empty() {
372 ptr::null()
373 } else {
374 pre_row_nums.as_ptr()
375 },
376 pre_row_nums.len(),
377 if post_handles.is_empty() {
378 ptr::null()
379 } else {
380 post_handles.as_ptr()
381 },
382 if post_name_ptrs.is_empty() {
383 ptr::null()
384 } else {
385 post_name_ptrs.as_ptr()
386 },
387 if post_name_lens.is_empty() {
388 ptr::null()
389 } else {
390 post_name_lens.as_ptr()
391 },
392 post_handles.len(),
393 post_row_count,
394 if post_row_nums.is_empty() {
395 ptr::null()
396 } else {
397 post_row_nums.as_ptr()
398 },
399 post_row_nums.len(),
400 )
401 };
402 if code != 0 {
403 return Err(SdkError::Other(format!("emit_diff failed: {}", code)));
404 }
405 Ok(())
406 }
407}