Add missing deregister methods to SessionContext

crates/core/src/context.rs

@@ -434,11 +434,25 @@ impl PySessionContext {
  &upstream_host
  };
  let url_string = format!("{scheme}{derived_host}");
- let url = Url::parse(&url_string).unwrap();
+ let url = Url::parse(&url_string).map_err(|e| PyValueError::new_err(e.to_string()))?;
  self.ctx.runtime_env().register_object_store(&url, store);
  Ok(())
  }
 
+ /// Deregister an object store with the given url
+ #[pyo3(signature = (scheme, host=None))]
+ pub fn deregister_object_store(
+ &self,
+ scheme: &str,
+ host: Option<&str>,
+ ) -> PyDataFusionResult<()> {
+ let host = host.unwrap_or("");
+ let url_string = format!("{scheme}{host}");
+ let url = Url::parse(&url_string).map_err(|e| PyDataFusionError::Common(e.to_string()))?;
+ self.ctx.runtime_env().deregister_object_store(&url)?;
+ Ok(())
+ }
+
  #[allow(clippy::too_many_arguments)]
  #[pyo3(signature = (name, path, table_partition_cols=vec![],
  file_extension=".parquet",
@@ -492,6 +506,10 @@ impl PySessionContext {
  self.ctx.register_udtf(&name, func);
  }
 
+ pub fn deregister_udtf(&self, name: &str) {
+ self.ctx.deregister_udtf(name);
+ }
+
  #[pyo3(signature = (query, options=None, param_values=HashMap::default(), param_strings=HashMap::default()))]
  pub fn sql_with_options(
  &self,
@@ -975,16 +993,28 @@ impl PySessionContext {
  Ok(())
  }
 
+ pub fn deregister_udf(&self, name: &str) {
+ self.ctx.deregister_udf(name);
+ }
+
  pub fn register_udaf(&self, udaf: PyAggregateUDF) -> PyResult<()> {
  self.ctx.register_udaf(udaf.function);
  Ok(())
  }
 
+ pub fn deregister_udaf(&self, name: &str) {
+ self.ctx.deregister_udaf(name);
+ }
+
  pub fn register_udwf(&self, udwf: PyWindowUDF) -> PyResult<()> {
  self.ctx.register_udwf(udwf.function);
  Ok(())
  }
 
+ pub fn deregister_udwf(&self, name: &str) {
+ self.ctx.deregister_udwf(name);
+ }
+
  #[pyo3(signature = (name="datafusion"))]
  pub fn catalog(&self, py: Python, name: &str) -> PyResult<Py<PyAny>> {
  let catalog = self.ctx.catalog(name).ok_or(PyKeyError::new_err(format!(

python/datafusion/context.py

@@ -568,6 +568,15 @@ def register_object_store(
  """
  self.ctx.register_object_store(schema, store, host)
 
+ def deregister_object_store(self, schema: str, host: str | None = None) -> None:
+ """Remove an object store from the session.
+
+ Args:
+ schema: The data source schema (e.g. ``"s3://"``).
+ host: URL for the host (e.g. bucket name).
+ """
+ self.ctx.deregister_object_store(schema, host)
+
  def register_listing_table(
  self,
  name: str,
@@ -894,6 +903,14 @@ def register_udtf(self, func: TableFunction) -> None:
  """Register a user defined table function."""
  self.ctx.register_udtf(func._udtf)
 
+ def deregister_udtf(self, name: str) -> None:
+ """Remove a user-defined table function from the session.
+
+ Args:
+ name: Name of the UDTF to deregister.
+ """
+ self.ctx.deregister_udtf(name)
+
  def register_record_batches(
  self, name: str, partitions: list[list[pa.RecordBatch]]
  ) -> None:
@@ -1105,14 +1122,38 @@ def register_udf(self, udf: ScalarUDF) -> None:
  """Register a user-defined function (UDF) with the context."""
  self.ctx.register_udf(udf._udf)
 
+ def deregister_udf(self, name: str) -> None:
+ """Remove a user-defined scalar function from the session.
+
+ Args:
+ name: Name of the UDF to deregister.
+ """
+ self.ctx.deregister_udf(name)
+
  def register_udaf(self, udaf: AggregateUDF) -> None:
  """Register a user-defined aggregation function (UDAF) with the context."""
  self.ctx.register_udaf(udaf._udaf)
 
+ def deregister_udaf(self, name: str) -> None:
+ """Remove a user-defined aggregate function from the session.
+
+ Args:
+ name: Name of the UDAF to deregister.
+ """
+ self.ctx.deregister_udaf(name)
+
  def register_udwf(self, udwf: WindowUDF) -> None:
  """Register a user-defined window function (UDWF) with the context."""
  self.ctx.register_udwf(udwf._udwf)
 
+ def deregister_udwf(self, name: str) -> None:
+ """Remove a user-defined window function from the session.
+
+ Args:
+ name: Name of the UDWF to deregister.
+ """
+ self.ctx.deregister_udwf(name)
+
  def catalog(self, name: str = "datafusion") -> Catalog:
  """Retrieve a catalog by name."""
  return Catalog(self.ctx.catalog(name))

python/tests/test_context.py

@@ -31,6 +31,7 @@
  Table,
  column,
  literal,
+ udf,
 )
 
 
@@ -351,6 +352,125 @@ def test_deregister_table(ctx, database):
  assert public.names() == {"csv1", "csv2"}
 
 
+def test_deregister_udf():
+ ctx = SessionContext()
+
+ is_null = udf(
+ lambda x: x.is_null(),
+ [pa.float64()],
+ pa.bool_(),
+ volatility="immutable",
+ name="my_is_null",
+ )
+ ctx.register_udf(is_null)
+
+ # Verify it works
+ df = ctx.from_pydict({"a": [1.0, None]})
+ ctx.register_table("t", df.into_view())
+ result = ctx.sql("SELECT my_is_null(a) FROM t").collect()
+ assert result[0].column(0) == pa.array([False, True])
+
+ # Deregister and verify it's gone
+ ctx.deregister_udf("my_is_null")
+ with pytest.raises(ValueError):
+ ctx.sql("SELECT my_is_null(a) FROM t").collect()
+
+
+def test_deregister_udaf():
+ import pyarrow.compute as pc
+
+ ctx = SessionContext()
+ from datafusion import Accumulator, udaf
+
+ class MySum(Accumulator):
+ def __init__(self):
+ self._sum = 0.0
+
+ def update(self, values: pa.Array) -> None:
+ self._sum += pc.sum(values).as_py()
+
+ def merge(self, states: list[pa.Array]) -> None:
+ self._sum += pc.sum(states[0]).as_py()
+
+ def state(self) -> list:
+ return [self._sum]
+
+ def evaluate(self) -> pa.Scalar:
+ return self._sum
+
+ my_sum = udaf(
+ MySum,
+ [pa.float64()],
+ pa.float64(),
+ [pa.float64()],
+ volatility="immutable",
+ name="my_sum",
+ )
+ ctx.register_udaf(my_sum)
+ df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})
+ ctx.register_table("t", df.into_view())
+
+ result = ctx.sql("SELECT my_sum(a) FROM t").collect()
+ assert result[0].column(0) == pa.array([6.0])
+
+ ctx.deregister_udaf("my_sum")
+ with pytest.raises(ValueError):
+ ctx.sql("SELECT my_sum(a) FROM t").collect()
+
+
+def test_deregister_udwf():
+ ctx = SessionContext()
+ from datafusion import udwf
+ from datafusion.user_defined import WindowEvaluator
+
+ class MyRowNumber(WindowEvaluator):
+ def __init__(self):
+ self._row = 0
+
+ def evaluate_all(self, values, num_rows):
+ return pa.array(list(range(1, num_rows + 1)), type=pa.uint64())
+
+ my_row_number = udwf(
+ MyRowNumber,
+ [pa.float64()],
+ pa.uint64(),
+ volatility="immutable",
+ name="my_row_number",
+ )
+ ctx.register_udwf(my_row_number)
+ df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]})
+ ctx.register_table("t", df.into_view())
+
+ result = ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect()
+ assert result[0].column(0) == pa.array([1, 2, 3], type=pa.uint64())
+
+ ctx.deregister_udwf("my_row_number")
+ with pytest.raises(ValueError):
+ ctx.sql("SELECT my_row_number(a) OVER () FROM t").collect()
+
+
+def test_deregister_udtf():
+ import pyarrow.dataset as ds
+
+ ctx = SessionContext()
+ from datafusion import Table, udtf
+
+ class MyTable:
+ def __call__(self):
+ batch = pa.RecordBatch.from_pydict({"x": [1, 2, 3]})
+ return Table(ds.dataset([batch]))
+
+ my_table = udtf(MyTable(), "my_table")
+ ctx.register_udtf(my_table)
+
+ result = ctx.sql("SELECT * FROM my_table()").collect()
+ assert result[0].column(0) == pa.array([1, 2, 3])
+
+ ctx.deregister_udtf("my_table")
+ with pytest.raises(ValueError):
+ ctx.sql("SELECT * FROM my_table()").collect()
+
+
 def test_register_table_from_dataframe(ctx):
  df = ctx.from_pydict({"a": [1, 2]})
  ctx.register_table("df_tbl", df)