feat: Base64 Encoding (#5158)

Author: srilman

.github/workflows/pr-test-suite.yml

@@ -1254,7 +1254,7 @@ jobs:
     - name: Python And Rust Style Check
       run: |
         source .venv/bin/activate
-        pre-commit run --all-files
+        pre-commit run --all-files -v
 
     - name: Send Slack notification on failure
       uses: slackapi/[email protected]

Cargo.lock

@@ -71,8 +71,8 @@ test: .venv build  ## Run tests
 	HYPOTHESIS_MAX_EXAMPLES=$(HYPOTHESIS_MAX_EXAMPLES) $(VENV_BIN)/pytest --hypothesis-seed=$(HYPOTHESIS_SEED) --ignore tests/integration $(EXTRA_ARGS)
 
 .PHONY: doctests
-doctests:
-	DAFT_BOLD_TABLE_HEADERS=0 pytest --doctest-modules --continue-on-collection-errors --ignore=daft/functions/llm.py daft/dataframe/dataframe.py daft/expressions/expressions.py daft/convert.py daft/udf/__init__.py daft/functions/ daft/datatype.py
+doctests: .venv
+	DAFT_BOLD_TABLE_HEADERS=0 DAFT_PROGRESS_BAR=0 $(VENV_BIN)/pytest --doctest-modules --continue-on-collection-errors --ignore=daft/functions/llm.py daft/dataframe/dataframe.py daft/expressions/expressions.py daft/convert.py daft/udf/__init__.py daft/functions/ daft/datatype.py
 
 .PHONY: dsdgen
 dsdgen: .venv ## Generate TPC-DS data

Makefile

@@ -47,7 +47,7 @@
     from daft.udf.legacy import BoundUDFArgs, InitArgsType, UninitializedUdf
     from daft.window import Window
 
-    EncodingCodec = Literal["deflate", "gzip", "gz", "utf-8", "utf8" "zlib"]
+    EncodingCodec = Literal["base64", "deflate", "gzip", "gz", "utf-8", "utf8", "zlib"]
 
 
 def lit(value: object) -> Expression:
@@ -1595,7 +1595,7 @@ def encode(self, codec: EncodingCodec) -> Expression:
         r"""Encodes the expression (binary strings) using the specified codec.
 
         Args:
-            codec (str): encoding codec (deflate, gzip, zlib)
+            codec (str): encoding codec (base64, deflate, gzip, zlib)
 
         Returns:
             Expression: A new expression, of type `binary`, with the encoded value.
@@ -1641,7 +1641,7 @@ def decode(self, codec: EncodingCodec) -> Expression:
         """Decodes the expression (binary strings) using the specified codec.
 
         Args:
-            codec (str): decoding codec (deflate, gzip, zlib)
+            codec (str): decoding codec (base64, deflate, gzip, zlib)
 
         Returns:
             Expression: A new expression with the decoded values.
@@ -1651,6 +1651,20 @@ def decode(self, codec: EncodingCodec) -> Expression:
             only decoding with 'utf-8' returns a string.
 
         Examples:
+            >>> import daft
+            >>> from daft import col
+            >>> df = daft.from_pydict({"bytes": [b"aGVsbG8sIHdvcmxkIQ=="]})
+            >>> df.select(col("bytes").decode("base64")).show()
+            ╭──────────────────╮
+            │ bytes            │
+            │ ---              │
+            │ Binary           │
+            ╞══════════════════╡
+            │ b"hello, world!" │
+            ╰──────────────────╯
+            <BLANKLINE>
+            (Showing first 1 of 1 rows)
+
             >>> import daft
             >>> import zlib
             >>> from daft import col

daft/expressions/expressions.py

@@ -1,5 +1,6 @@
 [dependencies]
 arrow2 = {workspace = true}
+base64 = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 common-macros = {path = "../common/macros"}
 daft-core = {path = "../daft-core", default-features = false}

src/daft-functions-binary/Cargo.toml

@@ -11,6 +11,7 @@ use simdutf8::basic::from_utf8;
 /// Supported codecs for the decode and encode functions.
 #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub enum Codec {
+    Base64,
     Deflate,
     Gzip,
     Utf8,
@@ -20,6 +21,7 @@ pub enum Codec {
 impl Display for Codec {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.write_str(match self {
+            Self::Base64 => "base64",
             Self::Deflate => "deflate",
             Self::Gzip => "gzip",
             Self::Utf8 => "utf8",
@@ -61,6 +63,7 @@ pub(crate) type Transform = fn(input: &[u8]) -> DaftResult<Vec<u8>>;
 impl Codec {
     pub(crate) fn encoder(&self) -> Transform {
         match self {
+            Self::Base64 => base64_encoder,
             Self::Deflate => deflate_encoder,
             Self::Gzip => gzip_encoder,
             Self::Utf8 => utf8_encoder,
@@ -70,6 +73,7 @@ impl Codec {
 
     pub(crate) fn decoder(&self) -> Transform {
         match self {
+            Self::Base64 => base64_decoder,
             Self::Deflate => deflate_decoder,
             Self::Gzip => gzip_decoder,
             Self::Utf8 => utf8_decoder,
@@ -79,6 +83,7 @@ impl Codec {
 
     pub(crate) fn kind(&self) -> CodecKind {
         match self {
+            Self::Base64 => CodecKind::Binary,
             Self::Deflate => CodecKind::Binary,
             Self::Gzip => CodecKind::Binary,
             Self::Utf8 => CodecKind::Text,
@@ -99,6 +104,7 @@ impl FromStr for Codec {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s.to_lowercase().as_str() {
+            "base64" => Ok(Self::Base64),
             "deflate" => Ok(Self::Deflate),
             "gzip" | "gz" => Ok(Self::Gzip),
             "zlib" => Ok(Self::Zlib),
@@ -115,6 +121,13 @@ impl FromStr for Codec {
 // ENCODERS
 //
 
+#[inline]
+fn base64_encoder(input: &[u8]) -> DaftResult<Vec<u8>> {
+    use base64::{Engine, engine::general_purpose::STANDARD};
+
+    Ok(STANDARD.encode(input).into_bytes())
+}
+
 #[inline]
 fn deflate_encoder(input: &[u8]) -> DaftResult<Vec<u8>> {
     use std::io::Write;
@@ -160,6 +173,14 @@ fn zlib_encoder(input: &[u8]) -> DaftResult<Vec<u8>> {
 // DECODERS
 //
 
+#[inline]
+fn base64_decoder(input: &[u8]) -> DaftResult<Vec<u8>> {
+    use base64::{Engine, engine::general_purpose::STANDARD};
+    STANDARD
+        .decode(input)
+        .map_err(|e| DaftError::ValueError(format!("Invalid base64 input: {}", e)))
+}
+
 #[inline]
 fn deflate_decoder(input: &[u8]) -> DaftResult<Vec<u8>> {
     use std::io::Read;
@@ -223,6 +244,8 @@ mod tests {
         assert_eq!("zlib".parse::<Codec>().unwrap(), Codec::Zlib);
         assert_eq!("ZLIB".parse::<Codec>().unwrap(), Codec::Zlib);
         assert_eq!("ZlIb".parse::<Codec>().unwrap(), Codec::Zlib);
+        assert_eq!("base64".parse::<Codec>().unwrap(), Codec::Base64);
+        assert_eq!("BASE64".parse::<Codec>().unwrap(), Codec::Base64);
         assert!("unknown".parse::<Codec>().is_err());
     }
 }

src/daft-functions-binary/src/codecs.rs

@@ -32,7 +32,7 @@ impl ScalarUDF for BinaryEncode {
         inputs: FunctionArgs<ExprRef>,
         schema: &Schema,
     ) -> DaftResult<Field> {
-        let Args { input, codec: _ } = inputs.try_into()?;
+        let Args { input, codec } = inputs.try_into()?;
         let input = input.to_field(schema)?;
 
         ensure!(
@@ -44,7 +44,7 @@ impl ScalarUDF for BinaryEncode {
             input.dtype
         );
 
-        Ok(Field::new(input.name, DataType::Binary))
+        Ok(Field::new(input.name, codec.returns()))
     }
 
     fn call(&self, inputs: daft_dsl::functions::FunctionArgs<Series>) -> DaftResult<Series> {

src/daft-functions-binary/src/encode.rs

@@ -105,8 +105,9 @@ def test_codec_zlib():
 
 
 def test_codec_base64():
-    with pytest.raises(Exception, match="unsupported codec"):
-        _test_codec("base64", None)
+    import base64
+
+    _test_codec("base64", buff=base64.b64encode(TEXT))
 
 
 def test_codec_zstd():