diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aac6be9..b69843d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -154,7 +154,7 @@ jobs: /o:"npiontko" \ /d:sonar.token="$SONAR_TOKEN" \ /d:sonar.cs.vscoveragexml.reportsPaths=coverage.xml \ - /d:sonar.coverage.exclusions="docs/**,examples/**,native/**"" + /d:sonar.coverage.exclusions="docs/**,examples/**,native/**" # Build the project without building the native library, it should be already built in the check job and cached dotnet build -c Release -p:BuildNativeLib=false diff --git a/DataFusionSharp.slnx b/DataFusionSharp.slnx index 779e11f..e4ba9b5 100644 --- a/DataFusionSharp.slnx +++ b/DataFusionSharp.slnx @@ -70,6 +70,7 @@ + diff --git a/README.md b/README.md index 364122c..bacdc69 100644 --- a/README.md +++ b/README.md @@ -11,8 +11,8 @@ ## Features -| Component | Feature | Status | Notes | -|------------------|----------------------------------------------|---------|---------------------------------------------------| +| Component | Feature | Status | Notes | +|------------------|----------------------------------------------|--------|---------------------------------------------------| | **Runtime** | Tokio runtime | ✅ | Configurable threads, supports multiple instances | | | Logger | ✅ | Configurable with log levels | | **Session** | Create session context | ✅ | | @@ -25,13 +25,13 @@ | | Amazon S3 | ✅ | | | | Azure Blob Storage | ✅ | | | | Google Cloud Storage | ✅ | | -| | InMemory | ❌ | | +| | InMemory | ✅ | | | **DataFrame** | Count rows | ✅ | `CountAsync()` | | | Get schema | ✅ | `GetSchema()` → Arrow Schema | | | Collect all data | ✅ | `CollectAsync()` → RecordBatches | | | Stream results | ✅ | `ExecuteStreamAsync()` → IAsyncEnumerable | | | Show/print | ✅ | `ShowAsync()`, `ToStringAsync()` | -| | Select, Aggregate, Join, Filter, Limit, Sort | 🟡 | Use SQL instead | +| | Select, Aggregate, Join, Filter, Limit, Sort | 🟡 | Use SQL instead | | **Arrow** | Apache Arrow support | ✅ | Via Apache.Arrow nuget package | | | Zero copy support | ✅ | | | **Advanced** | UDF registration | ❌ | | diff --git a/examples/QueryInMemoryData/Program.cs b/examples/QueryInMemoryData/Program.cs new file mode 100644 index 0000000..71c9820 --- /dev/null +++ b/examples/QueryInMemoryData/Program.cs @@ -0,0 +1,58 @@ +using Dapper; +using DataFusionSharp; +using DataFusionSharp.Data; +using Microsoft.Extensions.Logging; + +// Setup logging +var loggerFactory = LoggerFactory.Create(builder => +{ + builder + .SetMinimumLevel(LogLevel.Information) + .AddSimpleConsole(o => o.IncludeScopes = true); +}); +DataFusionNativeLogger.ConfigureLogger(loggerFactory.CreateLogger("DataFusionSharp"), LogLevel.Information); + +// Init runtime and single session +using var runtime = DataFusionRuntime.Create(); +using var session = runtime.CreateSessionContext(); + +// Load customers CSV into an in-memory object store +using var store = runtime.CreateInMemoryStore(); + +var customersCsvBytes = await File.ReadAllBytesAsync(Path.Combine("Data", "orders", "csv", "customers.csv")); +await store.PutAsync("customers.csv", customersCsvBytes); + +session.RegisterInMemoryObjectStore("memory://", store); +await session.RegisterCsvAsync("customers", "memory:///customers.csv"); + +// Register orders from a regular file on disk +await session.RegisterCsvAsync("orders", Path.Combine("Data", "orders", "csv", "orders.csv")); + +// Create a connection wrapper around the SessionContext for ADO.NET operations. +await using var connection = session.AsConnection(); + + +// QueryAsync – map rows to a strongly-typed record +Console.WriteLine("=== Completed orders per customer (in-memory JOIN file) ==="); + +var summaries = await connection.QueryAsync( + """ + SELECT + c.customer_name AS CustomerName, + c.country AS Country, + COUNT(*) AS OrderCount, + SUM(o.order_amount) AS TotalAmount + FROM customers AS c + JOIN orders AS o ON c.customer_id = o.customer_id + WHERE o.order_status = @status + GROUP BY c.customer_name, c.country + ORDER BY TotalAmount DESC + """, + new { status = "Completed" }); + +foreach (var s in summaries) + Console.WriteLine($" {s.CustomerName,-25} {s.Country,-10} orders: {s.OrderCount,3} total: {s.TotalAmount,10:N0}"); + +// Model types + +record OrderSummary(string CustomerName, string Country, long OrderCount, long TotalAmount); diff --git a/examples/QueryInMemoryData/QueryInMemoryData.csproj b/examples/QueryInMemoryData/QueryInMemoryData.csproj new file mode 100644 index 0000000..4003fec --- /dev/null +++ b/examples/QueryInMemoryData/QueryInMemoryData.csproj @@ -0,0 +1,28 @@ + + + + Exe + net10.0 + 14 + enable + enable + false + + + + + + + + + + + + + + + + + diff --git a/native/Cargo.toml b/native/Cargo.toml index 422b477..88b84e2 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -16,6 +16,7 @@ panic = "unwind" [dependencies] anyhow = "1.0.102" arrow-array = { version = "57.3.0", features = ["ffi"] } +bytes = "1.11.1" datafusion = "52.4.0" datafusion-proto = "52.4.0" futures = "0.3.32" diff --git a/native/src/common.rs b/native/src/common.rs index 783e096..f6a8fd5 100644 --- a/native/src/common.rs +++ b/native/src/common.rs @@ -20,16 +20,30 @@ impl BytesData { } pub(crate) fn as_slice(&self) -> &[u8] { - unsafe { std::slice::from_raw_parts( self.data, self.len as usize) } + unsafe { + std::slice::from_raw_parts( self.data, self.len as usize) + } + } + + pub(crate) fn as_slice_static(&self) -> &'static [u8] { + unsafe { + std::slice::from_raw_parts( self.data, self.len as usize) + } } pub(crate) fn as_opt_slice(&self) -> Option<&[u8]> { if self.data.is_null() { None } else { - Some(unsafe { std::slice::from_raw_parts( self.data, self.len as usize) }) + Some(unsafe { + std::slice::from_raw_parts( self.data, self.len as usize) + }) } } + + pub(crate) fn len(&self) -> usize { + self.len as usize + } } #[repr(C)] diff --git a/native/src/context.rs b/native/src/context.rs index 3038967..79dd36d 100644 --- a/native/src/context.rs +++ b/native/src/context.rs @@ -623,6 +623,39 @@ pub unsafe extern "C" fn datafusion_context_register_object_store_http( ErrorCode::Ok } +/// Registers an in-memory store. +/// +/// This is a synchronous operation. The callback is invoked with the result. +/// +/// # Safety +/// - `context_ptr` must be a valid pointer returned by `datafusion_context_new` +/// - `base_url_ptr` must be a valid null-terminated UTF-8 string (e.g. "memory://") +/// - `store_ptr` must be a valid pointer returned by `datafusion_in_memory_store_new` +/// - `callback` must be valid to call from any thread +#[unsafe(no_mangle)] +pub unsafe extern "C" fn datafusion_context_register_object_store_in_memory( + context_ptr: *mut SessionContextWrapper, + base_url_ptr: *const std::ffi::c_char, + store_ptr: *const crate::memory_store::InMemoryStoreWrapper, + callback: crate::Callback, + user_data: u64, +) -> ErrorCode { + let context = ffi_ref!(context_ptr); + + let base_url = ffi_cstr_to_string!(base_url_ptr); + let Ok(url) = url::Url::parse(&base_url) else { return ErrorCode::InvalidArgument }; + + let store = ffi_ref!(store_ptr); + + debug!("Registering in-memory object store {store_ptr:p} for '{url}' on session {context_ptr:p}"); + + context.inner.register_object_store(&url, store.inner()); + + crate::invoke_callback_null_result(callback, user_data); + + ErrorCode::Ok +} + /// Deregisters an object store for the given URL. /// /// This is a synchronous operation. The callback is invoked with the result. diff --git a/native/src/error.rs b/native/src/error.rs index 5b8c802..792f5c1 100644 --- a/native/src/error.rs +++ b/native/src/error.rs @@ -8,7 +8,8 @@ pub enum ErrorCode { RuntimeShutdownFailed = 4, TableRegistrationFailed = 5, SqlError = 6, - DataFrameError = 7 + DataFrameError = 7, + ObjectStoreError = 8 } #[derive(Debug, Clone)] diff --git a/native/src/lib.rs b/native/src/lib.rs index 47814f6..0080b7c 100644 --- a/native/src/lib.rs +++ b/native/src/lib.rs @@ -12,6 +12,7 @@ pub mod logger; pub mod runtime; pub mod context; pub mod dataframe; +pub mod memory_store; pub use proto::*; pub use error::*; diff --git a/native/src/memory_store.rs b/native/src/memory_store.rs new file mode 100644 index 0000000..c36be27 --- /dev/null +++ b/native/src/memory_store.rs @@ -0,0 +1,174 @@ +use std::sync::Arc; + +use log::{debug, error, warn}; + +use object_store::path::Path; +use object_store::{ObjectStore, PutPayload}; +use object_store::memory::InMemory; + +use crate::{BytesData, Callback, ErrorCode}; +use crate::error::ErrorInfo; + +pub struct InMemoryStoreWrapper { + runtime: crate::RuntimeHandle, + inner: Arc +} + +impl InMemoryStoreWrapper { + pub(crate) fn new(runtime: &crate::RuntimeHandle) -> Self { + Self { + runtime: Arc::clone(runtime), + inner: Arc::new(InMemory::new()) + } + } + + pub(crate) fn inner(&self) -> Arc { + Arc::clone(&self.inner) + } +} + +/// Creates a new in-memory store. +/// +/// # Safety +/// - `runtime_ptr` must be a valid pointer returned by `datafusion_runtime_new` +/// - `store_ptr` must be a valid pointer to a pointer that will receive the store pointer +/// - Caller must call `datafusion_in_memory_store_destroy` exactly once with the returned pointer +#[unsafe(no_mangle)] +pub unsafe extern "C" fn datafusion_in_memory_store_new( + runtime_ptr: *mut crate::RuntimeHandle, + store_ptr: *mut *mut InMemoryStoreWrapper +) -> ErrorCode { + if store_ptr.is_null() { + error!("Received null output pointer for store"); + return ErrorCode::InvalidArgument; + } + + let runtime_handle = ffi_ref!(runtime_ptr); + + let store = Box::new(InMemoryStoreWrapper::new(runtime_handle)); + let store_raw_ptr = Box::into_raw(store); + + unsafe { + *store_ptr = store_raw_ptr; + } + + debug!("Created in-memory store {store_raw_ptr:p}"); + + ErrorCode::Ok +} + +/// Destroys a `InMemoryStoreWrapper` created by `datafusion_in_memory_store_new`. +/// +/// # Safety +/// - `store_ptr` must be a valid pointer returned by `datafusion_in_memory_store_new`, or null +/// - Caller must not use `store_ptr` after this call +#[unsafe(no_mangle)] +pub unsafe extern "C" fn datafusion_in_memory_store_destroy(store_ptr: *mut InMemoryStoreWrapper) -> ErrorCode { + debug!("Destroying in-memory store {store_ptr:p}"); + + if store_ptr.is_null() { + warn!("Received null output pointer for store"); + } else { + unsafe { + drop(Box::from_raw(store_ptr)); + }; + } + + ErrorCode::Ok +} + +/// Puts data into the in-memory store at the specified path. +/// +/// # Safety +/// - `store_ptr` must be a valid pointer returned by `datafusion_in_memory_store_new` +/// - `path_ptr` must be a valid null-terminated C string representing the object path +/// - `data_bytes` must point to valid memory for the duration of this call if copy is `true` or for the duration of the store if copy is `false` +/// - If `copy` is false, the data behind `data_bytes` must remain valid for the lifetime of the store +/// - `callback` will be invoked exactly once when the operation completes +#[unsafe(no_mangle)] +pub unsafe extern "C" fn datafusion_in_memory_store_put( + store_ptr: *mut InMemoryStoreWrapper, + path_ptr: *const std::ffi::c_char, + data_bytes: BytesData, + copy: bool, + callback: Callback, + user_data: u64 +) -> ErrorCode { + let store_wrapper = ffi_ref!(store_ptr); + let path_str = ffi_cstr_to_string!(path_ptr); + + debug!("Putting data to in-memory store {store_ptr:p} at path '{path_str}' with data length {}, copy={copy}", data_bytes.len()); + + let store = Arc::clone(&store_wrapper.inner); + + let bytes = if copy { + bytes::Bytes::from(data_bytes.as_slice().to_owned()) + } else { + bytes::Bytes::from_static(data_bytes.as_slice_static()) + }; + + store_wrapper.runtime.spawn(async move { + let Some(path) = parse_path(&path_str, callback, user_data) else { return }; + + let payload = PutPayload::from_bytes(bytes); + + let result = store + .put(&path, payload) + .await + .map(drop) + .map_err(|e| ErrorInfo::new(ErrorCode::ObjectStoreError, e)); + + crate::invoke_callback(result, callback, user_data); + }); + + ErrorCode::Ok +} + +/// Deletes an object from the in-memory store at the specified path. +/// +/// # Safety +/// - `store_ptr` must be a valid pointer returned by `datafusion_in_memory_store_new` +/// - `path_ptr` must be a valid null-terminated C string representing the object path to delete +/// - `callback` will be invoked exactly once when the operation completes +#[unsafe(no_mangle)] +pub unsafe extern "C" fn datafusion_in_memory_store_delete( + store_ptr: *mut InMemoryStoreWrapper, + path_ptr: *const std::ffi::c_char, + callback: Callback, + user_data: u64 +) -> ErrorCode { + let store_wrapper = ffi_ref!(store_ptr); + let path_str = ffi_cstr_to_string!(path_ptr); + + debug!("Deleting data from in-memory store {store_ptr:p} at path '{path_str}'"); + + let store = Arc::clone(&store_wrapper.inner); + + store_wrapper.runtime.spawn(async move { + let Some(path) = parse_path(&path_str, callback, user_data) else { return }; + + let result = store + .delete(&path) + .await + .map_err(|e| ErrorInfo::new(ErrorCode::ObjectStoreError, e)); + + crate::invoke_callback(result, callback, user_data); + }); + + ErrorCode::Ok +} + +fn parse_path(path_str: &str, callback: Callback, user_data: u64) -> Option { + match Path::parse(path_str) { + Ok(path) => Some(path), + Err(e) => { + error!("Failed to parse path: {e}"); + crate::invoke_callback_error( + &ErrorInfo::new(ErrorCode::InvalidArgument, e), + callback, + user_data, + ); + None + } + } +} diff --git a/src/DataFusionSharp/DataFrame.cs b/src/DataFusionSharp/DataFrame.cs index 61e9a69..3239c63 100644 --- a/src/DataFusionSharp/DataFrame.cs +++ b/src/DataFusionSharp/DataFrame.cs @@ -47,7 +47,7 @@ public DataFrame WithParameters(IEnumerable paramet { ArgumentNullException.ThrowIfNull(parameters); - using var paramValuesData = PinnedProtobufData.FromMessage(parameters.ToProto()); + using var paramValuesData = PinnedBytesData.FromMessage(parameters.ToProto()); var id = SyncOperations.Instance.Create(); var result = NativeMethods.DataFrameWithParameters(_handle, paramValuesData.ToBytesData(), GenericCallbacks.CallbackForVoidSyncHandle, id); if (result != DataFusionErrorCode.Ok) @@ -186,8 +186,8 @@ public Task WriteCsvAsync(string path, DataFrameWriteOptions? dataFrameWriteOpti { ArgumentException.ThrowIfNullOrEmpty(path); - using var dataFrameOptionsData = PinnedProtobufData.FromMessage(dataFrameWriteOptions?.ToProto()); - using var csvOptionsData = PinnedProtobufData.FromMessage(csvWriteOptions?.ToProto()); + using var dataFrameOptionsData = PinnedBytesData.FromMessage(dataFrameWriteOptions?.ToProto()); + using var csvOptionsData = PinnedBytesData.FromMessage(csvWriteOptions?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.DataFrameWriteCsv(_handle, path, @@ -214,8 +214,8 @@ public Task WriteJsonAsync(string path, DataFrameWriteOptions? dataFrameWriteOpt { ArgumentException.ThrowIfNullOrEmpty(path); - using var dataFrameOptionsData = PinnedProtobufData.FromMessage(dataFrameWriteOptions?.ToProto()); - using var optionsData = PinnedProtobufData.FromMessage(jsonWriteOptions?.ToProto()); + using var dataFrameOptionsData = PinnedBytesData.FromMessage(dataFrameWriteOptions?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(jsonWriteOptions?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.DataFrameWriteJson(_handle, path, @@ -242,8 +242,8 @@ public Task WriteParquetAsync(string path, DataFrameWriteOptions? dataFrameWrite { ArgumentException.ThrowIfNullOrEmpty(path); - using var dataFrameOptionsData = PinnedProtobufData.FromMessage(dataFrameWriteOptions?.ToProto()); - using var parquetOptionsData = PinnedProtobufData.FromMessage(parquetWriteOptions?.ToProto()); + using var dataFrameOptionsData = PinnedBytesData.FromMessage(dataFrameWriteOptions?.ToProto()); + using var parquetOptionsData = PinnedBytesData.FromMessage(parquetWriteOptions?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.DataFrameWriteParquet(_handle, path, diff --git a/src/DataFusionSharp/DataFusionErrorCode.cs b/src/DataFusionSharp/DataFusionErrorCode.cs index 1cc8e37..b44221e 100644 --- a/src/DataFusionSharp/DataFusionErrorCode.cs +++ b/src/DataFusionSharp/DataFusionErrorCode.cs @@ -20,5 +20,7 @@ public enum DataFusionErrorCode /// An error occurred while executing SQL. SqlError = 6, /// An error occurred during DataFrame operations. - DataFrameError = 7 + DataFrameError = 7, + /// An error occurred in the object store. + ObjectStoreError = 8 } diff --git a/src/DataFusionSharp/DataFusionRuntime.cs b/src/DataFusionSharp/DataFusionRuntime.cs index 3562353..bda3310 100644 --- a/src/DataFusionSharp/DataFusionRuntime.cs +++ b/src/DataFusionSharp/DataFusionRuntime.cs @@ -1,4 +1,5 @@ using DataFusionSharp.Interop; +using DataFusionSharp.ObjectStore; namespace DataFusionSharp; @@ -53,6 +54,19 @@ public SessionContext CreateSessionContext() return new SessionContext(this, new SessionContextSafeHandle(contextHandle)); } + + /// + /// Creates a new in-memory object store. + /// + /// A new instance. + /// Thrown when object store creation fails. + public InMemoryObjectStore CreateInMemoryStore() + { + var errorCode = NativeMethods.InMemoryStoreNew(_handle, out var storeHandle); + DataFusionException.ThrowIfError(errorCode, "Failed to create in-memory object store"); + + return new InMemoryObjectStore(this, new InMemoryStoreSafeHandle(storeHandle)); + } /// /// Shuts down the runtime and releases all resources. diff --git a/src/DataFusionSharp/Interop/NativeMethods.cs b/src/DataFusionSharp/Interop/NativeMethods.cs index 64f736c..61beea4 100644 --- a/src/DataFusionSharp/Interop/NativeMethods.cs +++ b/src/DataFusionSharp/Interop/NativeMethods.cs @@ -83,6 +83,9 @@ internal enum NativeLogLevel : uint [LibraryImport(LibraryName, EntryPoint = "datafusion_context_register_object_store_http")] public static partial DataFusionErrorCode ContextRegisterObjectStoreHttp(SessionContextSafeHandle contextHandle, [MarshalAs(UnmanagedType.LPUTF8Str)] string url, BytesData optionsData, IntPtr callback, ulong userData); + [LibraryImport(LibraryName, EntryPoint = "datafusion_context_register_object_store_in_memory")] + public static partial DataFusionErrorCode ContextRegisterObjectStoreInMemory(SessionContextSafeHandle contextHandle, [MarshalAs(UnmanagedType.LPUTF8Str)] string url, InMemoryStoreSafeHandle inMemoryStoreHandle, IntPtr callback, ulong userData); + [LibraryImport(LibraryName, EntryPoint = "datafusion_context_deregister_object_store")] public static partial DataFusionErrorCode ContextDeregisterObjectStore(SessionContextSafeHandle contextHandle, [MarshalAs(UnmanagedType.LPUTF8Str)] string url, IntPtr callback, ulong userData); @@ -131,4 +134,18 @@ internal enum NativeLogLevel : uint [LibraryImport(LibraryName, EntryPoint = "datafusion_dataframe_stream_next")] public static partial DataFusionErrorCode DataFrameStreamNext(DataFrameStreamSafeHandle streamHandle, IntPtr callback, ulong userData); + + // InMemory Store + + [LibraryImport(LibraryName, EntryPoint = "datafusion_in_memory_store_new")] + public static partial DataFusionErrorCode InMemoryStoreNew(RuntimeSafeHandle runtimeHandle, out IntPtr inMemoryHandle); + + [LibraryImport(LibraryName, EntryPoint = "datafusion_in_memory_store_destroy")] + public static partial DataFusionErrorCode InMemoryStoreDestroy(IntPtr inMemoryHandle); + + [LibraryImport(LibraryName, EntryPoint = "datafusion_in_memory_store_put")] + public static partial DataFusionErrorCode InMemoryStorePut(InMemoryStoreSafeHandle inMemoryHandle, [MarshalAs(UnmanagedType.LPUTF8Str)] string path, BytesData data, [MarshalAs(UnmanagedType.I1)] bool copy, IntPtr callback, ulong userData); + + [LibraryImport(LibraryName, EntryPoint = "datafusion_in_memory_store_delete")] + public static partial DataFusionErrorCode InMemoryStoreDelete(InMemoryStoreSafeHandle inMemoryHandle, [MarshalAs(UnmanagedType.LPUTF8Str)] string path, IntPtr callback, ulong userData); } diff --git a/src/DataFusionSharp/Interop/PinnedProtobufData.cs b/src/DataFusionSharp/Interop/PinnedBytesData.cs similarity index 53% rename from src/DataFusionSharp/Interop/PinnedProtobufData.cs rename to src/DataFusionSharp/Interop/PinnedBytesData.cs index 7bdade0..c005b49 100644 --- a/src/DataFusionSharp/Interop/PinnedProtobufData.cs +++ b/src/DataFusionSharp/Interop/PinnedBytesData.cs @@ -4,14 +4,14 @@ namespace DataFusionSharp.Interop; /// -/// A helper struct that pins the byte array of a protobuf message in memory and provides a way to convert it to a BytesData struct for interop. +/// A helper struct that pins the byte array in memory and provides a way to convert it to a BytesData struct for interop. /// -internal readonly ref struct PinnedProtobufData : IDisposable +internal readonly ref struct PinnedBytesData : IDisposable { private readonly MemoryHandle? _handle; private readonly int _length; - private PinnedProtobufData(MemoryHandle? message, int length) + private PinnedBytesData(MemoryHandle? message, int length) { _handle = message; _length = length; @@ -29,14 +29,20 @@ public void Dispose() _handle?.Dispose(); } - public static PinnedProtobufData FromMessage(TMessage? message) + public static PinnedBytesData FromMemory(Memory bytes) + { + var handle = bytes.Pin(); + return new PinnedBytesData(handle, bytes.Length); + } + + public static PinnedBytesData FromMessage(TMessage? message) where TMessage : class, IMessage { if (message is null) - return new PinnedProtobufData(null, 0); + return new PinnedBytesData(null, 0); var bytes = message.ToByteArray(); var handle = bytes.AsMemory().Pin(); - return new PinnedProtobufData(handle, bytes.Length); + return new PinnedBytesData(handle, bytes.Length); } } \ No newline at end of file diff --git a/src/DataFusionSharp/Interop/SafeHandles.cs b/src/DataFusionSharp/Interop/SafeHandles.cs index bb249c1..ada86bb 100644 --- a/src/DataFusionSharp/Interop/SafeHandles.cs +++ b/src/DataFusionSharp/Interop/SafeHandles.cs @@ -86,3 +86,16 @@ protected override bool ReleaseHandle() return NativeMethods.DataFrameStreamDestroy(handle) == DataFusionErrorCode.Ok; } } + +internal sealed class InMemoryStoreSafeHandle : DataFusionSafeHandle +{ + internal InMemoryStoreSafeHandle(IntPtr handle) + : base(handle) + { + } + + protected override bool ReleaseHandle() + { + return NativeMethods.InMemoryStoreDestroy(handle) == DataFusionErrorCode.Ok; + } +} diff --git a/src/DataFusionSharp/ObjectStore/InMemoryObjectStore.cs b/src/DataFusionSharp/ObjectStore/InMemoryObjectStore.cs new file mode 100644 index 0000000..80a7542 --- /dev/null +++ b/src/DataFusionSharp/ObjectStore/InMemoryObjectStore.cs @@ -0,0 +1,119 @@ +using System.Buffers; +using DataFusionSharp.Interop; + +namespace DataFusionSharp.ObjectStore; + +/// +/// Represents an in-memory object store that can be used with the DataFusion runtime. +/// +public sealed class InMemoryObjectStore : IDisposable +{ + /// + /// Gets the DataFusion runtime associated with this in-memory object store. + /// + public DataFusionRuntime Runtime { get; } + + /// + /// Gets the handle to the in-memory store. + /// + internal InMemoryStoreSafeHandle Handle { get; } + + internal InMemoryObjectStore(DataFusionRuntime runtime, InMemoryStoreSafeHandle inMemoryStoreHandle) + { + Runtime = runtime; + Handle = inMemoryStoreHandle; + } + + /// + /// Puts data into the in-memory store at the specified path. + /// The data is provided as a byte array. + /// + /// + /// The native code copies the data from the provided byte array, so the caller can safely modify or dispose of the byte array after this method returns. + /// + /// Path to data + /// Bytes data to put + /// A task that completes when the put operation is finished. + /// Invalid path + /// Null data + /// Failed to put object into in-memory store + public Task PutAsync(string path, Memory data) + { + ArgumentException.ThrowIfNullOrEmpty(path); + if (data.IsEmpty) + throw new ArgumentException($"{nameof(data)} must not be empty.", nameof(data)); + + using var pinnedData = PinnedBytesData.FromMemory(data); + var bytesData = pinnedData.ToBytesData(); + + var (id, tcs) = AsyncOperations.Instance.Create(); + var result = NativeMethods.InMemoryStorePut(Handle, path, bytesData, true, GenericCallbacks.CallbackForVoidHandle, id); + if (result != DataFusionErrorCode.Ok) + { + AsyncOperations.Instance.Abort(id); + throw new DataFusionException(result, "Failed to put object into in-memory store."); + } + + return tcs.Task; + } + + /// + /// Puts data into the in-memory store at the specified path. + /// The data is provided as a pinned byte array. + /// + /// + /// The native code takes provided array as is without copying. + /// IMPORTANT: the caller must ensure that the memory is not modified or disposed for whole life of instance. + /// + /// Path to data + /// Pinned memory handle to byte data to put + /// Length of the data in bytes + /// A task that completes when the put operation is finished. + /// Invalid path + /// Null data + /// Failed to put object into in-memory store + public Task PutAsStaticAsync(string path, MemoryHandle memoryHandle, int length) + { + ArgumentException.ThrowIfNullOrEmpty(path); + + var (id, tcs) = AsyncOperations.Instance.Create(); + var bytesData = BytesData.FromPinned(memoryHandle, length); + var result = NativeMethods.InMemoryStorePut(Handle, path, bytesData, false, GenericCallbacks.CallbackForVoidHandle, id); + if (result != DataFusionErrorCode.Ok) + { + AsyncOperations.Instance.Abort(id); + throw new DataFusionException(result, "Failed to put object into in-memory store."); + } + + return tcs.Task; + } + + /// + /// Deletes data at the specified path from the in-memory store. + /// + /// Path to data to delete + /// A task that completes when the delete operation is finished. + /// Invalid path + /// Null data + /// Failed to delete object from in-memory store + public Task DeleteAsync(string path) + { + ArgumentException.ThrowIfNullOrEmpty(path); + + var (id, tcs) = AsyncOperations.Instance.Create(); + var result = NativeMethods.InMemoryStoreDelete(Handle, path, GenericCallbacks.CallbackForVoidHandle, id); + if (result != DataFusionErrorCode.Ok) + { + AsyncOperations.Instance.Abort(id); + throw new DataFusionException(result, "Failed to delete object from in-memory store."); + } + + return tcs.Task; + } + + /// + public void Dispose() + { + Handle.Dispose(); + } +} \ No newline at end of file diff --git a/src/DataFusionSharp/SessionContext.cs b/src/DataFusionSharp/SessionContext.cs index 41c8107..73ca01a 100644 --- a/src/DataFusionSharp/SessionContext.cs +++ b/src/DataFusionSharp/SessionContext.cs @@ -44,7 +44,7 @@ public Task RegisterCsvAsync(string tableName, string filePath, CsvReadOptions? ArgumentNullException.ThrowIfNull(tableName); ArgumentNullException.ThrowIfNull(filePath); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterCsv(_handle, tableName, filePath, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidHandle, id); @@ -70,7 +70,7 @@ public Task RegisterJsonAsync(string tableName, string filePath, JsonReadOptions ArgumentNullException.ThrowIfNull(tableName); ArgumentNullException.ThrowIfNull(filePath); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterJson(_handle, tableName, filePath, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidHandle, id); @@ -95,7 +95,7 @@ public Task RegisterParquetAsync(string tableName, string filePath, ParquetReadO ArgumentNullException.ThrowIfNull(tableName); ArgumentNullException.ThrowIfNull(filePath); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterParquet(_handle, tableName, filePath, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidHandle, id); @@ -204,7 +204,7 @@ public async Task SqlAsync(string sql, IEnumerable task; - using (var paramValuesData = PinnedProtobufData.FromMessage(parameters.ToProto())) + using (var paramValuesData = PinnedBytesData.FromMessage(parameters.ToProto())) { var (id, tcs) = AsyncOperations.Instance.Create(); var result = NativeMethods.ContextSql(_handle, sql, paramValuesData.ToBytesData(), CallbackForSqlAsyncHandle, id); @@ -253,7 +253,7 @@ public void RegisterS3ObjectStore(string url, S3ObjectStoreOptions? options = nu #pragma warning restore CA1054 { ArgumentNullException.ThrowIfNull(url); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var id = SyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterObjectStoreS3(_handle, url, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidSyncHandle, id); if (result != DataFusionErrorCode.Ok) @@ -275,7 +275,7 @@ public void RegisterAzureBlobStorage(string url, AzureBlobStorageOptions? option #pragma warning restore CA1054 { ArgumentNullException.ThrowIfNull(url); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var id = SyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterObjectStoreAzure(_handle, url, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidSyncHandle, id); if (result != DataFusionErrorCode.Ok) @@ -297,7 +297,7 @@ public void RegisterGoogleCloudStorage(string url, GoogleCloudStorageOptions? op #pragma warning restore CA1054 { ArgumentNullException.ThrowIfNull(url); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var id = SyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterObjectStoreGcs(_handle, url, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidSyncHandle, id); if (result != DataFusionErrorCode.Ok) @@ -319,7 +319,7 @@ public void RegisterHttpObjectStore(string url, HttpObjectStoreOptions? options #pragma warning restore CA1054 { ArgumentNullException.ThrowIfNull(url); - using var optionsData = PinnedProtobufData.FromMessage(options?.ToProto()); + using var optionsData = PinnedBytesData.FromMessage(options?.ToProto()); var id = SyncOperations.Instance.Create(); var result = NativeMethods.ContextRegisterObjectStoreHttp(_handle, url, optionsData.ToBytesData(), GenericCallbacks.CallbackForVoidSyncHandle, id); if (result != DataFusionErrorCode.Ok) @@ -329,6 +329,30 @@ public void RegisterHttpObjectStore(string url, HttpObjectStoreOptions? options } SyncOperations.Instance.TakeResult(id); } + + /// + /// Registers an in-memory object store for the given URL. + /// + /// The URL scheme to register (e.g., "memory://"). + /// The in-memory object store instance to register. + /// Null data + /// Thrown when registration fails. +#pragma warning disable CA1054 // URL is passed as-is to DataFusion's native, System.Uri would add redundant conversion. + public void RegisterInMemoryObjectStore(string url, InMemoryObjectStore store) +#pragma warning restore CA1054 + { + ArgumentNullException.ThrowIfNull(url); + ArgumentNullException.ThrowIfNull(store); + + var id = SyncOperations.Instance.Create(); + var result = NativeMethods.ContextRegisterObjectStoreInMemory(_handle, url, store.Handle, GenericCallbacks.CallbackForVoidSyncHandle, id); + if (result != DataFusionErrorCode.Ok) + { + SyncOperations.Instance.Abort(id); + throw new DataFusionException(result, "Failed to register in-memory object store"); + } + SyncOperations.Instance.TakeResult(id); + } /// /// Deregisters an object store for the given URL. diff --git a/tests/DataFusionSharp.Tests/InMemoryObjectStoreTests.cs b/tests/DataFusionSharp.Tests/InMemoryObjectStoreTests.cs new file mode 100644 index 0000000..c43635f --- /dev/null +++ b/tests/DataFusionSharp.Tests/InMemoryObjectStoreTests.cs @@ -0,0 +1,237 @@ +namespace DataFusionSharp.Tests; + +public sealed class InMemoryObjectStoreTests : IDisposable +{ + private readonly DataFusionRuntime _runtime = DataFusionRuntime.Create(); + + [Fact] + public void CreateInMemoryStore_ReturnsNonNullInstance() + { + // Act + using var store = _runtime.CreateInMemoryStore(); + + // Assert + Assert.NotNull(store); + Assert.Same(_runtime, store.Runtime); + } + + [Fact] + public async Task PutAsync_AndRegister_ThenQueryCsv_ReturnsData() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var csvBytes = await File.ReadAllBytesAsync(DataSet.CustomersCsvPath); + + // Act + await store.PutAsync("customers.csv", csvBytes); + context.RegisterInMemoryObjectStore("memory://", store); + await context.RegisterCsvAsync("customers", "memory:///customers.csv"); + + // Assert + using var df = await context.SqlAsync("SELECT * FROM customers"); + var count = await df.CountAsync(); + Assert.Equal(10UL, count); + } + + [Fact] + public async Task PutAsync_AndRegister_ThenQueryJson_ReturnsData() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var jsonBytes = await File.ReadAllBytesAsync(DataSet.CustomersJsonPath); + + // Act + await store.PutAsync("customers.json", jsonBytes); + context.RegisterInMemoryObjectStore("memory://", store); + await context.RegisterJsonAsync("customers", "memory:///customers.json"); + + // Assert + using var df = await context.SqlAsync("SELECT * FROM customers"); + var count = await df.CountAsync(); + Assert.Equal(10UL, count); + } + + [Fact] + public async Task PutAsync_AndRegister_ThenQueryParquet_ReturnsData() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var parquetBytes = await File.ReadAllBytesAsync(DataSet.CustomersParquetPath); + + // Act + await store.PutAsync("customers.parquet", parquetBytes); + context.RegisterInMemoryObjectStore("memory://", store); + await context.RegisterParquetAsync("customers", "memory:///customers.parquet"); + + // Assert + using var df = await context.SqlAsync("SELECT * FROM customers"); + var count = await df.CountAsync(); + Assert.Equal(10UL, count); + } + + [Fact] + public async Task DeleteAsync_RemovesData_QueryFails() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var csvBytes = await File.ReadAllBytesAsync(DataSet.CustomersCsvPath); + + await store.PutAsync("customers.csv", csvBytes); + + context.RegisterInMemoryObjectStore("memory://", store); + await context.RegisterCsvAsync("customers", "memory:///customers.csv"); + + // Sanity check — data is queryable before delete + using (var dfBefore = await context.SqlAsync("SELECT * FROM customers")) + { + var countBefore = await dfBefore.CountAsync(); + Assert.Equal(10UL, countBefore); + } + + // Act + await store.DeleteAsync("customers.csv"); + + // Assert + // Query should return no data after delete + using (var dfAfter = await context.SqlAsync("SELECT * FROM customers")) + { + var countAfter = await dfAfter.CountAsync(); + Assert.Equal(0UL, countAfter); + } + } + + [Fact] + public async Task PutAsync_NullArguments_Throw() + { + // Arrange + using var store = _runtime.CreateInMemoryStore(); + + // Act & Assert + await Assert.ThrowsAsync(() => store.PutAsync(null!, new byte[] {1, 2, 3})); + await Assert.ThrowsAsync(() => store.PutAsync("path", null!)); + } + + [Fact] + public async Task DeleteAsync_NullOrEmptyPath_Throws() + { + // Arrange + using var store = _runtime.CreateInMemoryStore(); + + // Act & Assert + await Assert.ThrowsAsync(() => store.DeleteAsync("")); + await Assert.ThrowsAsync(() => store.DeleteAsync(null!)); + } + + [Fact] + public async Task SameStore_RegisteredInTwoSessions_BothCanQuery() + { + // Arrange + using var contextA = _runtime.CreateSessionContext(); + using var contextB = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var csvBytes = await File.ReadAllBytesAsync(DataSet.CustomersCsvPath); + await store.PutAsync("customers.csv", csvBytes); + + // Act + contextA.RegisterInMemoryObjectStore("memory://", store); + contextB.RegisterInMemoryObjectStore("memory://", store); + + await contextA.RegisterCsvAsync("customers", "memory:///customers.csv"); + await contextB.RegisterCsvAsync("customers", "memory:///customers.csv"); + + // Assert + using (var dfA = await contextA.SqlAsync("SELECT * FROM customers")) + { + var countA = await dfA.CountAsync(); + Assert.Equal(10UL, countA); + } + + using (var dfB = await contextB.SqlAsync("SELECT * FROM customers")) + { + var countB = await dfB.CountAsync(); + Assert.Equal(10UL, countB); + } + } + + [Fact] + public async Task PutAsStaticAsync_AndRegister_ThenQueryCsv_ReturnsData() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var store = _runtime.CreateInMemoryStore(); + + var csvBytes = await File.ReadAllBytesAsync(DataSet.CustomersCsvPath); + using var memoryHandle = csvBytes.AsMemory().Pin(); + + // Act + await store.PutAsStaticAsync("customers.csv", memoryHandle, csvBytes.Length); + context.RegisterInMemoryObjectStore("memory://", store); + await context.RegisterCsvAsync("customers", "memory:///customers.csv"); + + // Assert + using var df = await context.SqlAsync("SELECT * FROM customers"); + var count = await df.CountAsync(); + Assert.Equal(10UL, count); + } + + [Fact] + public async Task TwoStores_RegisterBoth_ThenJoinQuery_ReturnsData() + { + // Arrange + using var context = _runtime.CreateSessionContext(); + using var storeA = _runtime.CreateInMemoryStore(); + using var storeB = _runtime.CreateInMemoryStore(); + + var customersCsv = await File.ReadAllBytesAsync(DataSet.CustomersCsvPath); + var ordersCsv = await File.ReadAllBytesAsync(DataSet.OrdersCsvPath); + + // Act + await storeA.PutAsync("customers.csv", customersCsv); + await storeB.PutAsync("orders.csv", ordersCsv); + + context.RegisterInMemoryObjectStore("mem-a://", storeA); + context.RegisterInMemoryObjectStore("mem-b://", storeB); + + await context.RegisterCsvAsync("customers", "mem-a:///customers.csv"); + await context.RegisterCsvAsync("orders", "mem-b:///orders.csv"); + + // Assert + using (var dfCustomers = await context.SqlAsync("SELECT * FROM customers")) + { + var count = await dfCustomers.CountAsync(); + Assert.Equal(10UL, count); + } + + using (var dfOrders = await context.SqlAsync("SELECT * FROM orders")) + { + var count = await dfOrders.CountAsync(); + Assert.Equal(50UL, count); + } + + using (var dfJoin = await context.SqlAsync( + """ + SELECT c.customer_name, o.order_id + FROM customers c + JOIN orders o ON c.customer_id = o.customer_id + """)) + { + var count = await dfJoin.CountAsync(); + Assert.Equal(50UL, count); + } + } + + public void Dispose() + { + _runtime.Dispose(); + } +} +