feat: add cloning support for SessionContext (#91)

Author: nazarii-piontko

native/src/context.rs

@@ -6,6 +6,7 @@ use tokio_util::sync::CancellationToken;
 
 use crate::{BytesData, Callback, ErrorCode, ErrorInfo, mappers, proto};
 
+#[derive(Clone)]
 pub struct SessionContextWrapper {
     runtime: crate::RuntimeHandle,
     inner: Arc<datafusion::prelude::SessionContext>,
@@ -22,28 +23,30 @@ impl SessionContextWrapper {
 
 /// Creates a new `SessionContext` bound to a runtime.
 ///
+/// This is a synchronous operation.
+///
 /// # Safety
 /// - `runtime_ptr` must be a valid pointer returned by `datafusion_runtime_new`
-/// - `context_ptr` must be a valid, aligned, non-null pointer to writable memory
+/// - `context_out_ptr` must be a valid, aligned, non-null pointer to writable memory
 /// - Caller must call `datafusion_context_destroy` exactly once with the returned pointer
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn datafusion_context_new(
     runtime_ptr: *mut crate::RuntimeHandle,
-    context_ptr: *mut *mut SessionContextWrapper,
+    context_out_ptr: *mut *mut SessionContextWrapper,
 ) -> ErrorCode {
-    if context_ptr.is_null() {
+    if context_out_ptr.is_null() {
         return ErrorCode::InvalidArgument;
     }
 
     let runtime_handle = ffi_ref!(runtime_ptr);
 
     let context = Box::new(SessionContextWrapper::new(Arc::clone(runtime_handle)));
-    let raw_ptr = Box::into_raw(context);
+    let context_ptr = Box::into_raw(context);
     unsafe {
-        *context_ptr = raw_ptr;
+        *context_out_ptr = context_ptr;
     }
 
-    debug!("Created session context {raw_ptr:p}");
+    debug!("Created session context {context_ptr:p}");
 
     ErrorCode::Ok
 }
@@ -66,6 +69,37 @@ pub unsafe extern "C" fn datafusion_context_destroy(
     ErrorCode::Ok
 }
 
+/// Clones a `SessionContext`, creating a new wrapper that shares the same underlying context.
+///
+/// This is a synchronous operation.
+///
+/// # Safety
+/// - `context_ptr` must be a valid pointer returned by `datafusion_context_new` or `datafusion_context_clone`
+/// - `context_out_ptr` must be a valid, aligned, non-null pointer to writable memory
+/// - Caller must call `datafusion_context_destroy` exactly once with the returned pointer
+#[unsafe(no_mangle)]
+pub unsafe extern "C" fn datafusion_context_clone(
+    context_ptr: *mut SessionContextWrapper,
+    context_out_ptr: *mut *mut SessionContextWrapper,
+) -> ErrorCode {
+    if context_out_ptr.is_null() {
+        return ErrorCode::InvalidArgument;
+    }
+
+    let context = ffi_ref!(context_ptr);
+
+    let cloned_context = Box::new(context.clone());
+    let cloned_context_ptr = Box::into_raw(cloned_context);
+
+    unsafe {
+        *context_out_ptr = cloned_context_ptr;
+    }
+
+    debug!("Cloned session context {context_ptr:p} -> {cloned_context_ptr:p}");
+
+    ErrorCode::Ok
+}
+
 /// Registers a CSV file as a table in the `SessionContext`.
 ///
 /// This is an async operation. The callback is invoked on completion with no result data.

src/DataFusionSharp/Interop/NativeMethods.cs

@@ -66,6 +66,9 @@ public static partial DataFusionErrorCode Ping(
     [LibraryImport(LibraryName, EntryPoint = "datafusion_context_destroy")]
     public static partial DataFusionErrorCode ContextDestroy(IntPtr contextHandle);
 
+    [LibraryImport(LibraryName, EntryPoint = "datafusion_context_clone")]
+    public static partial DataFusionErrorCode ContextClone(SessionContextSafeHandle contextHandle, out IntPtr clonedContextHandle);
+    
     [LibraryImport(LibraryName, EntryPoint = "datafusion_context_register_csv")]
     public static partial DataFusionErrorCode ContextRegisterCsv(
         SessionContextSafeHandle contextHandle,

src/DataFusionSharp/SessionContext.cs

@@ -32,6 +32,22 @@ internal SessionContext(DataFusionRuntime runtime, SessionContextSafeHandle hand
         _handle = handle;
     }
 
+    /// <summary>
+    /// Creates a clone of this session context.
+    /// </summary>
+    /// <remarks>
+    /// The cloned context is an independent wrapper with its own lifecycle.
+    /// </remarks>
+    /// <returns>A cloned <see cref="SessionContext"/>.</returns>
+    /// <exception cref="DataFusionException">Thrown when cloning fails.</exception>
+    public SessionContext Clone()
+    {
+        var errorCode = NativeMethods.ContextClone(_handle, out var clonedContextHandle);
+        DataFusionException.ThrowIfError(errorCode, "Failed to clone DataFusion context");
+
+        return new SessionContext(Runtime, new SessionContextSafeHandle(clonedContextHandle));
+    }
+
     /// <summary>
     /// Registers a CSV file as a table in this session.
     /// </summary>

tests/DataFusionSharp.Tests/SessionContextTests.cs

@@ -173,6 +173,46 @@ public async Task SqlAsync_WithMissingParameter_Throws()
         Assert.Contains("$value_b", ex.Message, StringComparison.Ordinal);
         Assert.DoesNotContain("$value_a", ex.Message, StringComparison.Ordinal);
     }
+    
+    [Fact]
+    public async Task Clone_RegisteredTableIsAvailableInClonedContext()
+    {
+        // Arrange
+        using var context = _runtime.CreateSessionContext();
+        await context.RegisterCsvAsync("customers", DataSet.CustomersCsvPath);
+
+        // Act
+        using var clonedContext = context.Clone();
+        using var clonedDataFrame = await clonedContext.SqlAsync("SELECT * FROM customers");
+        var count = await clonedDataFrame.CountAsync();
+
+        // Assert
+        Assert.NotSame(context, clonedContext);
+        Assert.Same(context.Runtime, clonedContext.Runtime);
+        Assert.Equal(10UL, count);
+    }
+
+    [Fact]
+    public async Task Clone_AfterOriginalDisposed_RemainsValid()
+    {
+        // Arrange
+        SessionContext clonedContext;
+        using (var context = _runtime.CreateSessionContext())
+        {
+            await context.RegisterCsvAsync("customers", DataSet.CustomersCsvPath);
+            clonedContext = context.Clone();
+        }
+
+        // Act
+        using (clonedContext)
+        {
+            using var clonedDataFrame = await clonedContext.SqlAsync("SELECT * FROM customers");
+            var count = await clonedDataFrame.CountAsync();
+
+            // Assert
+            Assert.Equal(10UL, count);
+        }
+    }
 
     public void Dispose()
     {