Add support for deregistering a table (#20)

Fix #15
datafusion-contrib · Aug 15, 2022 · e401f02 · e401f02
1 parent 4132420
commit e401f02
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 2 deletions.
diff --git a/datafusion-glib/session-context.c b/datafusion-glib/session-context.c
@@ -188,6 +188,41 @@ df_arrow_array_release_destroy(gpointer array)
   df_arrow_array_release(array);
 }
 
+/**
+ * gdf_session_context_deregister:
+ * @context: A #GDFSessionContext.
+ * @name: A name to be deregistered.
+ * @error: (nullable): Return location for a #GError or %NULL.
+ *
+ * Deregisters a registered table. If there isn't a registered table
+ * associated with @name, this function does nothing.
+ *
+ * Returns: %TRUE on success, %FALSE otherwise.
+ *
+ * Since: 10.0.0
+ */
+gboolean
+gdf_session_context_deregister(GDFSessionContext *context,
+                               const gchar *name,
+                               GError **error)
+{
+  GDFSessionContextPrivate *priv =
+    gdf_session_context_get_instance_private(context);
+  DFError *df_error = NULL;
+  bool success = df_session_context_deregister(priv->context, name, &df_error);
+  if (success) {
+    g_hash_table_remove(priv->registered_record_batches, name);
+  } else {
+    g_set_error(error,
+                GDF_ERROR,
+                df_error_get_code(df_error),
+                "[session-context][deregister] %s",
+                df_error_get_message(df_error));
+    df_error_free(df_error);
+  }
+  return success;
+}
+
 /**
  * gdf_session_context_register_record_batch:
  * @context: A #GDFSessionContext.
@@ -233,7 +268,6 @@ gdf_session_context_register_record_batch(GDFSessionContext *context,
     df_arrow_schema_release(c_abi_schema);
     df_arrow_array_release(c_abi_array);
   } else {
-    /* TODO: Remove this on unregister */
     g_hash_table_insert(priv->registered_record_batches,
                         g_strdup(name),
                         g_list_prepend(NULL, g_object_ref(record_batch)));
@@ -324,7 +358,6 @@ gdf_session_context_register_table(GDFSessionContext *context,
     df_arrow_schema_release(c_abi_schema);
   } else {
     g_ptr_array_set_free_func(c_abi_arrays, NULL);
-    /* TODO: Remove this on unregister */
     g_hash_table_insert(priv->registered_record_batches,
                         g_strdup(name),
                         record_batches);

diff --git a/datafusion-glib/session-context.h b/datafusion-glib/session-context.h
@@ -46,6 +46,11 @@ gdf_session_context_sql(GDFSessionContext *context,
                         GError **error);
 GDF_AVAILABLE_IN_10_0
 gboolean
+gdf_session_context_deregister(GDFSessionContext *context,
+                               const gchar *name,
+                               GError **error);
+GDF_AVAILABLE_IN_10_0
+gboolean
 gdf_session_context_register_record_batch(GDFSessionContext *context,
                                           const gchar *name,
                                           GArrowRecordBatch *record_batch,

diff --git a/src/capi.rs b/src/capi.rs
@@ -404,6 +404,25 @@ pub extern "C" fn df_session_context_sql(
     maybe_data_frame.map(|data_frame| Box::new(DFDataFrame::new(data_frame)))
 }
 
+#[no_mangle]
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+pub extern "C" fn df_session_context_deregister(
+    context: &mut DFSessionContext,
+    name: *const libc::c_char,
+    error: *mut *mut DFError,
+) -> bool {
+    let option = || -> Option<bool> {
+        let cstr_name = unsafe { CStr::from_ptr(name) };
+        let rs_name = cstr_name.to_str().into_df_error(error, None)?;
+        context
+            .context
+            .deregister_table(rs_name)
+            .into_df_error(error, None)?;
+        Some(true)
+    }();
+    option.unwrap_or(false)
+}
+
 #[no_mangle]
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 pub extern "C" fn df_session_context_register_record_batches(

diff --git a/test/test-session-context.rb b/test/test-session-context.rb
@@ -26,6 +26,24 @@ def test_sql
                  @context.sql("SELECT 1").to_table)
   end
 
+  def test_deregister
+    record_batch = Arrow::RecordBatch.new(boolean: [true, false, nil],
+                                          integer: [1, nil, 3])
+    assert do
+      @context.register_record_batch("data", record_batch)
+    end
+    data_frame = @context.sql("SELECT * FROM data")
+    assert_equal(record_batch.to_table, data_frame.to_table)
+    assert do
+      @context.deregister("data")
+    end
+    message = "[session-context][sql] Error during planning: " +
+              "'datafusion.public.data' not found"
+    assert_raise(DataFusion::Error::Plan.new(message)) do
+      @context.sql("SELECT * FROM data")
+    end
+  end
+
   def test_register_record_batch
     record_batch = Arrow::RecordBatch.new(boolean: [true, false, nil],
                                           integer: [1, nil, 3])