feat: Add support for Binary column types (#146)

Co-authored-by: Oliver Borchert <[email protected]>

Author: frewsxcv

dataframely/__init__.py

@@ -23,6 +23,7 @@
 from .columns import (
     Any,
     Array,
+    Binary,
     Bool,
     Categorical,
     Column,
@@ -77,6 +78,7 @@
     "read_parquet_metadata_schema",
     "read_parquet_metadata_collection",
     "Any",
+    "Binary",
     "Bool",
     "Categorical",
     "Column",

dataframely/columns/__init__.py

@@ -5,6 +5,7 @@
 from ._registry import column_from_dict
 from .any import Any
 from .array import Array
+from .binary import Binary
 from .bool import Bool
 from .categorical import Categorical
 from .datetime import Date, Datetime, Duration, Time
@@ -22,6 +23,7 @@
     "column_from_dict",
     "Any",
     "Array",
+    "Binary",
     "Bool",
     "Categorical",
     "Date",

dataframely/columns/binary.py

@@ -0,0 +1,38 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+from __future__ import annotations
+
+import polars as pl
+
+from dataframely._compat import pa, sa, sa_TypeEngine
+from dataframely.random import Generator
+
+from ._base import Column
+from ._registry import register
+
+
+@register
+class Binary(Column):
+    """A column of binary values."""
+
+    @property
+    def dtype(self) -> pl.DataType:
+        return pl.Binary()
+
+    def sqlalchemy_dtype(self, dialect: sa.Dialect) -> sa_TypeEngine:
+        if dialect.name == "mssql":
+            return sa.VARBINARY()
+        return sa.LargeBinary()
+
+    @property
+    def pyarrow_dtype(self) -> pa.DataType:
+        return pa.large_binary()
+
+    def _sample_unchecked(self, generator: Generator, n: int) -> pl.Series:
+        return generator.sample_binary(
+            n,
+            min_bytes=0,
+            max_bytes=32,
+            null_probability=self._null_probability,
+        )

dataframely/random.py

@@ -148,6 +148,31 @@ def sample_string(
             pl.Series(samples, dtype=pl.String), null_probability
         )
 
+    def sample_binary(
+        self,
+        n: int = 1,
+        *,
+        min_bytes: int,
+        max_bytes: int,
+        null_probability: float = 0.0,
+    ) -> pl.Series:
+        """Sample a list of binary values in the specified length range.
+
+        Args:
+            n: The number of binary values to sample.
+            min_bytes: The minimum number of bytes for each value.
+            max_bytes: The maximum number of bytes for each value.
+            null_probability: The probability of an element being ``null``.
+
+        Returns:
+            A series with ``n`` elements of dtype ``Binary``.
+        """
+        lengths = self.numpy_generator.integers(min_bytes, max_bytes + 1, size=n)
+        samples = [self.numpy_generator.bytes(length) for length in lengths]
+        return self._apply_null_mask(
+            pl.Series(samples, dtype=pl.Binary), null_probability
+        )
+
     def sample_choice(
         self,
         n: int = 1,

dataframely/testing/const.py

@@ -22,6 +22,7 @@
     dc.UInt64,
     dc.String,
     dc.Categorical,
+    dc.Binary,
 ]
 INTEGER_COLUMN_TYPES: list[type[dc.Column]] = [
     dc.Integer,

tests/column_types/test_binary.py

@@ -0,0 +1,24 @@
+# Copyright (c) QuantCo 2025-2025
+# SPDX-License-Identifier: BSD-3-Clause
+
+import polars as pl
+import pytest
+
+import dataframely as dy
+from dataframely.columns import Column
+
+
+class BinarySchema(dy.Schema):
+    a = dy.Binary()
+
+
+@pytest.mark.parametrize(
+    ("column", "dtype", "is_valid"),
+    [
+        (dy.Binary(), pl.Binary(), True),
+        (dy.Binary(), pl.String(), False),
+        (dy.Binary(), pl.Null(), False),
+    ],
+)
+def test_validate_dtype(column: Column, dtype: pl.DataType, is_valid: bool) -> None:
+    assert column.validate_dtype(dtype) == is_valid

tests/columns/test_sql_schema.py

@@ -15,6 +15,7 @@
     ("column", "datatype"),
     [
         (dy.Any(), "SQL_VARIANT"),
+        (dy.Binary(), "VARBINARY(max)"),
         (dy.Bool(), "BIT"),
         (dy.Date(), "DATE"),
         (dy.Datetime(), "DATETIME2(6)"),
@@ -61,6 +62,7 @@ def test_mssql_datatype(column: Column, datatype: str) -> None:
 @pytest.mark.parametrize(
     ("column", "datatype"),
     [
+        (dy.Binary(), "BYTEA"),
         (dy.Bool(), "BOOLEAN"),
         (dy.Date(), "DATE"),
         (dy.Datetime(), "TIMESTAMP WITHOUT TIME ZONE"),

tests/test_random.py

@@ -40,6 +40,7 @@ def test_seeding_nonconstant() -> None:
         lambda generator, n: generator.sample_float(n, min=0, max=5),
         lambda generator, n: generator.sample_string(n, regex="[abc]"),
         lambda generator, n: generator.sample_choice(n, choices=[1, 2, 3]),
+        lambda generator, n: generator.sample_binary(n, min_bytes=1, max_bytes=10),
         lambda generator, n: generator.sample_time(n, min=dt.time(0, 0), max=None),
         lambda generator, n: generator.sample_date(
             n, min=dt.date(1970, 1, 1), max=None
@@ -75,6 +76,9 @@ def test_sample_correct_n(
         lambda generator, n, prob: generator.sample_choice(
             n, choices=[1, 2, 3], null_probability=prob
         ),
+        lambda generator, n, prob: generator.sample_binary(
+            n, min_bytes=1, max_bytes=10, null_probability=prob
+        ),
         lambda generator, n, prob: generator.sample_time(
             n, min=dt.time(0, 0), max=None, null_probability=prob
         ),
@@ -131,6 +135,16 @@ def test_sample_string(generator: Generator) -> None:
     assert (samples.str.len_bytes() == 2).all()
 
 
+def test_sample_binary(generator: Generator) -> None:
+    samples = generator.sample_binary(100, min_bytes=1, max_bytes=10)
+    assert (
+        samples.to_frame("s").select(pl.col("s").bin.size("b") >= 1).to_series().all()
+    )
+    assert (
+        samples.to_frame("s").select(pl.col("s").bin.size("b") <= 10).to_series().all()
+    )
+
+
 def test_sample_choice(generator: Generator) -> None:
     samples = generator.sample_choice(100_000, choices=[1, 2, 3])
     assert np.allclose(