use std::{
error::Error as StdError,
ffi::{CStr, CString, c_char},
fs,
path::Path,
sync::Arc,
};
use crate::{
core::{
error::{self, Error, ErrorCode, FACT_COLUMN_INDEX, Result},
table::{Table, TableView},
},
runtime::{Runtime, execution_resolver::ResolvedExecution},
};
pub mod parquet;
pub mod source;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum KvikioCompatMode {
Auto,
On,
Off,
}
impl KvikioCompatMode {
const fn as_ffi(self) -> i32 {
match self {
Self::Auto => cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_AUTO as i32,
Self::On => cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_ON as i32,
Self::Off => cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_OFF as i32,
}
}
fn from_ffi(value: i32) -> Option<Self> {
match value {
value if value == cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_AUTO as i32 => Some(Self::Auto),
value if value == cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_ON as i32 => Some(Self::On),
value if value == cudf_sys::CUDF_IO_KVIKIO_COMPAT_MODE_OFF as i32 => Some(Self::Off),
_ => None,
}
}
pub const fn as_str(self) -> &'static str {
match self {
Self::Auto => "auto",
Self::On => "on",
Self::Off => "off",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct KvikioProcessConfig {
pub compat_mode: KvikioCompatMode,
pub thread_count: u32,
pub task_size_bytes: u64,
}
impl KvikioProcessConfig {
pub const DEFAULT: Self =
Self { compat_mode: KvikioCompatMode::On, thread_count: 64, task_size_bytes: 4 * 1024 * 1024 };
fn from_ffi(raw: &cudf_sys::cudf_io_kvikio_process_config_t) -> Option<Self> {
Some(Self {
compat_mode: KvikioCompatMode::from_ffi(raw.compat_mode)?,
thread_count: raw.thread_count,
task_size_bytes: raw.task_size_bytes,
})
}
const fn as_ffi(&self) -> cudf_sys::cudf_io_kvikio_process_config_t {
cudf_sys::cudf_io_kvikio_process_config_t {
compat_mode: self.compat_mode.as_ffi(),
thread_count: self.thread_count,
task_size_bytes: self.task_size_bytes,
}
}
}
const fn zeroed_kvikio_process_config() -> cudf_sys::cudf_io_kvikio_process_config_t {
cudf_sys::cudf_io_kvikio_process_config_t { compat_mode: 0, thread_count: 0, task_size_bytes: 0 }
}
pub fn configure_kvikio_process(device_ordinal: usize, defaults: &KvikioProcessConfig) -> Result<KvikioProcessConfig> {
const OPERATION: &str = "io::configure_kvikio_process";
let device_ordinal = i32::try_from(device_ordinal).map_err(|_| {
Error::new_with(ErrorCode::ForeignInvalidArgument, OPERATION, "device ordinal exceeds the CUDA device range")
.with_detail("device_ordinal", device_ordinal.to_string())
})?;
let raw_defaults = defaults.as_ffi();
let mut raw_effective = zeroed_kvikio_process_config();
let mut active_io = 0u64;
error::check_status_in_with_details(
unsafe {
cudf_sys::cudf_io_kvikio_configure_process(
device_ordinal,
&raw_defaults,
&mut raw_effective,
&mut active_io,
)
},
OPERATION,
[("resource", "kvikio_defaults".to_string())],
)?;
kvikio_configure_result(&raw_effective, active_io)
}
fn kvikio_configure_result(
raw_effective: &cudf_sys::cudf_io_kvikio_process_config_t,
active_io: u64,
) -> Result<KvikioProcessConfig> {
const OPERATION: &str = "io::configure_kvikio_process";
if active_io != 0 {
return Err(kvikio_busy_error(
OPERATION,
"cannot configure KvikIO while I/O is active",
"kvikio_defaults",
active_io,
"stop and join all I/O workers before configuring KvikIO",
));
}
KvikioProcessConfig::from_ffi(raw_effective).ok_or_else(|| {
Error::new_with(
ErrorCode::ForeignInvariantViolation,
OPERATION,
"KvikIO reported an unknown compatibility mode",
)
.with_detail("compat_mode", raw_effective.compat_mode.to_string())
})
}
pub fn kvikio_process_config() -> Option<KvikioProcessConfig> {
let mut raw = zeroed_kvikio_process_config();
let status = unsafe { cudf_sys::cudf_io_kvikio_process_config(&mut raw) };
(status == cudf_sys::CUDF_SUCCESS).then(|| KvikioProcessConfig::from_ffi(&raw)).flatten()
}
fn kvikio_busy_error(
operation: &'static str,
message: &'static str,
resource: &'static str,
active_io: u64,
remedy: &'static str,
) -> Error {
Error::new_with(ErrorCode::IoBusy, operation, message)
.with_detail("resource", resource)
.with_detail("reason", "active_io")
.with_detail("active_io", active_io.to_string())
.with_detail("remedy", remedy)
}
pub fn drain_kvikio_bounce_buffer_pools() -> Result<u64> {
const OPERATION: &str = "io::drain_kvikio_bounce_buffer_pools";
let mut cleared_bytes = 0u64;
let mut active_io = 0u64;
error::check_status_in_with_details(
unsafe { cudf_sys::cudf_io_kvikio_drain_bounce_buffer_pools(&mut cleared_bytes, &mut active_io) },
OPERATION,
[("resource", "kvikio_bounce_buffer_pools".to_string())],
)?;
kvikio_drain_result(cleared_bytes, active_io)
}
fn kvikio_drain_result(cleared_bytes: u64, active_io: u64) -> Result<u64> {
if active_io == 0 {
return Ok(cleared_bytes);
}
Err(kvikio_busy_error(
"io::drain_kvikio_bounce_buffer_pools",
"cannot drain KvikIO bounce buffers while I/O is active",
"kvikio_bounce_buffer_pools",
active_io,
"stop and join all I/O workers before retrying bounce-buffer drain",
))
}
pub fn gds_driver_available() -> Result<()> {
const OPERATION: &str = "io::gds_driver_available";
error::check_status_in_with_details(
unsafe { cudf_sys::cudf_gds_driver_available() },
OPERATION,
[("resource", "cufile_driver".to_string())],
)
}
pub struct TableReadOutput {
pub table: Table,
pub column_names: Option<Vec<String>>,
}
fn resolve_write_execution(
table: &TableView<'_>,
runtime: Option<&Runtime>,
operation: &'static str,
) -> Result<ResolvedExecution> {
let write_domain = runtime
.map(|runtime| Arc::clone(runtime.execution_domain()))
.unwrap_or_else(|| Arc::clone(table.execution_domain()));
ResolvedExecution::with_output_domain(write_domain, [table.execution_domain()], operation)
}
fn map_path_io_error(
source: impl StdError + Send + Sync + 'static,
path: &Path,
operation: &'static str,
message: &'static str,
) -> Error {
Error::new_with(ErrorCode::IoFailure, operation, message)
.with_detail("path", path.display().to_string())
.with_source(source)
}
fn validate_input_path(path: &Path, operation: &'static str) -> Result<()> {
if path.to_string_lossy().contains("://") {
return Ok(());
}
let metadata = fs::metadata(path)
.map_err(|source| map_path_io_error(source, path, operation, "failed to access input path"))?;
if metadata.is_file() {
return Ok(());
}
if metadata.is_dir() {
return Err(Error::new_with(ErrorCode::PathIsDirectory, operation, "input path is a directory")
.with_detail("path", path.display().to_string())
.with_detail("resource", "input_path")
.with_detail("reason", "path_is_directory")
.with_detail("remedy", "provide a regular file path instead of a directory"));
}
Err(Error::new_with(ErrorCode::IoFailure, operation, "input path is not a regular file")
.with_detail("path", path.display().to_string())
.with_detail("resource", "input_path")
.with_detail("reason", "path_not_regular_file")
.with_detail("remedy", "provide a regular file path"))
}
fn path_to_cstring(path: &Path, operation: &'static str) -> Result<CString> {
let path_str = path.to_str().ok_or_else(|| {
Error::new_with(ErrorCode::IoFailure, operation, "path is not valid UTF-8")
.with_detail("resource", "path")
.with_detail("reason", "invalid_utf8")
})?;
CString::new(path_str).map_err(|_| {
error::interior_null_byte_error(operation, "path", "path contains interior null byte")
.with_detail("path", path.to_string_lossy().into_owned())
})
}
unsafe fn collect_column_names(
names_ptr: *mut *mut c_char,
num_names: i32,
operation: &'static str,
) -> Result<Vec<String>> {
if num_names < 0 {
return Err(error::foreign_invariant_violation(
operation,
"ffi returned negative column name count",
"column_name_count",
"negative_name_count",
"ensure the foreign reader reports a non-negative column name count",
)
.with_detail("column_name_count", num_names.to_string()));
}
if num_names > 0 && names_ptr.is_null() {
return Err(error::foreign_invariant_violation(
operation,
"ffi returned null column name array with positive count",
"column_names",
"null_name_array",
"ensure the foreign reader returns a name array whenever column_name_count is positive",
)
.with_detail("column_name_count", num_names.to_string()));
}
let mut names = Vec::with_capacity(num_names as usize);
for i in 0..num_names {
let name_ptr = unsafe { *names_ptr.add(i as usize) };
if name_ptr.is_null() {
return Err(error::foreign_invariant_violation(
operation,
"ffi returned null column name entry",
"column_name",
"null_name_entry",
"ensure the foreign reader returns one non-null name entry per reported column",
)
.with_detail("column_name_count", num_names.to_string())
.with_fact(FACT_COLUMN_INDEX, i as usize));
}
let cstr = unsafe { CStr::from_ptr(name_ptr) };
let name = cstr.to_str().map_err(|_| {
error::foreign_invariant_violation(
operation,
"column name from FFI is not valid UTF-8",
"column_name",
"invalid_utf8",
"ensure the foreign reader returns valid UTF-8 column names",
)
.with_detail("column_name_count", num_names.to_string())
.with_fact(FACT_COLUMN_INDEX, i as usize)
})?;
names.push(name.to_owned());
}
Ok(names)
}
#[cfg(test)]
mod tests {
use std::{ffi::CString, fs, os::raw::c_char, ptr};
use crate::core::error::{ErrorCode, ErrorKind, ErrorStatus, FactKeyName, FactValue};
#[test]
fn kvikio_bounce_buffer_drain_busy_error_is_flat_and_retryable() {
let error = super::kvikio_drain_result(0, 2).unwrap_err();
assert_eq!(error.code(), ErrorCode::IoBusy);
assert_eq!(error.kind(), ErrorKind::ResourceExhausted);
assert_eq!(error.status(), ErrorStatus::Temporary);
assert_eq!(error.detail_value("active_io"), Some("2"));
assert_eq!(error.detail_value("reason"), Some("active_io"));
}
#[test]
fn kvikio_bounce_buffer_drain_idle_result_is_idempotent() {
assert_eq!(super::kvikio_drain_result(16 * 1024 * 1024, 0).unwrap(), 16 * 1024 * 1024);
assert_eq!(super::kvikio_drain_result(0, 0).unwrap(), 0);
}
#[test]
fn test_collect_column_names_rejects_negative_count() {
let err = unsafe { super::collect_column_names(ptr::null_mut(), -1, "test") }.unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvariantViolation);
assert_eq!(err.code(), ErrorCode::ForeignInvariantViolation);
assert_eq!(err.operation(), "test");
assert!(err.details().iter().any(|(k, v)| *k == "column_name_count" && v == "-1"));
assert!(!err.details().iter().any(|(k, _)| *k == "num_names"));
assert!(err.details().iter().any(|(k, v)| *k == "resource" && v == "column_name_count"));
assert!(err.details().iter().any(|(k, v)| *k == "reason" && v == "negative_name_count"));
assert!(
err.details()
.iter()
.any(|(k, v)| *k == "remedy"
&& v == "ensure the foreign reader reports a non-negative column name count")
);
}
#[test]
fn test_collect_column_names_rejects_null_names_ptr_with_positive_count() {
let err = unsafe { super::collect_column_names(ptr::null_mut(), 1, "test") }.unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvariantViolation);
assert_eq!(err.code(), ErrorCode::ForeignInvariantViolation);
assert_eq!(err.operation(), "test");
assert!(err.details().iter().any(|(k, v)| *k == "column_name_count" && v == "1"));
assert!(err.details().iter().any(|(k, v)| *k == "resource" && v == "column_names"));
assert!(err.details().iter().any(|(k, v)| *k == "reason" && v == "null_name_array"));
assert!(err.details().iter().any(|(k, v)| *k == "remedy"
&& v == "ensure the foreign reader returns a name array whenever column_name_count is positive"));
assert!(!err.details().iter().any(|(k, _)| *k == "names_ptr_null"));
}
#[test]
fn test_collect_column_names_rejects_null_entry_pointer() {
let mut ptrs: Vec<*mut c_char> = vec![ptr::null_mut()];
let err = unsafe { super::collect_column_names(ptrs.as_mut_ptr(), 1, "test") }.unwrap_err();
assert_eq!(err.kind(), ErrorKind::InvariantViolation);
assert_eq!(err.code(), ErrorCode::ForeignInvariantViolation);
assert_eq!(err.operation(), "test");
assert_eq!(err.fact_value(FactKeyName::ColumnIndex), Some(&FactValue::Usize(0)));
assert!(!err.details().iter().any(|(k, _)| *k == "column_index"));
assert!(err.details().iter().any(|(k, v)| *k == "resource" && v == "column_name"));
assert!(err.details().iter().any(|(k, v)| *k == "reason" && v == "null_name_entry"));
assert!(err.details().iter().any(|(k, v)| *k == "remedy"
&& v == "ensure the foreign reader returns one non-null name entry per reported column"));
assert!(!err.details().iter().any(|(k, _)| *k == "name_index"));
}
#[test]
fn test_collect_column_names_success() {
let mut raw = [CString::new("alpha").unwrap(), CString::new("beta").unwrap()];
let mut ptrs: Vec<*mut c_char> = raw.iter_mut().map(|s| s.as_ptr() as *mut c_char).collect();
let names = unsafe { super::collect_column_names(ptrs.as_mut_ptr(), ptrs.len() as i32, "test") }.unwrap();
assert_eq!(names, vec!["alpha", "beta"]);
}
#[test]
fn test_validate_input_path_reports_io_failure() {
let path = std::env::temp_dir().join("cudf_rs_missing_input_path_test_12345.csv");
let err = super::validate_input_path(&path, "test").unwrap_err();
assert_eq!(err.code(), ErrorCode::IoFailure);
assert_eq!(err.kind(), ErrorKind::DependencyFailure);
assert_eq!(err.operation(), "test");
assert!(err.details().iter().any(|(k, v)| *k == "path" && v == &path.display().to_string()));
}
#[test]
fn test_validate_input_path_rejects_directory() {
let path = std::env::temp_dir().join(format!("cudf_rs_directory_input_path_test_{}", std::process::id()));
fs::create_dir_all(&path).unwrap();
let result = super::validate_input_path(&path, "test");
let _ = fs::remove_dir(&path);
let err = result.unwrap_err();
let path_display = path.display().to_string();
assert_eq!(err.code(), ErrorCode::PathIsDirectory);
assert_eq!(err.code().to_string(), "PathIsDirectory");
assert_eq!(err.kind(), ErrorKind::InvalidInput);
assert_eq!(err.status(), ErrorStatus::Permanent);
assert_eq!(err.operation(), "test");
assert_eq!(err.detail_value("path"), Some(path_display.as_str()));
assert_eq!(err.detail_value("resource"), Some("input_path"));
assert_eq!(err.detail_value("reason"), Some("path_is_directory"));
assert_eq!(err.detail_value("remedy"), Some("provide a regular file path instead of a directory"));
}
}