Merge branch main into colin/hide-series-lit

Author: colin-ho

Cargo.lock

@@ -64,7 +64,7 @@ build-release: check-toolchain frontend .venv  ## Compile and install a faster D
 
 .PHONY: test
 test: .venv build  ## Run tests
-	HYPOTHESIS_MAX_EXAMPLES=$(HYPOTHESIS_MAX_EXAMPLES) $(VENV_BIN)/pytest --hypothesis-seed=$(HYPOTHESIS_SEED)
+	HYPOTHESIS_MAX_EXAMPLES=$(HYPOTHESIS_MAX_EXAMPLES) $(VENV_BIN)/pytest --hypothesis-seed=$(HYPOTHESIS_SEED) --ignore tests/integration
 
 .PHONY: doctests
 doctests:

Makefile

@@ -113,10 +113,6 @@ def lit(value: object) -> Expression:
         sign, digits, exponent = value.as_tuple()
         assert isinstance(exponent, int)
         lit_value = _decimal_lit(sign == 1, digits, exponent)
-    elif isinstance(value, ImageFormat):
-        lit_value = _lit(str(value))
-    elif isinstance(value, ImageMode):
-        lit_value = _lit(str(value))
     else:
         lit_value = _lit(value)
     return Expression._from_pyexpr(lit_value)

daft/expressions/expressions.py

@@ -1,6 +1,6 @@
 [dependencies]
 aws-credential-types = {version = "0.55.3"}
-chrono = {workspace = true}
+chrono = {workspace = true, features = ["serde"]}
 common-error = {path = "../error", default-features = false}
 common-py-serde = {path = "../py-serde", default-features = false}
 derivative = {workspace = true}

src/common/io-config/Cargo.toml

@@ -4,7 +4,7 @@ use common_error::DaftResult;
 use common_file_formats::{FileFormat, WriteMode};
 use daft_catalog::TableSource;
 use daft_context::get_context;
-use daft_dsl::LiteralValue;
+use daft_dsl::{literals_to_series, LiteralValue};
 use daft_logical_plan::LogicalPlanBuilder;
 use daft_micropartition::MicroPartition;
 use daft_recordbatch::RecordBatch;
@@ -366,9 +366,7 @@ impl ConnectSession {
         let tbl = RecordBatch::concat(&tbls)?;
         let output = tbl.to_comfy_table(None).to_string();
 
-        let s = LiteralValue::Utf8(output)
-            .into_single_value_series()?
-            .rename("show_string");
+        let s = literals_to_series(&[LiteralValue::Utf8(output)])?.rename("show_string");
 
         let tbl = RecordBatch::from_nonempty_columns(vec![s])?;
         response_builder.arrow_batch_response(&tbl)

src/daft-connect/src/execute.rs

@@ -2,6 +2,7 @@
 bincode = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 common-hashable-float-wrapper = {path = "../common/hashable-float-wrapper"}
+common-io-config = {path = "../common/io-config", default-features = false}
 common-py-serde = {path = "../common/py-serde", default-features = false}
 common-resource-request = {path = "../common/resource-request", default-features = false}
 common-treenode = {path = "../common/treenode", default-features = false}
@@ -10,12 +11,21 @@ daft-sketch = {path = "../daft-sketch", default-features = false}
 derive_more = {workspace = true}
 indexmap = {workspace = true}
 itertools = {workspace = true}
+num-traits = {workspace = true}
 pyo3 = {workspace = true, optional = true}
 serde = {workspace = true}
 typetag = {workspace = true}
 
 [features]
-python = ["dep:pyo3", "common-error/python", "daft-core/python", "common-treenode/python", "common-py-serde/python", "common-resource-request/python"]
+python = [
+  "dep:pyo3",
+  "common-error/python",
+  "daft-core/python",
+  "common-io-config/python",
+  "common-treenode/python",
+  "common-py-serde/python",
+  "common-resource-request/python"
+]
 test-utils = []
 
 [lints]

src/daft-dsl/Cargo.toml

@@ -1,8 +1,11 @@
 use std::sync::Arc;
 
 use common_error::{DaftError, DaftResult};
+use daft_core::series::Series;
 use serde::{Deserialize, Serialize};
 
+use crate::{lit::FromLiteral, ExprRef, LiteralValue};
+
 /// Wrapper around T to hold either a named or an unnamed argument.
 #[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq, Hash)]
 pub enum FunctionArg<T> {
@@ -388,11 +391,106 @@ impl<T> FunctionArgs<T> {
     }
 }
 
+impl FunctionArgs<Series> {
+    /// Uses serde to extract a scalar value from the function args.
+    /// This will error if the series is not scalar. (len ==1) or if it is not deserializable to the provided type.
+    /// It will also return an error if the value does not exist.
+    pub fn extract<V: FromLiteral, Key: FunctionArgKey>(&self, position: Key) -> DaftResult<V> {
+        let value = position.required(self).map_err(|_| {
+            DaftError::ValueError(format!(
+                "Expected a value for the required argument at position `{position:?}`"
+            ))
+        })?;
+
+        let lit = LiteralValue::try_from_single_value_series(value)?;
+
+        let res = V::try_from_literal(&lit);
+        res.map_err(|e| e.into())
+    }
+
+    /// Uses serde to extract an optional scalar value from the function args.
+    /// This will error if the series is not scalar. (len ==1) or if it is not deserializable to the provided type.
+    /// if the value does not exist, None is returned.
+    pub fn extract_optional<V: FromLiteral, Key: FunctionArgKey>(
+        &self,
+        position: Key,
+    ) -> DaftResult<Option<V>> {
+        let value = position.optional(self).map_err(|_| {
+            DaftError::ValueError(format!(
+                "Expected a value for the optional argument at position `{position:?}`"
+            ))
+        })?;
+
+        match value {
+            Some(value) => {
+                let lit = LiteralValue::try_from_single_value_series(value)?;
+                let res = V::try_from_literal(&lit);
+                res.map_err(|e| e.into()).map(Some)
+            }
+            None => Ok(None),
+        }
+    }
+}
+
+impl FunctionArgs<ExprRef> {
+    /// Uses serde to extract a scalar value from the function args.
+    /// This will error if the the expr is not a literal, or if it is not deserializable to the provided type.
+    /// It will also error if the value does not exist.
+    pub fn extract<V: FromLiteral, Key: FunctionArgKey>(&self, position: Key) -> DaftResult<V> {
+        let value = position.required(self).map_err(|_| {
+            DaftError::ValueError(format!(
+                "Expected a value for the required argument at position `{position:?}`"
+            ))
+        })?;
+
+        match value.as_literal() {
+            Some(lit) => {
+                let res = V::try_from_literal(lit);
+                res.map_err(|e| e.into())
+            }
+            None => Err(DaftError::ValueError(format!(
+                "Expected a literal value for the optional argument at position `{position:?}`"
+            ))),
+        }
+    }
+    /// Uses serde to extract an optional scalar value from the function args.
+    /// This will error if the the expr is not a literal, or if it is not deserializable to the provided type.
+    /// if the value does not exist, None is returned.
+    pub fn extract_optional<V: FromLiteral, Key: FunctionArgKey>(
+        &self,
+        position: Key,
+    ) -> DaftResult<Option<V>> {
+        let value = position.optional(self).map_err(|_| {
+            DaftError::ValueError(format!(
+                "Expected a value for the optional argument at position `{position:?}`"
+            ))
+        })?;
+
+        match value {
+            Some(value) => match value.as_literal() {
+                Some(lit) => {
+                    let res = V::try_from_literal(lit);
+                    res.map_err(|e| e.into()).map(Some)
+                }
+                None => Err(DaftError::ValueError(format!(
+                    "Expected a literal value for the optional argument at position `{position:?}`"
+                ))),
+            },
+            None => Ok(None),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use common_error::DaftResult;
+    use common_io_config::IOConfig;
+    use daft_core::prelude::CountMode;
 
-    use crate::functions::function_args::{FunctionArg, FunctionArgs};
+    use crate::{
+        functions::function_args::{FunctionArg, FunctionArgs},
+        lit, Literal,
+    };
     #[test]
     fn test_function_args_ordering() {
         let res = FunctionArgs::try_new(vec![
@@ -533,4 +631,26 @@ mod tests {
         assert!(args.is_empty());
         assert_eq!(args.len(), 0);
     }
+
+    #[test]
+    fn test_extract() -> DaftResult<()> {
+        let io_conf = IOConfig::default();
+        let count_mode = CountMode::Valid;
+        let args = FunctionArgs::new_unchecked(vec![
+            FunctionArg::unnamed(100i64.lit()),
+            FunctionArg::unnamed(222i32.lit()),
+            FunctionArg::named("io_config", lit(io_conf.clone())),
+            FunctionArg::named("arg2", lit(count_mode.clone())),
+        ]);
+
+        let res: usize = args.extract(0)?;
+        assert_eq!(res, 100);
+        let second_pos: usize = args.extract(1)?;
+        assert_eq!(second_pos, 222);
+        let io_conf_extracted: IOConfig = args.extract("io_config")?;
+        assert_eq!(io_conf_extracted, io_conf);
+        let count_mode_extracted: CountMode = args.extract("arg2")?;
+        assert_eq!(count_mode_extracted, count_mode);
+        Ok(())
+    }
 }

src/daft-dsl/src/functions/function_args.rs

@@ -0,0 +1,157 @@
+use common_error::{DaftError, DaftResult};
+use common_io_config::IOConfig;
+use daft_core::{
+    datatypes::IntervalValue,
+    prelude::{CountMode, ImageFormat, ImageMode},
+    series::Series,
+};
+use num_traits::cast;
+use serde::de::DeserializeOwned;
+
+use super::{deserializer, serializer, FromLiteral, Literal, LiteralValue};
+#[cfg(feature = "python")]
+use crate::pyobj_serde::PyObjectWrapper;
+
+impl Literal for IntervalValue {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Interval(self)
+    }
+}
+
+impl Literal for String {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Utf8(self)
+    }
+}
+
+impl FromLiteral for String {
+    fn try_from_literal(lit: &LiteralValue) -> DaftResult<Self> {
+        match lit {
+            LiteralValue::Utf8(s) => Ok(s.clone()),
+            _ => Err(DaftError::TypeError(format!(
+                "Cannot convert {:?} to String",
+                lit
+            ))),
+        }
+    }
+}
+
+impl Literal for &'_ str {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Utf8(self.to_owned())
+    }
+}
+
+impl Literal for &'_ [u8] {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Binary(self.to_vec())
+    }
+}
+
+impl Literal for Series {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Series(self)
+    }
+}
+
+#[cfg(feature = "python")]
+impl Literal for pyo3::PyObject {
+    fn literal_value(self) -> LiteralValue {
+        LiteralValue::Python(PyObjectWrapper(std::sync::Arc::new(self)))
+    }
+}
+
+impl<T> Literal for Option<T>
+where
+    T: Literal,
+{
+    fn literal_value(self) -> LiteralValue {
+        match self {
+            Some(val) => val.literal_value(),
+            None => LiteralValue::Null,
+        }
+    }
+}
+
+macro_rules! make_literal {
+    ($TYPE:ty, $SCALAR:ident) => {
+        impl Literal for $TYPE {
+            fn literal_value(self) -> LiteralValue {
+                LiteralValue::$SCALAR(self)
+            }
+        }
+        impl FromLiteral for $TYPE {
+            fn try_from_literal(lit: &LiteralValue) -> DaftResult<Self> {
+                match lit {
+                    LiteralValue::$SCALAR(v) => Ok(*v),
+                    _ => Err(DaftError::TypeError(format!(
+                        "Expected {} literal",
+                        stringify!($TYPE)
+                    ))),
+                }
+            }
+        }
+    };
+}
+make_literal!(bool, Boolean);
+make_literal!(i8, Int8);
+make_literal!(u8, UInt8);
+make_literal!(i16, Int16);
+make_literal!(u16, UInt16);
+make_literal!(i32, Int32);
+make_literal!(u32, UInt32);
+make_literal!(i64, Int64);
+make_literal!(u64, UInt64);
+make_literal!(f64, Float64);
+
+impl FromLiteral for usize {
+    fn try_from_literal(lit: &LiteralValue) -> DaftResult<Self> {
+        match lit {
+            LiteralValue::Int8(i8) => cast(*i8),
+            LiteralValue::UInt8(u8) => cast(*u8),
+            LiteralValue::Int16(i16) => cast(*i16),
+            LiteralValue::UInt16(u16) => cast(*u16),
+            LiteralValue::Int32(i32) => cast(*i32),
+            LiteralValue::UInt32(u32) => cast(*u32),
+            LiteralValue::Int64(i64) => cast(*i64),
+            LiteralValue::UInt64(u64) => cast(*u64),
+            _ => None,
+        }
+        .ok_or_else(|| DaftError::ValueError("Unsupported literal type".to_string()))
+    }
+}
+
+/// Marker trait to allowlist what can be converted to a literal via serde
+trait SerializableLiteral: serde::Serialize {}
+
+/// Marker trait to allowlist what can be converted from a literal via serde
+trait DeserializableLiteral: DeserializeOwned {}
+
+impl SerializableLiteral for IOConfig {}
+impl DeserializableLiteral for IOConfig {}
+impl SerializableLiteral for ImageMode {}
+impl DeserializableLiteral for ImageMode {}
+impl SerializableLiteral for ImageFormat {}
+impl DeserializableLiteral for ImageFormat {}
+impl SerializableLiteral for CountMode {}
+impl DeserializableLiteral for CountMode {}
+
+impl<D> FromLiteral for D
+where
+    D: DeserializableLiteral,
+{
+    fn try_from_literal(lit: &LiteralValue) -> DaftResult<Self> {
+        let deserializer = deserializer::LiteralValueDeserializer { lit };
+        D::deserialize(deserializer).map_err(|e| e.into())
+    }
+}
+
+impl<S> Literal for S
+where
+    S: SerializableLiteral,
+{
+    fn literal_value(self) -> LiteralValue {
+        self.serialize(serializer::LiteralValueSerializer)
+            .expect("serialization failed")
+    }
+}