xlog_cuda/provider/

io.rs

//! I/O operations: multi-column buffer creation, Arrow interop.

use std::sync::Arc;

use crate::{LaunchAsync, LaunchConfig};
use xlog_core::{Result, ScalarType, Schema, XlogError};

use super::{d4_kernels, pack_kernels, D4_MODULE, PACK_MODULE};
use crate::{CudaBuffer, CudaColumn};

fn host_columns_form_canonical_u32_set(schema: &Schema, columns: &[Vec<u8>], rows: usize) -> bool {
    if rows <= 1 {
        return true;
    }
    if schema.arity() == 0
        || !(0..schema.arity()).all(|index| {
            matches!(
                schema.column_type(index),
                Some(ScalarType::U32 | ScalarType::Symbol)
            )
        })
    {
        return false;
    }

    (1..rows).all(|row| {
        columns
            .iter()
            .map(|column| {
                let prior = (row - 1) * std::mem::size_of::<u32>();
                let current = row * std::mem::size_of::<u32>();
                let prior_value = u32::from_le_bytes(
                    column[prior..prior + 4]
                        .try_into()
                        .expect("validated u32 column"),
                );
                let current_value = u32::from_le_bytes(
                    column[current..current + 4]
                        .try_into()
                        .expect("validated u32 column"),
                );
                prior_value.cmp(&current_value)
            })
            .find(|ordering| *ordering != std::cmp::Ordering::Equal)
            == Some(std::cmp::Ordering::Less)
    })
}

impl super::CudaKernelProvider {
    // ============== Buffer Helper Methods ==============

    /// Create a CudaBuffer from multiple u32 column slices
    ///
    /// # Arguments
    /// * `columns` - Slice of column data slices (each column as &[u32])
    /// * `schema` - The schema for the buffer
    ///
    /// # Returns
    /// A new CudaBuffer containing all columns
    ///
    /// # Errors
    /// Returns `XlogError::Kernel` if upload fails or columns have mismatched lengths
    pub fn create_buffer_from_u32_columns(
        &self,
        columns: &[&[u32]],
        schema: Schema,
    ) -> Result<CudaBuffer> {
        if columns.is_empty() {
            return self.create_empty_buffer(schema);
        }

        let num_rows = columns[0].len();
        for (i, col) in columns.iter().enumerate() {
            if col.len() != num_rows {
                return Err(XlogError::Kernel(format!(
                    "Column {} has {} rows but expected {}",
                    i,
                    col.len(),
                    num_rows
                )));
            }
        }

        let mut cuda_columns = Vec::with_capacity(columns.len());
        for col_data in columns {
            let bytes: Vec<u8> = col_data.iter().flat_map(|v| v.to_le_bytes()).collect();
            let mut col = self.memory.alloc::<u8>(bytes.len())?;
            self.htod_sync_copy_into_tracked(&bytes, &mut col)
                .map_err(|e| XlogError::Kernel(format!("Failed to upload column: {}", e)))?;
            cuda_columns.push(col.into());
        }

        self.buffer_from_columns(cuda_columns, num_rows as u64, schema)
    }

    /// Create a buffer from multiple column slices (raw bytes)
    ///
    /// This is a generic version that works with any column type by accepting
    /// raw byte slices. Each slice should contain the column data in little-endian
    /// format with the correct size for the column's type.
    ///
    /// # Arguments
    /// * `slices` - Slice of raw byte slices, one per column
    /// * `schema` - The schema for the buffer
    ///
    /// # Returns
    /// A new CudaBuffer containing all columns
    ///
    /// # Errors
    /// Returns `XlogError::Kernel` if:
    /// - Number of slices doesn't match schema arity
    /// - Upload fails
    pub fn create_buffer_from_slices(
        &self,
        slices: &[&[u8]],
        schema: Schema,
    ) -> Result<CudaBuffer> {
        if slices.len() != schema.arity() {
            return Err(XlogError::Kernel(format!(
                "Slice count {} doesn't match schema arity {}",
                slices.len(),
                schema.arity()
            )));
        }

        if slices.is_empty() {
            return self.create_empty_buffer(schema);
        }

        let first_col_size = schema.column_type(0).map(|t| t.size_bytes()).unwrap_or(4);
        let num_rows = slices[0].len() / first_col_size;

        // Validate that all columns have consistent row counts
        for (i, slice) in slices.iter().enumerate() {
            let col_size = schema.column_type(i).map(|t| t.size_bytes()).unwrap_or(4);
            let col_rows = slice.len() / col_size;
            if col_rows != num_rows {
                return Err(XlogError::Kernel(format!(
                    "Column {} has {} rows but expected {} rows (based on first column)",
                    i, col_rows, num_rows
                )));
            }
            // Also verify the slice length is exactly divisible by the type size
            if slice.len() % col_size != 0 {
                return Err(XlogError::Kernel(format!(
                    "Column {} slice length {} is not divisible by type size {}",
                    i,
                    slice.len(),
                    col_size
                )));
            }
        }

        let mut columns = Vec::with_capacity(slices.len());

        for (i, slice) in slices.iter().enumerate() {
            let mut col = self.memory.alloc::<u8>(slice.len())?;
            self.htod_sync_copy_into_tracked(slice, &mut col)
                .map_err(|e| XlogError::Kernel(format!("Failed to upload column {}: {}", i, e)))?;
            columns.push(col.into());
        }

        self.buffer_from_columns(columns, num_rows as u64, schema)
    }

    /// Upload a group of host-columnar relations with one payload transfer and
    /// one row-count transfer. Each returned buffer owns disjoint views of the
    /// shared device allocations, so dropping one relation cannot invalidate
    /// another. Strictly sorted, unique U32/Symbol rows receive a canonical
    /// full-row set proof after their host columns are checked.
    pub fn create_buffers_from_host_columns(
        &self,
        relations: &[(Schema, Vec<Vec<u8>>, usize)],
    ) -> Result<Vec<CudaBuffer>> {
        if relations.is_empty() {
            return Ok(Vec::new());
        }

        let count_capacity = relations
            .len()
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| XlogError::Kernel("Host relation count size overflow".into()))?;
        let mut payload_bytes = 0usize;
        let mut spans = Vec::with_capacity(relations.len());
        let mut count_bytes = Vec::with_capacity(count_capacity);
        let mut canonical_sets = Vec::with_capacity(relations.len());
        for (schema, columns, rows) in relations {
            if columns.len() != schema.arity() {
                return Err(XlogError::Kernel(format!(
                    "Column count {} doesn't match schema arity {}",
                    columns.len(),
                    schema.arity()
                )));
            }
            let rows_u32 = u32::try_from(*rows)
                .map_err(|_| XlogError::Kernel(format!("Row count {} exceeds u32::MAX", rows)))?;
            count_bytes.extend_from_slice(&rows_u32.to_le_bytes());

            let mut column_spans = Vec::with_capacity(columns.len());
            for (index, column) in columns.iter().enumerate() {
                let width = schema
                    .column_type(index)
                    .map(|typ| typ.size_bytes())
                    .unwrap_or(4);
                let expected = rows
                    .checked_mul(width)
                    .ok_or_else(|| XlogError::Kernel("Host column byte length overflow".into()))?;
                if column.len() != expected {
                    return Err(XlogError::Kernel(format!(
                        "Column {} has {} bytes but expected {} for {} rows",
                        index,
                        column.len(),
                        expected,
                        rows
                    )));
                }
                // A later u64/f64 view may borrow this byte slice directly.
                // Keep every column start aligned even after a shorter column.
                let alignment = width.max(std::mem::align_of::<u64>());
                let remainder = payload_bytes % alignment;
                let padding = if remainder == 0 {
                    0
                } else {
                    alignment - remainder
                };
                let start = payload_bytes.checked_add(padding).ok_or_else(|| {
                    XlogError::Kernel("Host relation payload size overflow".into())
                })?;
                let end = start.checked_add(column.len()).ok_or_else(|| {
                    XlogError::Kernel("Host relation payload size overflow".into())
                })?;
                column_spans.push(start..end);
                payload_bytes = end;
            }
            canonical_sets.push(host_columns_form_canonical_u32_set(schema, columns, *rows));
            spans.push(column_spans);
        }

        let mut payload = Vec::with_capacity(payload_bytes);
        for ((_, columns, _), column_spans) in relations.iter().zip(&spans) {
            for (column, span) in columns.iter().zip(column_spans) {
                payload.resize(span.start, 0);
                payload.extend_from_slice(column);
            }
        }
        let payload_owner = if relations.iter().any(|(_, columns, _)| !columns.is_empty()) {
            let mut allocation = self.memory.alloc::<u8>(payload.len())?;
            if !payload.is_empty() {
                self.htod_sync_copy_into_tracked(&payload, &mut allocation)?;
            }
            Some(allocation)
        } else {
            None
        };
        let mut counts = self.memory.alloc::<u8>(count_bytes.len())?;
        self.htod_launch_metadata_sync_copy_into(&count_bytes, &mut counts)?;

        let mut buffers = Vec::with_capacity(relations.len());
        for (index, ((schema, _, rows), column_spans)) in relations.iter().zip(spans).enumerate() {
            let mut columns = Vec::<CudaColumn>::with_capacity(column_spans.len());
            for span in column_spans {
                let column = payload_owner
                    .as_ref()
                    .expect("non-nullary relation has a payload owner")
                    .try_owned_subslice::<u8>(span)
                    .ok_or_else(|| XlogError::Kernel("Host column view is invalid".into()))?;
                columns.push(column.into());
            }
            let start = index * std::mem::size_of::<u32>();
            let count = counts
                .try_owned_subslice::<u32>(start..start + std::mem::size_of::<u32>())
                .ok_or_else(|| XlogError::Kernel("Host row count view is invalid".into()))?;
            let mut buffer = CudaBuffer::from_columns_with_host_count(
                columns,
                *rows as u64,
                count,
                schema.clone(),
                *rows as u32,
            );
            if canonical_sets[index] {
                buffer.certify_canonical_full_row_set();
            }
            buffers.push(buffer);
        }
        Ok(buffers)
    }

    /// Export CudaBuffer to Arrow C Data Interface (device-resident).
    ///
    /// This is a zero-copy export: column buffers remain on device, and the
    /// returned ArrowDeviceArray describes CUDA-resident memory.
    ///
    /// The export requires that the device row count matches the host row_cap
    pub fn to_arrow_device_record_batch(
        &self,
        buffer: CudaBuffer,
    ) -> Result<crate::arrow_device::ArrowDeviceArrayOwned> {
        use crate::device_runtime::Access;
        use crate::memory::{with_memory_access, DeviceRead};
        use arrow::array::ArrayData;
        use arrow::datatypes::{DataType, Field};
        use arrow::ffi::to_ffi;

        use crate::arrow_device::{ArrowDeviceArray, ARROW_DEVICE_CUDA};

        let stream = self.device.inner().stream();
        let _ordinary = crate::cuda_graph::reserve_uncaptured_stream(stream).map_err(|error| {
            XlogError::Kernel(format!("Arrow export stream admission: {error}"))
        })?;
        let buffer = Arc::new(buffer);
        let row_cap = buffer.num_rows();
        let num_rows_u32 = u32::try_from(row_cap).map_err(|_| {
            XlogError::Kernel(format!(
                "Arrow device export supports at most {} rows, got {}",
                u32::MAX,
                row_cap
            ))
        })?;

        // GPU-side assertion: device row count must match row_cap (no host reads).
        let assert_fn = self
            .device
            .inner()
            .get_func(D4_MODULE, d4_kernels::D4_ASSERT_U32_EQ)
            .ok_or_else(|| XlogError::Kernel("d4_assert_u32_eq kernel not found".to_string()))?;
        // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
        unsafe {
            assert_fn.clone().launch(
                LaunchConfig {
                    grid_dim: (1, 1, 1),
                    block_dim: (1, 1, 1),
                    shared_mem_bytes: 0,
                },
                (buffer.num_rows_device(), num_rows_u32),
            )
        }
        .map_err(|e| XlogError::Kernel(format!("d4_assert_u32_eq failed: {}", e)))?;
        // A null Arrow sync_event promises ready buffers, not merely a ready
        // row count. Admit every exported column so nonblocking producer-stream
        // dependencies are included in the wait without copying host data.
        let accesses = buffer
            .columns()
            .iter()
            .map(|column| column.device_view().access(Access::Read))
            .collect::<crate::device_runtime::ResourceResult<Vec<_>>>()
            .map_err(|error| XlogError::Kernel(error.to_string()))?;
        with_memory_access(Arc::clone(stream), accesses, |_| Ok(stream.synchronize()?))
            .map_err(|error| XlogError::Kernel(format!("Arrow producer readiness: {error}")))?;

        let num_rows = usize::try_from(num_rows_u32)
            .map_err(|_| XlogError::Kernel("Arrow device export row count overflow".to_string()))?;

        let mut fields: Vec<Field> = Vec::with_capacity(buffer.arity());
        let mut children: Vec<ArrayData> = Vec::with_capacity(buffer.arity());

        for (col_idx, (name, scalar_type)) in buffer.schema().columns.iter().enumerate() {
            let (field, child) =
                self.build_arrow_device_child(&buffer, col_idx, name, *scalar_type, num_rows)?;
            fields.push(field);
            children.push(child);
        }

        let struct_type = DataType::Struct(fields.into());
        let struct_data = ArrayData::builder(struct_type)
            .len(num_rows)
            .child_data(children)
            .build()
            .map_err(|e| XlogError::Kernel(format!("Arrow device export failed: {}", e)))?;

        let (ffi_array, ffi_schema) =
            to_ffi(&struct_data).map_err(|e| XlogError::Kernel(format!("{}", e)))?;
        let array_ptr = Box::into_raw(Box::new(ffi_array));
        let schema_ptr = Box::into_raw(Box::new(ffi_schema));

        // SAFETY: these uniquely transferred Box allocations retain all GPU
        // buffers through their thread-safe Arrow custom allocation owners.
        Ok(unsafe {
            ArrowDeviceArray::new(
                ARROW_DEVICE_CUDA,
                self.device.ordinal() as i32,
                array_ptr,
                schema_ptr,
            )
        })
    }

    /// Import Arrow C Data Interface (device-resident) into a CudaBuffer (zero-copy).
    ///
    /// This is an experimental API gated behind the `arrow-device-import` feature.
    #[cfg(feature = "arrow-device-import")]
    pub fn from_arrow_device_record_batch(
        &self,
        device_array: crate::arrow_device::ArrowDeviceArrayOwned,
    ) -> Result<CudaBuffer> {
        use arrow::array::ArrayData;
        use arrow::datatypes::DataType;
        use arrow::ffi::from_ffi;
        use std::sync::Arc;

        use crate::arrow_device::{ArrowDeviceImport, ARROW_DEVICE_CUDA};
        use crate::memory::CudaColumn;

        let (device_type, device_id, outer_owner, ffi_array, ffi_schema) =
            // SAFETY: device_array is a valid ArrowDeviceArrayOwned; into_ffi_parts transfers ownership per Arrow C Data Interface
            unsafe { device_array.into_ffi_parts() };

        if device_type != ARROW_DEVICE_CUDA {
            return Err(XlogError::Kernel(format!(
                "Arrow device import requires CUDA device type={}, got {}",
                ARROW_DEVICE_CUDA, device_type
            )));
        }
        if device_id != self.device.ordinal() as i32 {
            return Err(XlogError::Kernel(format!(
                "Arrow device import device id mismatch: expected {}, got {}",
                self.device.ordinal(),
                device_id
            )));
        }

        // SAFETY: ffi_array and ffi_schema conform to the Arrow C Data Interface; device and type checks passed above
        let data: ArrayData = unsafe { from_ffi(ffi_array, &ffi_schema) }
            .map_err(|e| XlogError::Kernel(format!("Arrow device import failed: {}", e)))?;
        drop(ffi_schema);
        let keepalive = Arc::new(ArrowDeviceImport::from_ffi(data, outer_owner));
        let data = keepalive.data();

        let (fields, children) = match data.data_type() {
            DataType::Struct(fields) => (fields.clone(), data.child_data().to_vec()),
            other => {
                return Err(XlogError::Kernel(format!(
                    "Arrow device import expects Struct, got {:?}",
                    other
                )))
            }
        };

        if data.offset() != 0 {
            return Err(XlogError::Kernel(
                "Arrow device import does not support non-zero offsets".to_string(),
            ));
        }
        if data.null_count() > 0 || data.nulls().is_some() {
            return Err(XlogError::Kernel(
                "Arrow device import does not support nulls".to_string(),
            ));
        }

        let num_rows = data.len();
        if fields.len() != children.len() {
            return Err(XlogError::Kernel(
                "Arrow device import field/child length mismatch".to_string(),
            ));
        }

        let mut columns = Vec::with_capacity(children.len());
        let mut schema_cols = Vec::with_capacity(children.len());

        for (field, child) in fields.iter().zip(children.iter()) {
            if child.len() != num_rows {
                return Err(XlogError::Kernel(
                    "Arrow device import child length mismatch".to_string(),
                ));
            }
            if child.offset() != 0 {
                return Err(XlogError::Kernel(
                    "Arrow device import does not support child offsets".to_string(),
                ));
            }
            if child.null_count() > 0 || child.nulls().is_some() {
                return Err(XlogError::Kernel(
                    "Arrow device import does not support child nulls".to_string(),
                ));
            }

            let (scalar_type, elem_size) = Self::scalar_type_from_arrow_field(field)?;
            let buffers = child.buffers();
            let buf = buffers.first().ok_or_else(|| {
                XlogError::Kernel("Arrow device import missing value buffer".to_string())
            })?;
            let len_bytes = buf.len();
            let expected_bytes = num_rows.checked_mul(elem_size).ok_or_else(|| {
                XlogError::Kernel("Arrow device import size overflow".to_string())
            })?;
            if len_bytes != expected_bytes {
                return Err(XlogError::Kernel(format!(
                    "Arrow device import buffer size mismatch: expected {}, got {}",
                    expected_bytes, len_bytes
                )));
            }

            let ptr = buf.as_ptr();
            if ptr.is_null() && len_bytes > 0 {
                return Err(XlogError::Kernel(
                    "Arrow device import got null buffer pointer".to_string(),
                ));
            }
            let device_ptr = ptr as usize as cudarc::driver::sys::CUdeviceptr;
            // SAFETY: the validated FFI child is retained with its full outer
            // producer owner; from_raw's contract supplies producer readiness.
            columns.push(unsafe {
                CudaColumn::arrow_device(
                    device_ptr,
                    len_bytes,
                    self.device().inner().stream().clone(),
                    keepalive.clone(),
                )
            });
            schema_cols.push((field.name().to_string(), scalar_type));
        }

        let schema = Schema::new(schema_cols);
        self.buffer_from_columns(columns, num_rows as u64, schema)
    }

    /// Export CudaBuffer to Arrow RecordBatch
    ///
    /// Downloads data from GPU and converts it to an Arrow RecordBatch for
    /// interoperability with Arrow-based tools like cuDF, Polars, or DuckDB.
    ///
    /// # Arguments
    /// * `buffer` - The CudaBuffer to export
    ///
    /// # Returns
    /// An Arrow RecordBatch containing all columns from the buffer
    ///
    /// # Errors
    /// Returns `XlogError::Kernel` if:
    /// - Column download fails
    /// - RecordBatch creation fails
    pub fn to_arrow_record_batch(
        &self,
        buffer: &CudaBuffer,
    ) -> Result<arrow::record_batch::RecordBatch> {
        use arrow::array::*;
        use arrow::datatypes::{Field, Schema as ArrowSchema};
        use arrow::record_batch::{RecordBatch, RecordBatchOptions};

        let num_rows = self.device_row_count(buffer)?;

        let fields: Vec<Field> = buffer
            .schema
            .columns
            .iter()
            .map(|(name, scalar_type)| Field::new(name, scalar_type.to_arrow_type(), false))
            .collect();
        let arrow_schema = Arc::new(ArrowSchema::new(fields));

        let mut arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(buffer.arity());

        for (col_idx, (_, scalar_type)) in buffer.schema.columns.iter().enumerate() {
            let col = buffer
                .column(col_idx)
                .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;

            // Handle empty buffer case
            if num_rows == 0 {
                let array: Arc<dyn Array> = match scalar_type {
                    ScalarType::Bool => Arc::new(BooleanArray::from(Vec::<bool>::new())),
                    ScalarType::U32 => Arc::new(UInt32Array::from(Vec::<u32>::new())),
                    ScalarType::Symbol => Arc::new(xlog_core::symbol::to_arrow(&[])),
                    ScalarType::I32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
                    ScalarType::U64 => Arc::new(UInt64Array::from(Vec::<u64>::new())),
                    ScalarType::I64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
                    ScalarType::F32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
                    ScalarType::F64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
                };
                arrays.push(array);
                continue;
            }

            let elem_size = scalar_type.size_bytes();
            let num_bytes = num_rows
                .checked_mul(elem_size)
                .ok_or_else(|| XlogError::Kernel("Row byte size overflow".to_string()))?;
            let mut bytes = vec![0u8; num_bytes];
            let col_view = self.column_bytes_view(col, num_bytes)?;
            self.device
                .inner()
                .dtoh_sync_copy_into(&col_view, &mut bytes)
                .map_err(|e| XlogError::Kernel(format!("Failed to download column: {}", e)))?;

            let array: Arc<dyn Array> = match scalar_type {
                ScalarType::Bool => Arc::new(BooleanArray::from(
                    bytes.iter().map(|&b| b != 0).collect::<Vec<_>>(),
                )),
                ScalarType::U32 => {
                    let values: Vec<u32> = bytes
                        .as_chunks::<4>()
                        .0
                        .iter()
                        .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                        .collect();
                    Arc::new(UInt32Array::from(values))
                }
                ScalarType::Symbol => {
                    let values: Vec<u32> = bytes
                        .as_chunks::<4>()
                        .0
                        .iter()
                        .map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                        .collect();
                    Arc::new(xlog_core::symbol::to_arrow(&values))
                }
                ScalarType::I32 => {
                    let values: Vec<i32> = bytes
                        .as_chunks::<4>()
                        .0
                        .iter()
                        .map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                        .collect();
                    Arc::new(Int32Array::from(values))
                }
                ScalarType::U64 => {
                    let values: Vec<u64> = bytes
                        .as_chunks::<8>()
                        .0
                        .iter()
                        .map(|c| {
                            u64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])
                        })
                        .collect();
                    Arc::new(UInt64Array::from(values))
                }
                ScalarType::I64 => {
                    let values: Vec<i64> = bytes
                        .as_chunks::<8>()
                        .0
                        .iter()
                        .map(|c| {
                            i64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])
                        })
                        .collect();
                    Arc::new(Int64Array::from(values))
                }
                ScalarType::F32 => {
                    let values: Vec<f32> = bytes
                        .as_chunks::<4>()
                        .0
                        .iter()
                        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
                        .collect();
                    Arc::new(Float32Array::from(values))
                }
                ScalarType::F64 => {
                    let values: Vec<f64> = bytes
                        .as_chunks::<8>()
                        .0
                        .iter()
                        .map(|c| {
                            f64::from_le_bytes([c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]])
                        })
                        .collect();
                    Arc::new(Float64Array::from(values))
                }
            };

            arrays.push(array);
        }

        let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
        RecordBatch::try_new_with_options(arrow_schema, arrays, &options)
            .map_err(|e| XlogError::Kernel(format!("Failed to create RecordBatch: {}", e)))
    }

    fn build_arrow_device_child(
        &self,
        buffer: &Arc<CudaBuffer>,
        col_idx: usize,
        name: &str,
        scalar_type: ScalarType,
        num_rows: usize,
    ) -> Result<(arrow::datatypes::Field, arrow::array::ArrayData)> {
        use arrow::array::ArrayData;
        use arrow::buffer::Buffer;
        use arrow::datatypes::{DataType, Field};
        use std::collections::HashMap;
        use std::ptr::NonNull;

        use crate::arrow_device::ArrowCudaAllocation;

        let col = buffer
            .column(col_idx)
            .ok_or_else(|| XlogError::Kernel(format!("Column {} not found", col_idx)))?;

        let (dtype, metadata) = if scalar_type == ScalarType::Symbol {
            let mut meta = HashMap::new();
            meta.insert("xlog.symbol".to_string(), "true".to_string());
            meta.insert("xlog.symbol_encoding".to_string(), "u32".to_string());
            (DataType::UInt32, Some(meta))
        } else {
            (scalar_type.to_arrow_type(), None)
        };

        let field = match metadata {
            Some(meta) => Field::new(name, dtype.clone(), false).with_metadata(meta),
            None => Field::new(name, dtype.clone(), false),
        };

        let elem_size = match scalar_type {
            ScalarType::Symbol => 4usize,
            _ => scalar_type.size_bytes(),
        };

        let len_bytes = num_rows
            .checked_mul(elem_size)
            .ok_or_else(|| XlogError::Kernel("Arrow device export size overflow".to_string()))?;

        let mut extra = Vec::new();
        let (ptr, len) = match scalar_type {
            ScalarType::Bool => {
                let packed_len = num_rows.div_ceil(8);
                let mut packed = self.memory.alloc::<u8>(packed_len)?;
                let pack_fn = self
                    .device
                    .inner()
                    .get_func(PACK_MODULE, pack_kernels::PACK_BOOLS_TO_BITMAP)
                    .ok_or_else(|| {
                        XlogError::Kernel("pack_bools_to_bitmap kernel not found".to_string())
                    })?;
                let block_size = 256u32;
                let grid_size = (packed_len as u32).div_ceil(block_size);
                // SAFETY: kernel arguments match the PTX signature; device buffers were allocated with sufficient size
                unsafe {
                    pack_fn.clone().launch(
                        LaunchConfig {
                            grid_dim: (grid_size, 1, 1),
                            block_dim: (block_size, 1, 1),
                            shared_mem_bytes: 0,
                        },
                        (col, num_rows as u32, &mut packed),
                    )
                }
                .map_err(|e| XlogError::Kernel(format!("pack_bools_to_bitmap failed: {}", e)))?;
                self.device.synchronize()?;
                let ptr = *packed.device_ptr() as usize as *mut u8;
                extra.push(packed);
                (ptr, packed_len)
            }
            _ => {
                let ptr = *col.device_ptr() as usize as *mut u8;
                (ptr, len_bytes)
            }
        };

        let alloc = Arc::new(ArrowCudaAllocation::new(Arc::clone(buffer), extra));
        let nn = if len == 0 {
            NonNull::dangling()
        } else {
            NonNull::new(ptr).ok_or_else(|| {
                XlogError::Kernel("Arrow device export got null device pointer".to_string())
            })?
        };
        // SAFETY: nn points to CUDA device memory managed by ArrowCudaAllocation; len is the correct byte length
        let buf = unsafe { Buffer::from_custom_allocation(nn, len, alloc) };

        let data = ArrayData::builder(dtype)
            .len(num_rows)
            .add_buffer(buf)
            .build()
            .map_err(|e| XlogError::Kernel(format!("Arrow device export failed: {}", e)))?;

        Ok((field, data))
    }

    #[cfg(feature = "arrow-device-import")]
    fn scalar_type_from_arrow_field(
        field: &arrow::datatypes::Field,
    ) -> Result<(ScalarType, usize)> {
        use arrow::datatypes::DataType;

        // Arrow's Field metadata is a map (possibly empty), not an Option.
        let is_symbol = field
            .metadata()
            .get("xlog.symbol")
            .map(|v| v == "true")
            .unwrap_or(false);

        let scalar = match field.data_type() {
            DataType::Boolean => {
                return Err(XlogError::Kernel(
                    "Arrow device import does not support bit-packed bool yet".to_string(),
                ))
            }
            DataType::UInt32 if is_symbol => ScalarType::Symbol,
            dt => ScalarType::from_arrow_type(dt).ok_or_else(|| {
                XlogError::Kernel(format!("Arrow device import unsupported type {:?}", dt))
            })?,
        };

        let elem_size = match scalar {
            ScalarType::Symbol => 4usize,
            _ => scalar.size_bytes(),
        };

        Ok((scalar, elem_size))
    }

    /// Import Arrow RecordBatch to CudaBuffer
    ///
    /// Uploads Arrow data to GPU memory.
    ///
    /// # Arguments
    /// * `record_batch` - The Arrow RecordBatch to import
    ///
    /// # Returns
    /// A new CudaBuffer with the data on GPU
    ///
    /// # Errors
    /// Returns error if Arrow type is not supported or upload fails
    pub fn from_arrow_record_batch(
        &self,
        record_batch: &arrow::record_batch::RecordBatch,
    ) -> Result<CudaBuffer> {
        use arrow::array::*;
        use arrow::datatypes::DataType;

        let num_rows = record_batch.num_rows() as u64;

        if num_rows == 0 {
            let columns: Vec<(String, ScalarType)> = record_batch
                .schema()
                .fields()
                .iter()
                .filter_map(|f| {
                    ScalarType::from_arrow_type(f.data_type()).map(|st| (f.name().clone(), st))
                })
                .collect();
            return self.create_empty_buffer(Schema::new(columns));
        }

        let mut columns = Vec::with_capacity(record_batch.num_columns());
        let mut schema_cols = Vec::with_capacity(record_batch.num_columns());

        for (col_idx, field) in record_batch.schema().fields().iter().enumerate() {
            let array = record_batch.column(col_idx);
            let scalar_type = ScalarType::from_arrow_type(field.data_type()).ok_or_else(|| {
                XlogError::Kernel(format!("Unsupported Arrow type: {:?}", field.data_type()))
            })?;

            let bytes: Vec<u8> = match field.data_type() {
                DataType::Boolean => {
                    let arr = array.as_any().downcast_ref::<BooleanArray>().unwrap();
                    arr.iter()
                        .map(|v| if v.unwrap_or(false) { 1u8 } else { 0u8 })
                        .collect()
                }
                DataType::UInt32 => {
                    let arr = array.as_any().downcast_ref::<UInt32Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                DataType::Int32 => {
                    let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                DataType::UInt64 => {
                    let arr = array.as_any().downcast_ref::<UInt64Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                DataType::Int64 => {
                    let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                DataType::Float32 => {
                    let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                DataType::Float64 => {
                    let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
                    arr.values().iter().flat_map(|v| v.to_le_bytes()).collect()
                }
                _ => {
                    return Err(XlogError::Kernel(format!(
                        "Unsupported Arrow type: {:?}",
                        field.data_type()
                    )))
                }
            };

            let mut d_col = self.memory.alloc::<u8>(bytes.len())?;
            self.htod_sync_copy_into_tracked(&bytes, &mut d_col)
                .map_err(|e| XlogError::Kernel(format!("Failed to upload column: {}", e)))?;

            columns.push(d_col.into());
            schema_cols.push((field.name().clone(), scalar_type));
        }

        self.buffer_from_columns(columns, num_rows, Schema::new(schema_cols))
    }

    /// Export a CudaBuffer to an Arrow IPC stream (RecordBatchStream) as bytes.
    ///
    /// This is a convenience wrapper around `to_arrow_record_batch` that enables
    /// interoperability with tools like cuDF via standard Arrow IPC readers.
    ///
    /// Note: This is not zero-copy; data is downloaded from GPU to host memory.
    pub fn to_arrow_ipc_stream(&self, buffer: &CudaBuffer) -> Result<Vec<u8>> {
        use arrow::ipc::writer::StreamWriter;

        let batch = self.to_arrow_record_batch(buffer)?;
        let mut out = Vec::new();
        let mut writer = StreamWriter::try_new(&mut out, &batch.schema())
            .map_err(|e| XlogError::Kernel(format!("Failed to create Arrow IPC writer: {}", e)))?;
        writer
            .write(&batch)
            .map_err(|e| XlogError::Kernel(format!("Failed to write Arrow RecordBatch: {}", e)))?;
        writer
            .finish()
            .map_err(|e| XlogError::Kernel(format!("Failed to finish Arrow IPC stream: {}", e)))?;
        Ok(out)
    }

    /// Import a single-batch Arrow IPC stream (RecordBatchStream) into a CudaBuffer.
    ///
    /// Note: This uploads Arrow data from host to GPU memory.
    pub fn from_arrow_ipc_stream(&self, ipc: &[u8]) -> Result<CudaBuffer> {
        use arrow::ipc::reader::StreamReader;
        use std::io::Cursor;

        let cursor = Cursor::new(ipc);
        let mut reader = StreamReader::try_new(cursor, None)
            .map_err(|e| XlogError::Kernel(format!("Failed to create Arrow IPC reader: {}", e)))?;

        let batches = reader.by_ref();
        let first = batches
            .next()
            .ok_or_else(|| {
                XlogError::Kernel("Arrow IPC stream contained no record batches".to_string())
            })?
            .map_err(|e| XlogError::Kernel(format!("Failed to read Arrow RecordBatch: {}", e)))?;

        if batches.next().is_some() {
            return Err(XlogError::Kernel(
                "Arrow IPC stream contains multiple record batches; this API expects exactly one"
                    .to_string(),
            ));
        }

        self.from_arrow_record_batch(&first)
    }

    /// Write a CudaBuffer to a file as an Arrow IPC stream (RecordBatchStream).
    pub fn write_arrow_ipc_stream_file<P: AsRef<std::path::Path>>(
        &self,
        buffer: &CudaBuffer,
        path: P,
    ) -> Result<()> {
        let bytes = self.to_arrow_ipc_stream(buffer)?;
        std::fs::write(&path, bytes).map_err(|e| {
            XlogError::Kernel(format!(
                "Failed to write Arrow IPC stream to {}: {}",
                path.as_ref().display(),
                e
            ))
        })?;
        Ok(())
    }

    /// Read a CudaBuffer from a file containing an Arrow IPC stream (RecordBatchStream).
    pub fn read_arrow_ipc_stream_file<P: AsRef<std::path::Path>>(
        &self,
        path: P,
    ) -> Result<CudaBuffer> {
        let bytes = std::fs::read(&path).map_err(|e| {
            XlogError::Kernel(format!(
                "Failed to read Arrow IPC stream from {}: {}",
                path.as_ref().display(),
                e
            ))
        })?;
        self.from_arrow_ipc_stream(&bytes)
    }
}