Repository: Eventual-Inc/Daft
Pull request #4352: refactor(exprs): url upload/download to new pattern
Author: login: universalmind303
Labels: name: refactor
Created: 2025-05-14T18:07:30Z | Merged: 2025-05-17T02:51:14Z

## Changes Made

<!-- Describe what changes were made and why. Include implementation details if necessary. -->

## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [ ] Documented in API Docs (if applicable)
- [ ] Documented in User Guide (if applicable)
- [ ] If adding a new documentation page, doc is added to `docs/mkdocs.yml` navigation
- [ ] Documentation builds and is formatted properly (tag @/ccmao1130 for docs review)



Diff:
diff --git a/daft/expressions/expressions.py b/daft/expressions/expressions.py
index 1c0a2e56ba..0657ac1250 100644
--- a/daft/expressions/expressions.py
+++ b/daft/expressions/expressions.py
@@ -43,7 +43,6 @@
 from daft.daft import tokenize_decode as _tokenize_decode
 from daft.daft import tokenize_encode as _tokenize_encode
 from daft.daft import udf as _udf
-from daft.daft import url_download as _url_download
 from daft.daft import utf8_count_matches as _utf8_count_matches
 from daft.datatype import DataType, DataTypeLike, TimeUnit
 from daft.dependencies import pa
@@ -1889,21 +1888,23 @@ def download(
             ``max_connections`` to get higher throughput with additional parallelism
 
         """
-        raise_on_error = False
-        if on_error == "raise":
-            raise_on_error = True
-        elif on_error == "null":
-            raise_on_error = False
-        else:
-            raise NotImplementedError(f"Unimplemented on_error option: {on_error}.")
-
-        if not (isinstance(max_connections, int) and max_connections > 0):
-            raise ValueError(f"Invalid value for `max_connections`: {max_connections}")
-
         multi_thread = ExpressionUrlNamespace._should_use_multithreading_tokio_runtime()
         io_config = ExpressionUrlNamespace._override_io_config_max_connections(max_connections, io_config)
+        max_connections_expr = Expression._to_expression(max_connections)._expr
+        on_error_expr = Expression._to_expression(on_error)._expr
+        multi_thread_expr = Expression._to_expression(multi_thread)._expr
+        io_config_expr = Expression._to_expression(io_config)._expr
+
+        f = native.get_function_from_registry("url_download")
+
         return Expression._from_pyexpr(
-            _url_download(self._expr, max_connections, raise_on_error, multi_thread, io_config)
+            f(
+                self._expr,
+                multi_thread=multi_thread_expr,
+                on_error=on_error_expr,
+                max_connections=max_connections_expr,
+                io_config=io_config_expr,
+            )
         )
 
     def upload(
@@ -1936,32 +1937,28 @@ def upload(
             >>> col("data").url.upload(col("paths"))  # doctest: +SKIP
 
         """
-        if not (isinstance(max_connections, int) and max_connections > 0):
-            raise ValueError(f"Invalid value for `max_connections`: {max_connections}")
-
-        location_expr = Expression._to_expression(location)
-        raise_on_error = False
-        if on_error == "raise":
-            raise_on_error = True
-        elif on_error == "null":
-            raise_on_error = False
-        else:
-            raise NotImplementedError(f"Unimplemented on_error option: {on_error}.")
+        location_expr = Expression._to_expression(location)._expr
         multi_thread = ExpressionUrlNamespace._should_use_multithreading_tokio_runtime()
         # If the user specifies a single location via a string, we should upload to a single folder. Otherwise,
         # if the user gave an expression, we assume that each row has a specific url to upload to.
         # Consider moving the check for is_single_folder to a lower IR.
         is_single_folder = isinstance(location, str)
         io_config = ExpressionUrlNamespace._override_io_config_max_connections(max_connections, io_config)
+        max_connections_expr = Expression._to_expression(max_connections)._expr
+        on_error_expr = Expression._to_expression(on_error)._expr
+        multi_thread_expr = Expression._to_expression(multi_thread)._expr
+        io_config_expr = Expression._to_expression(io_config)._expr
+        is_single_folder_expr = Expression._to_expression(is_single_folder)._expr
+        f = native.get_function_from_registry("url_upload")
         return Expression._from_pyexpr(
-            native.url_upload(
+            f(
                 self._expr,
-                location_expr._expr,
-                max_connections,
-                raise_on_error,
-                multi_thread,
-                is_single_folder,
-                io_config,
+                location_expr,
+                max_connections=max_connections_expr,
+                on_error=on_error_expr,
+                multi_thread=multi_thread_expr,
+                is_single_folder=is_single_folder_expr,
+                io_config=io_config_expr,
             )
         )
 
diff --git a/src/daft-functions/src/uri/download.rs b/src/daft-functions/src/uri/download.rs
index 056211b3ca..c4c91df8d1 100644
--- a/src/daft-functions/src/uri/download.rs
+++ b/src/daft-functions/src/uri/download.rs
@@ -1,25 +1,28 @@
 use std::sync::Arc;
 
-use common_error::{DaftError, DaftResult};
+use common_error::{ensure, DaftError, DaftResult};
 use common_runtime::get_io_runtime;
 use daft_core::prelude::*;
-use daft_dsl::{functions::ScalarUDF, ExprRef};
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
 use daft_io::{get_io_client, Error, IOConfig, IOStatsContext, IOStatsRef};
 use futures::{StreamExt, TryStreamExt};
 use serde::Serialize;
-use snafu::prelude::*;
-
-use crate::InvalidArgumentSnafu;
 
 /// Container for the keyword arguments of `url_download`
 /// ex:
 /// ```text
-/// url_decode(input)
-/// url_decode(input, max_connections=32)
-/// url_decode(input, on_error='raise')
-/// url_decode(input, on_error='null')
-/// url_decode(input, max_connections=32, on_error='raise')
+/// url_download(input)
+/// url_download(input, max_connections=32)
+/// url_download(input, on_error='raise')
+/// url_download(input, on_error='null')
+/// url_download(input, max_connections=32, on_error='raise')
 /// ```
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlDownload;
+
 #[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
 pub struct UrlDownloadArgs {
     pub max_connections: usize,
@@ -56,67 +59,62 @@ impl Default for UrlDownloadArgs {
 }
 
 #[typetag::serde]
-impl ScalarUDF for UrlDownloadArgs {
-    fn evaluate(&self, inputs: daft_dsl::functions::FunctionArgs<Series>) -> DaftResult<Series> {
-        let inner = inputs.into_inner();
-        self.evaluate_from_series(&inner)
-    }
-
+impl ScalarUDF for UrlDownload {
     fn name(&self) -> &'static str {
-        "download"
+        "url_download"
     }
-
-    fn evaluate_from_series(&self, inputs: &[Series]) -> DaftResult<Series> {
-        let Self {
+    fn evaluate(&self, inputs: daft_dsl::functions::FunctionArgs<Series>) -> DaftResult<Series> {
+        let input = inputs.required((0, "input"))?;
+        let max_connections = inputs
+            .extract_optional::<usize, _>("max_connections")?
+            .unwrap_or(32);
+        let on_error = inputs
+            .extract_optional::<String, _>("on_error")?
+            .unwrap_or_else(|| "raise".to_string());
+
+        let raise_error_on_failure = match on_error.as_str() {
+            "raise" => true,
+            "null" => false,
+            _ => {
+                return Err(DaftError::ValueError(format!(
+                    "Invalid value for 'on_error': {}",
+                    on_error
+                )))
+            }
+        };
+
+        let multi_thread = inputs
+            .extract_optional::<bool, _>("multi_thread")?
+            .unwrap_or(true);
+        let io_config = inputs
+            .extract_optional::<IOConfig, _>("io_config")?
+            .unwrap_or_default();
+
+        let array = input.utf8()?;
+        let io_stats = IOStatsContext::new("download");
+        let result = url_download(
+            array,
             max_connections,
             raise_error_on_failure,
             multi_thread,
-            io_config,
-        } = self;
-
-        match inputs {
-            [input] => match input.data_type() {
-                DataType::Utf8 => {
-                    let array = input.utf8()?;
-                    let io_stats = IOStatsContext::new("download");
-                    let result = url_download(
-                        array,
-                        *max_connections,
-                        *raise_error_on_failure,
-                        *multi_thread,
-                        io_config.clone(),
-                        Some(io_stats),
-                    )?;
-                    Ok(result.into_series())
-                }
-                _ => Err(DaftError::TypeError(format!(
-                    "Download can only download uris from Utf8Array, got {input}"
-                ))),
-            },
-            _ => Err(DaftError::ValueError(format!(
-                "Expected 1 input arg, got {}",
-                inputs.len()
-            ))),
-        }
+            Arc::new(io_config),
+            Some(io_stats),
+        )?;
+        Ok(result.into_series())
     }
 
-    fn to_field(&self, inputs: &[ExprRef], schema: &Schema) -> DaftResult<Field> {
-        match inputs {
-            [input] => {
-                let field = input.to_field(schema)?;
-
-                match &field.dtype {
-                    DataType::Utf8 => Ok(Field::new(field.name, DataType::Binary)),
-                    _ => Err(DaftError::TypeError(format!(
-                        "Download can only download uris from Utf8Array, got {field}"
-                    ))),
-                }
-            }
-            _ => Err(DaftError::SchemaMismatch(format!(
-                "Expected 1 input arg, got {}",
-                inputs.len()
-            ))),
-        }
+    fn function_args_to_field(
+        &self,
+        inputs: FunctionArgs<ExprRef>,
+        schema: &Schema,
+    ) -> DaftResult<Field> {
+        let field = inputs.required((0, "input"))?.to_field(schema)?;
+        let _ = inputs.extract_optional::<bool, _>("multi_thread")?;
+        let _ = inputs.extract_optional::<IOConfig, _>("io_config")?;
+        let _ = inputs.extract_optional::<usize, _>("max_connections")?;
+        let _ = inputs.extract_optional::<String, _>("on_error")?;
+        ensure!(field.dtype.is_string(), TypeError: "Input must be a string");
+        Ok(Field::new(field.name, DataType::Binary))
     }
 }
 
@@ -131,9 +129,7 @@ fn url_download(
     let name = array.name();
     ensure!(
         max_connections > 0,
-        InvalidArgumentSnafu {
-            msg: "max_connections for url_download must be non-zero".to_owned()
-        }
+        ValueError: "max_connections for url_download must be non-zero"
     );
 
     let runtime_handle = get_io_runtime(true);
diff --git a/src/daft-functions/src/uri/mod.rs b/src/daft-functions/src/uri/mod.rs
index 541af8ef8d..5a10bfe2bb 100644
--- a/src/daft-functions/src/uri/mod.rs
+++ b/src/daft-functions/src/uri/mod.rs
@@ -1,18 +1,70 @@
 pub mod download;
 pub mod upload;
+use std::sync::Arc;
 
-use daft_dsl::{functions::ScalarFunction, ExprRef};
-use download::UrlDownloadArgs;
-use upload::UrlUploadArgs;
+use daft_dsl::{
+    functions::{FunctionModule, ScalarFunction},
+    lit, ExprRef,
+};
+use download::{UrlDownload, UrlDownloadArgs};
+use upload::{UrlUpload, UrlUploadArgs};
 
 /// Creates a `url_download` ExprRef from the positional and optional named arguments.
 #[must_use]
 pub fn download(input: ExprRef, args: Option<UrlDownloadArgs>) -> ExprRef {
-    ScalarFunction::new(args.unwrap_or_default(), vec![input]).into()
+    let args = if let Some(UrlDownloadArgs {
+        max_connections,
+        raise_error_on_failure,
+        multi_thread,
+        io_config,
+    }) = args
+    {
+        vec![
+            input,
+            lit(max_connections as u64),
+            lit(raise_error_on_failure),
+            lit(multi_thread),
+            lit(Arc::unwrap_or_clone(io_config)),
+        ]
+    } else {
+        vec![input]
+    };
+
+    ScalarFunction::new(UrlDownload, args).into()
 }
 
 /// Creates a `url_upload` ExprRef from the positional and optional named arguments.
 #[must_use]
 pub fn upload(input: ExprRef, location: ExprRef, args: Option<UrlUploadArgs>) -> ExprRef {
-    ScalarFunction::new(args.unwrap_or_default(), vec![input, location]).into()
+    let args = if let Some(UrlUploadArgs {
+        max_connections,
+        raise_error_on_failure,
+        multi_thread,
+        is_single_folder,
+        io_config,
+    }) = args
+    {
+        vec![
+            input,
+            location,
+            lit(max_connections as u64),
+            lit(raise_error_on_failure),
+            lit(multi_thread),
+            lit(is_single_folder),
+            lit(Arc::unwrap_or_clone(io_config)),
+        ]
+    } else {
+        vec![input, location]
+    };
+
+    ScalarFunction::new(UrlUpload, args).into()
+}
+
+pub struct UriFunctions;
+
+impl FunctionModule for UriFunctions {
+    fn register(parent: &mut daft_dsl::functions::FunctionRegistry) {
+        parent.add_fn(dbg!(UrlDownload));
+        parent.add_fn(UrlUpload);
+    }
 }
diff --git a/src/daft-functions/src/uri/upload.rs b/src/daft-functions/src/uri/upload.rs
index ac8aeb9c19..82941e9fbf 100644
--- a/src/daft-functions/src/uri/upload.rs
+++ b/src/daft-functions/src/uri/upload.rs
@@ -1,9 +1,12 @@
 use std::{collections::HashSet, iter::repeat_n, path::Path, sync::Arc};
 
-use common_error::{DaftError, DaftResult};
+use common_error::{ensure, DaftError, DaftResult};
 use common_runtime::get_io_runtime;
 use daft_core::prelude::*;
-use daft_dsl::{functions::ScalarUDF, ExprRef};
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
 use daft_io::{get_io_client, IOConfig, IOStatsRef, SourceType};
 use futures::{StreamExt, TryStreamExt};
 use serde::Serialize;
@@ -17,6 +20,9 @@ pub struct UrlUploadArgs {
     pub io_config: Arc<IOConfig>,
 }
 
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlUpload;
+
 impl UrlUploadArgs {
     pub fn new(
         max_connections: usize,
@@ -48,65 +54,70 @@ impl Default for UrlUploadArgs {
 }
 
 #[typetag::serde]
-impl ScalarUDF for UrlUploadArgs {
+impl ScalarUDF for UrlUpload {
     fn evaluate(&self, inputs: daft_dsl::functions::FunctionArgs<Series>) -> DaftResult<Series> {
-        let inner = inputs.into_inner();
-        self.evaluate_from_series(&inner)
-    }
+        let input = inputs.required((0, "input"))?;
+        let location = inputs.required((1, "location"))?;
 
-    fn name(&self) -> &'static str {
-        "upload"
-    }
+        let max_connections = inputs.extract_optional("max_connections")?.unwrap_or(32);
+
+        let on_error = inputs
+            .extract_optional("on_error")?
+            .unwrap_or_else(|| "raise".to_string());
+
+        let raise_error_on_failure = match on_error.as_str() {
+            "raise" => true,
+            "null" => false,
+            _ => {
+                return Err(DaftError::ValueError(format!(
+                    "Invalid value for 'on_error': {}",
+                    on_error
+                )))
+            }
+        };
+        let multi_thread = inputs.extract_optional("multi_thread")?.unwrap_or(true);
+
+        let is_single_folder = inputs
+            .extract_optional("is_single_folder")?
+            .unwrap_or(false);
 
-    fn evaluate_from_series(&self, inputs: &[Series]) -> DaftResult<Series> {
-        let Self {
+        let io_config = inputs.extract_optional("io_config")?.unwrap_or_default();
+
+        url_upload(
+            input,
+            location,
             max_connections,
             raise_error_on_failure,
             multi_thread,
             is_single_folder,
-            io_config,
-        } = self;
+            Arc::new(io_config),
+            None,
+        )
+    }
 
-        match inputs {
-            [data, location] => url_upload(
-                data,
-                location,
-                *max_connections,
-                *raise_error_on_failure,
-                *multi_thread,
-                *is_single_folder,
-                io_config.clone(),
-                None,
-            ),
-            _ => Err(DaftError::ValueError(format!(
-                "Expected 2 input args, got {}",
-                inputs.len()
-            ))),
-        }
+    fn name(&self) -> &'static str {
+        "url_upload"
     }
 
-    fn to_field(&self, inputs: &[ExprRef], schema: &Schema) -> DaftResult<Field> {
-        match inputs {
-            [data, location] => {
-                let data_field = data.to_field(schema)?;
-                let location_field = location.to_field(schema)?;
-                match data_field.dtype {
-                    DataType::Binary | DataType::FixedSizeBinary(..) | DataType::Utf8 => (),
-                    _ => return Err(DaftError::TypeError(format!("Expects input to url_upload to be Binary, FixedSizeBinary or String, but received {data_field}"))),
-                }
-                if !location_field.dtype.is_string() {
-                    return Err(DaftError::TypeError(format!(
-                        "Expected location to be string, received: {}",
-                        location_field.dtype
-                    )));
-                }
-                Ok(Field::new(data_field.name, DataType::Utf8))
-            }
-            _ => Err(DaftError::SchemaMismatch(format!(
-                "Expected 2 input args, got {}",
-                inputs.len()
-            ))),
-        }
+    fn function_args_to_field(
+        &self,
+        inputs: FunctionArgs<ExprRef>,
+        schema: &Schema,
+    ) -> DaftResult<Field> {
+        let field = inputs.required((0, "input"))?.to_field(schema)?;
+        ensure!(
+            field.dtype.is_binary() || field.dtype.is_fixed_size_binary() || field.dtype.is_string(),
+            TypeError: "input must be a binary, fixed-size binary or utf8"
+        );
+
+        let location = inputs.required((1, "location"))?.to_field(schema)?;
+        ensure!(location.dtype.is_string(), TypeError: "location must be a string");
+        let _ = inputs.extract_optional::<usize, _>("max_connections")?;
+        let _ = inputs.extract_optional::<String, _>("on_error")?;
+        let _ = inputs.extract_optional::<bool, _>("multi_thread")?;
+        let _ = inputs.extract_optional::<bool, _>("is_single_folder")?;
+        let _ = inputs.extract_optional::<IOConfig, _>("io_config")?;
+        Ok(Field::new(field.name, DataType::Utf8))
     }
 }
 
diff --git a/src/daft-sql/src/functions.rs b/src/daft-sql/src/functions.rs
index 3e2ea7e174..3f8dd68760 100644
--- a/src/daft-sql/src/functions.rs
+++ b/src/daft-sql/src/functions.rs
@@ -19,8 +19,7 @@ use crate::{
     modules::{
         coalesce::SQLCoalesce, hashing::SQLModuleHashing, SQLModule, SQLModuleAggs,
         SQLModuleConfig, SQLModuleList, SQLModuleMap, SQLModulePartitioning, SQLModulePython,
-        SQLModuleSketch, SQLModuleStructs, SQLModuleTemporal, SQLModuleUri, SQLModuleUtf8,
-        SQLModuleWindow,
+        SQLModuleSketch, SQLModuleStructs, SQLModuleTemporal, SQLModuleUtf8, SQLModuleWindow,
     },
     planner::SQLPlanner,
     unsupported_sql_err,
@@ -38,7 +37,6 @@ pub(crate) static SQL_FUNCTIONS: LazyLock<SQLFunctions> = LazyLock::new(|| {
     functions.register::<SQLModuleSketch>();
     functions.register::<SQLModuleStructs>();
     functions.register::<SQLModuleTemporal>();
-    functions.register::<SQLModuleUri>();
     functions.register::<SQLModuleUtf8>();
     functions.register::<SQLModuleConfig>();
     functions.register::<SQLModuleWindow>();
@@ -624,19 +622,7 @@ pub(crate) mod args {
     use common_io_config::IOConfig;
 
     use super::SQLFunctionArguments;
-    use crate::{error::PlannerError, modules::config::expr_to_iocfg, unsupported_sql_err};
-
-    /// Parses on_error => Literal['raise', 'null'] = 'raise' or err.
-    pub(crate) fn parse_on_error(args: &SQLFunctionArguments) -> Result<bool, PlannerError> {
-        match args.try_get_named::<String>("on_error")?.as_deref() {
-            None => Ok(true),
-            Some("raise") => Ok(true),
-            Some("null") => Ok(false),
-            Some(other) => {
-                unsupported_sql_err!("Expected on_error to be 'raise' or 'null', found '{other}'")
-            }
-        }
-    }
+    use crate::{error::PlannerError, modules::config::expr_to_iocfg};
 
     /// Parses io_config which is used in several SQL functions.
     pub(crate) fn parse_io_config(args: &SQLFunctionArguments) -> Result<IOConfig, PlannerError> {
diff --git a/src/daft-sql/src/modules/mod.rs b/src/daft-sql/src/modules/mod.rs
index 42d7a67bbe..fef2320dfa 100644
--- a/src/daft-sql/src/modules/mod.rs
+++ b/src/daft-sql/src/modules/mod.rs
@@ -11,7 +11,6 @@ pub mod python;
 pub mod sketch;
 pub mod structs;
 pub mod temporal;
-pub mod uri;
 pub mod utf8;
 pub mod window;
 
@@ -24,7 +23,6 @@ pub use python::SQLModulePython;
 pub use sketch::SQLModuleSketch;
 pub use structs::SQLModuleStructs;
 pub use temporal::SQLModuleTemporal;
-pub use uri::SQLModuleUri;
 pub use utf8::SQLModuleUtf8;
 pub use window::SQLModuleWindow;
 
diff --git a/src/daft-sql/src/modules/uri/mod.rs b/src/daft-sql/src/modules/uri/mod.rs
deleted file mode 100644
index 5c748c9ece..0000000000
--- a/src/daft-sql/src/modules/uri/mod.rs
+++ /dev/null
@@ -1,14 +0,0 @@
-use super::SQLModule;
-use crate::functions::SQLFunctions;
-
-mod url_download;
-mod url_upload;
-
-pub struct SQLModuleUri;
-
-impl SQLModule for SQLModuleUri {
-    fn register(parent: &mut SQLFunctions) {
-        parent.add_fn("url_download", url_download::SqlUrlDownload);
-        parent.add_fn("url_upload", url_upload::SqlUrlUpload);
-    }
-}
diff --git a/src/daft-sql/src/modules/uri/url_download.rs b/src/daft-sql/src/modules/uri/url_download.rs
deleted file mode 100644
index e7f0706929..0000000000
--- a/src/daft-sql/src/modules/uri/url_download.rs
+++ /dev/null
@@ -1,58 +0,0 @@
-use daft_dsl::ExprRef;
-use daft_functions::uri::{self, download::UrlDownloadArgs};
-use sqlparser::ast::FunctionArg;
-
-use crate::{
-    error::{PlannerError, SQLPlannerResult},
-    functions::{self, SQLFunction, SQLFunctionArguments},
-    unsupported_sql_err, SQLPlanner,
-};
-
-/// The Daft-SQL `url_download` definition.
-pub struct SqlUrlDownload;
-
-impl TryFrom<SQLFunctionArguments> for UrlDownloadArgs {
-    type Error = PlannerError;
-
-    fn try_from(args: SQLFunctionArguments) -> Result<Self, Self::Error> {
-        let max_connections: usize = args.try_get_named("max_connections")?.unwrap_or(32);
-        let raise_error_on_failure = functions::args::parse_on_error(&args)?;
-        let io_config = functions::args::parse_io_config(&args)?;
-        Ok(Self {
-            max_connections,
-            raise_error_on_failure,
-            multi_thread: true, // TODO always true
-            io_config: io_config.into(),
-        })
-    }
-}
-
-impl SQLFunction for SqlUrlDownload {
-    fn to_expr(&self, inputs: &[FunctionArg], planner: &SQLPlanner) -> SQLPlannerResult<ExprRef> {
-        match inputs {
-            [input] => {
-                let input = planner.plan_function_arg(input)?.into_inner();
-                Ok(uri::download(input, None))
-            }
-            [input, args @ ..] => {
-                let input = planner.plan_function_arg(input)?.into_inner();
-                let args = planner.plan_function_args(
-                    args,
-                    &["max_connections", "on_error", "io_config"],
-                    0,
-                )?;
-                Ok(uri::download(input, Some(args)))
-            }
-            _ => unsupported_sql_err!("Invalid arguments for url_download: '{inputs:?}'"),
-        }
-    }
-
-    fn docstrings(&self, _alias: &str) -> String {
-        "Treats each string as a URL, and downloads the bytes contents as a bytes column."
-            .to_string()
-    }
-
-    fn arg_names(&self) -> &'static [&'static str] {
-        &["input", "max_connections", "on_error", "io_config"]
-    }
-}
diff --git a/src/daft-sql/src/modules/uri/url_upload.rs b/src/daft-sql/src/modules/uri/url_upload.rs
deleted file mode 100644
index 381819be5f..0000000000
--- a/src/daft-sql/src/modules/uri/url_upload.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-use daft_dsl::{Expr, ExprRef, LiteralValue};
-use daft_functions::uri::{self, upload::UrlUploadArgs};
-use sqlparser::ast::FunctionArg;
-
-use crate::{
-    error::{PlannerError, SQLPlannerResult},
-    functions::{self, SQLFunction, SQLFunctionArguments},
-    unsupported_sql_err, SQLPlanner,
-};
-
-/// The Daft-SQL `url_upload` definition.
-pub struct SqlUrlUpload;
-
-impl TryFrom<SQLFunctionArguments> for UrlUploadArgs {
-    type Error = PlannerError;
-
-    fn try_from(args: SQLFunctionArguments) -> Result<Self, Self::Error> {
-        let max_connections: usize = args.try_get_named("max_connections")?.unwrap_or(32);
-        let raise_error_on_failure = functions::args::parse_on_error(&args)?;
-        let io_config = functions::args::parse_io_config(&args)?;
-        Ok(Self {
-            max_connections,
-            raise_error_on_failure,
-            multi_thread: true, // TODO always true
-            is_single_folder: true,
-            io_config: io_config.into(),
-        })
-    }
-}
-
-impl SQLFunction for SqlUrlUpload {
-    fn to_expr(&self, inputs: &[FunctionArg], planner: &SQLPlanner) -> SQLPlannerResult<ExprRef> {
-        match inputs {
-            [input, location] => {
-                let input = planner.plan_function_arg(input)?.into_inner();
-                let location = planner.plan_function_arg(location)?.into_inner();
-                Ok(uri::upload(input, location, None))
-            }
-            [input, location, args @ ..] => {
-                let input = planner.plan_function_arg(input)?.into_inner();
-                let location = planner.plan_function_arg(location)?.into_inner();
-                let mut args: UrlUploadArgs = planner.plan_function_args(
-                    args,
-                    &["max_connections", "on_error", "io_config"],
-                    0,
-                )?;
-                // TODO consider moving the calculation of "is_single_folder" deeper so that both the SQL and Python sides don't compute this independently.
-                // is_single_folder = true iff the location is a string.
-                // in python, this is isinstance(location, str)
-                // We perform the check here (with mut args) because TryFrom<SQLFunctionArguments> does not have access to `location`.
-                args.is_single_folder =
-                    matches!(location.as_ref(), Expr::Literal(LiteralValue::Utf8(_)));
-                Ok(uri::upload(input, location, Some(args)))
-            }
-            _ => unsupported_sql_err!("Invalid arguments for url_upload: '{inputs:?}'"),
-        }
-    }
-
-    fn docstrings(&self, _alias: &str) -> String {
-        "Uploads a column of binary data to the provided location(s) (also supports S3, local etc)."
-            .to_string()
-    }
-
-    fn arg_names(&self) -> &'static [&'static str] {
-        &[
-            "input",
-            "location",
-            "max_connections",
-            "on_error",
-            "io_config",
-        ]
-    }
-}
diff --git a/src/lib.rs b/src/lib.rs
index 2be556eef4..fd39900377 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -153,6 +153,7 @@ pub mod pylib {
             .expect("Failed to acquire write lock on function registry");
         functions_registry.register::<daft_functions::numeric::NumericFunctions>();
         functions_registry.register::<daft_functions::float::FloatFunctions>();
+        functions_registry.register::<daft_functions::uri::UriFunctions>();
         functions_registry.register::<daft_image::functions::ImageFunctions>();
         functions_registry.register::<daft_functions_json::JsonFunctions>();
 
diff --git a/tests/io/test_url_upload_local.py b/tests/io/test_url_upload_local.py
index 3a23d0d8a5..8e75b1a76f 100644
--- a/tests/io/test_url_upload_local.py
+++ b/tests/io/test_url_upload_local.py
@@ -76,7 +76,7 @@ def test_upload_local_no_write_permissions(tmpdir):
     data = {"data": bytes_data, "paths": paths}
     df = daft.from_pydict(data)
     df_raise_error = df.with_column("files", df["data"].url.upload(df["paths"]))
-    with pytest.raises(ValueError, match="Unable to write data to file"):
+    with pytest.raises(daft.exceptions.DaftCoreException):
         df_raise_error.collect()
     # Retry with `on_error` set to `null`.
     df_null = df.with_column("files", df["data"].url.upload(df["paths"], on_error="null"))


Reviews:
user: login: srilman
body: LGTM, just had 1 question. Will let Kevin do the final approval
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-14T19:15:59Z
user: login: universalmind303
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-14T19:25:46Z
user: login: kevinzwang
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-16T20:43:59Z
user: login: universalmind303
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-17T01:10:09Z

Review comments:
diff_hunk: @@ -1,25 +1,28 @@
 use std::sync::Arc;
 
-use common_error::{DaftError, DaftResult};
+use common_error::{ensure, DaftError, DaftResult};
 use common_runtime::get_io_runtime;
 use daft_core::prelude::*;
-use daft_dsl::{functions::ScalarUDF, ExprRef};
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
 use daft_io::{get_io_client, Error, IOConfig, IOStatsContext, IOStatsRef};
 use futures::{StreamExt, TryStreamExt};
 use serde::Serialize;
-use snafu::prelude::*;
-
-use crate::InvalidArgumentSnafu;
 
 /// Container for the keyword arguments of `url_download`
 /// ex:
 /// ```text
-/// url_decode(input)
-/// url_decode(input, max_connections=32)
-/// url_decode(input, on_error='raise')
-/// url_decode(input, on_error='null')
-/// url_decode(input, max_connections=32, on_error='raise')
+/// url_download(input)
+/// url_download(input, max_connections=32)
+/// url_download(input, on_error='raise')
+/// url_download(input, on_error='null')
+/// url_download(input, max_connections=32, on_error='raise')
 /// ```
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlDownload;
+
 #[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
 pub struct UrlDownloadArgs {
path: src/daft-functions/src/uri/download.rs
user: login: srilman
body: Is this class still necessary if the udf can track its inputs now?
created_at: 2025-05-14T19:14:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,25 +1,28 @@
 use std::sync::Arc;
 
-use common_error::{DaftError, DaftResult};
+use common_error::{ensure, DaftError, DaftResult};
 use common_runtime::get_io_runtime;
 use daft_core::prelude::*;
-use daft_dsl::{functions::ScalarUDF, ExprRef};
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
 use daft_io::{get_io_client, Error, IOConfig, IOStatsContext, IOStatsRef};
 use futures::{StreamExt, TryStreamExt};
 use serde::Serialize;
-use snafu::prelude::*;
-
-use crate::InvalidArgumentSnafu;
 
 /// Container for the keyword arguments of `url_download`
 /// ex:
 /// ```text
-/// url_decode(input)
-/// url_decode(input, max_connections=32)
-/// url_decode(input, on_error='raise')
-/// url_decode(input, on_error='null')
-/// url_decode(input, max_connections=32, on_error='raise')
+/// url_download(input)
+/// url_download(input, max_connections=32)
+/// url_download(input, on_error='raise')
+/// url_download(input, on_error='null')
+/// url_download(input, max_connections=32, on_error='raise')
 /// ```
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlDownload;
+
 #[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
 pub struct UrlDownloadArgs {
path: src/daft-functions/src/uri/download.rs
user: login: universalmind303
body: We create unit types like this for registering the function, so yes it's still needed. 
created_at: 2025-05-14T19:25:45Z
author_association: MEMBER
diff_hunk: @@ -1,25 +1,28 @@
 use std::sync::Arc;
 
-use common_error::{DaftError, DaftResult};
+use common_error::{ensure, DaftError, DaftResult};
 use common_runtime::get_io_runtime;
 use daft_core::prelude::*;
-use daft_dsl::{functions::ScalarUDF, ExprRef};
+use daft_dsl::{
+    functions::{FunctionArgs, ScalarUDF},
+    ExprRef,
+};
 use daft_io::{get_io_client, Error, IOConfig, IOStatsContext, IOStatsRef};
 use futures::{StreamExt, TryStreamExt};
 use serde::Serialize;
-use snafu::prelude::*;
-
-use crate::InvalidArgumentSnafu;
 
 /// Container for the keyword arguments of `url_download`
 /// ex:
 /// ```text
-/// url_decode(input)
-/// url_decode(input, max_connections=32)
-/// url_decode(input, on_error='raise')
-/// url_decode(input, on_error='null')
-/// url_decode(input, max_connections=32, on_error='raise')
+/// url_download(input)
+/// url_download(input, max_connections=32)
+/// url_download(input, on_error='raise')
+/// url_download(input, on_error='null')
+/// url_download(input, max_connections=32, on_error='raise')
 /// ```
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlDownload;
+
 #[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
 pub struct UrlDownloadArgs {
path: src/daft-functions/src/uri/download.rs
user: login: kevinzwang
body: Wait @universalmind303 I don't see `UrlDownloadArgs` being used anywhere except for a test and a, now unused, `url_download` pyfunction. Are you sure this is needed?

When calling URL download through the function registry, it seems like we bypass the creation of `UrlDownloadArgs`
created_at: 2025-05-16T20:41:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -17,6 +20,9 @@ pub struct UrlUploadArgs {
     pub io_config: Arc<IOConfig>,
 }
 
+#[derive(Debug, Clone, Serialize, serde::Deserialize, PartialEq, Eq, Hash)]
+pub struct UrlUpload;
+
 impl UrlUploadArgs {
path: src/daft-functions/src/uri/upload.rs
user: login: kevinzwang
body: same comment as `UrlDownloadArgs`
created_at: 2025-05-16T20:42:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -154,6 +154,7 @@ pub mod pylib {
             .expect("Failed to acquire write lock on function registry");
         functions_registry.register::<daft_functions::numeric::NumericFunctions>();
         functions_registry.register::<daft_functions::float::FloatFunctions>();
+        functions_registry.register::<daft_functions::uri::UriFunctions>();
         functions_registry.register::<daft_image::functions::ImageFunctions>();
path: src/lib.rs
user: login: kevinzwang
body: Btw, would this be the right place to populate the function registry? `FUNCTION_REGISTRY` would be empty if the Python feature was not enabled
created_at: 2025-05-16T20:43:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -154,6 +154,7 @@ pub mod pylib {
             .expect("Failed to acquire write lock on function registry");
         functions_registry.register::<daft_functions::numeric::NumericFunctions>();
         functions_registry.register::<daft_functions::float::FloatFunctions>();
+        functions_registry.register::<daft_functions::uri::UriFunctions>();
         functions_registry.register::<daft_image::functions::ImageFunctions>();
path: src/lib.rs
user: login: universalmind303
body: we could add a [ctor](https://docs.rs/ctor/latest/ctor/attr.ctor.html) to the `src/lib`, but since our build commands all expect python anyways, this seemed like a reasonable approach. The only time we **need** all of the functions is when we do all of the linking at the end here. 
created_at: 2025-05-17T01:10:09Z
author_association: MEMBER

PR comments:
user: login: srilman
created_at: 2025-05-16T20:39:11Z
body: @kevinzwang bump, need this to unblock another PR of mine
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-05-17T02:51:35Z
body: ## [Codecov](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) Report
Attention: Patch coverage is `78.41727%` with `30 lines` in your changes missing coverage. Please review.
> Project coverage is 78.03%. Comparing base [(`d58ce8e`)](https://app.codecov.io/gh/Eventual-Inc/Daft/commit/d58ce8ebcbf4ddb0e1a2021328213bb2a9776e55?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) to head [(`7a53547`)](https://app.codecov.io/gh/Eventual-Inc/Daft/commit/7a53547153a6a3039c978c6fa23efa77e03415d5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc).
> Report is 2 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) | Patch % | Lines |
|---|---|---|
| [src/daft-functions/src/uri/mod.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fmod.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvbW9kLnJz) | 47.36% | [20 Missing :warning: ](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) |
| [src/daft-functions/src/uri/download.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fdownload.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvZG93bmxvYWQucnM=) | 88.37% | [5 Missing :warning: ](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) |
| [src/daft-functions/src/uri/upload.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fupload.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvdXBsb2FkLnJz) | 88.88% | [5 Missing :warning: ](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) |

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352/graphs/tree.svg?width=650&height=150&src=pr&token=J430QVFE89&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc)](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc)

```diff
@@            Coverage Diff             @@
##             main    #4352      +/-   ##
==========================================
- Coverage   78.11%   78.03%   -0.09%     
==========================================
  Files         828      826       -2     
  Lines      110934   111272     +338     
==========================================
+ Hits        86658    86832     +174     
- Misses      24276    24440     +164     
```

| [Files with missing lines](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc) | Coverage Δ | |
|---|---|---|
| [daft/expressions/expressions.py](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=daft%2Fexpressions%2Fexpressions.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-ZGFmdC9leHByZXNzaW9ucy9leHByZXNzaW9ucy5weQ==) | `96.32% <100.00%> (+0.41%)` | :arrow_up: |
| [src/daft-sql/src/functions.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-sql%2Fsrc%2Ffunctions.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtc3FsL3NyYy9mdW5jdGlvbnMucnM=) | `76.36% <ø> (-3.59%)` | :arrow_down: |
| [src/lib.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Flib.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2xpYi5ycw==) | `91.75% <100.00%> (+0.08%)` | :arrow_up: |
| [src/daft-functions/src/uri/download.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fdownload.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvZG93bmxvYWQucnM=) | `88.96% <88.37%> (+2.47%)` | :arrow_up: |
| [src/daft-functions/src/uri/upload.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fupload.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvdXBsb2FkLnJz) | `68.67% <88.88%> (-6.04%)` | :arrow_down: |
| [src/daft-functions/src/uri/mod.rs](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352?src=pr&el=tree&filepath=src%2Fdaft-functions%2Fsrc%2Furi%2Fmod.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc#diff-c3JjL2RhZnQtZnVuY3Rpb25zL3NyYy91cmkvbW9kLnJz) | `48.78% <47.36%> (-51.22%)` | :arrow_down: |

... and [18 files with indirect coverage changes](https://app.codecov.io/gh/Eventual-Inc/Daft/pull/4352/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Eventual-Inc)

</details>
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
author_association: NONE