duckdb 1.5.0.dev32__cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl

duckdb/experimental/spark/sql/udf.py

@@ -0,0 +1,37 @@
+# https://sparkbyexamples.com/pyspark/pyspark-udf-user-defined-function/
+from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union
+
+from .types import DataType
+
+if TYPE_CHECKING:
+    from .session import SparkSession
+
+DataTypeOrString = Union[DataType, str]
+UserDefinedFunctionLike = TypeVar("UserDefinedFunctionLike")
+
+
+class UDFRegistration:
+    def __init__(self, sparkSession: "SparkSession"):
+        self.sparkSession = sparkSession
+
+    def register(
+        self,
+        name: str,
+        f: Union[Callable[..., Any], "UserDefinedFunctionLike"],
+        returnType: Optional["DataTypeOrString"] = None,
+    ) -> "UserDefinedFunctionLike":
+        self.sparkSession.conn.create_function(name, f, return_type=returnType)
+
+    def registerJavaFunction(
+        self,
+        name: str,
+        javaClassName: str,
+        returnType: Optional["DataTypeOrString"] = None,
+    ) -> None:
+        raise NotImplementedError
+
+    def registerJavaUDAF(self, name: str, javaClassName: str) -> None:
+        raise NotImplementedError
+
+
+__all__ = ["UDFRegistration"]

duckdb/filesystem.py

@@ -0,0 +1,23 @@
+from fsspec import filesystem, AbstractFileSystem
+from fsspec.implementations.memory import MemoryFileSystem, MemoryFile
+from .bytes_io_wrapper import BytesIOWrapper
+from io import TextIOBase
+
+def is_file_like(obj):
+    # We only care that we can read from the file
+    return hasattr(obj, "read") and hasattr(obj, "seek")
+
+
+class ModifiedMemoryFileSystem(MemoryFileSystem):
+    protocol = ('DUCKDB_INTERNAL_OBJECTSTORE',)
+    # defer to the original implementation that doesn't hardcode the protocol
+    _strip_protocol = classmethod(AbstractFileSystem._strip_protocol.__func__)
+
+    def add_file(self, object, path):
+        if not is_file_like(object):
+            raise ValueError("Can not read from a non file-like object")
+        path = self._strip_protocol(path)
+        if isinstance(object, TextIOBase):
+            # Wrap this so that we can return a bytes object from 'read'
+            object = BytesIOWrapper(object)
+        self.store[path] = MemoryFile(self, path, object.read())

duckdb/functional/__init__.py

@@ -0,0 +1,17 @@
+from _duckdb.functional import (
+	FunctionNullHandling,
+	PythonUDFType,
+	SPECIAL,
+	DEFAULT,
+	NATIVE,
+	ARROW
+)
+
+__all__ = [
+	"FunctionNullHandling",
+	"PythonUDFType",
+	"SPECIAL",
+	"DEFAULT",
+	"NATIVE",
+	"ARROW"
+]

duckdb/functional/__init__.pyi

@@ -0,0 +1,31 @@
+from typing import Dict
+
+SPECIAL: FunctionNullHandling
+DEFAULT: FunctionNullHandling
+
+NATIVE: PythonUDFType
+ARROW: PythonUDFType
+
+class FunctionNullHandling:
+    DEFAULT: FunctionNullHandling
+    SPECIAL: FunctionNullHandling
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
+    @property
+    def __members__(self) -> Dict[str, FunctionNullHandling]: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class PythonUDFType:
+    NATIVE: PythonUDFType
+    ARROW: PythonUDFType
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
+    @property
+    def __members__(self) -> Dict[str, PythonUDFType]: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...

duckdb/polars_io.py

@@ -0,0 +1,237 @@
+import duckdb
+import polars as pl
+from typing import Iterator, Optional
+
+from polars.io.plugins import register_io_source
+from duckdb import SQLExpression
+import json
+from decimal import Decimal
+import datetime
+
+def _predicate_to_expression(predicate: pl.Expr) -> Optional[SQLExpression]:
+    """
+    Convert a Polars predicate expression to a DuckDB-compatible SQL expression.
+    
+    Parameters:
+        predicate (pl.Expr): A Polars expression (e.g., col("foo") > 5)
+    
+    Returns:
+        SQLExpression: A DuckDB SQL expression string equivalent.
+        None: If conversion fails.
+
+    Example:
+        >>> _predicate_to_expression(pl.col("foo") > 5)
+        SQLExpression("(foo > 5)")
+    """
+    # Serialize the Polars expression tree to JSON
+    tree = json.loads(predicate.meta.serialize(format="json"))
+    
+    try:
+        # Convert the tree to SQL
+        sql_filter = _pl_tree_to_sql(tree)
+        return SQLExpression(sql_filter)
+    except:
+        # If the conversion fails, we return None
+        return None
+
+
+def _pl_operation_to_sql(op: str) -> str:
+    """
+    Map Polars binary operation strings to SQL equivalents.
+    
+    Example:
+        >>> _pl_operation_to_sql("Eq")
+        '='
+    """
+    try:
+        return {
+            "Lt": "<",
+            "LtEq": "<=",
+            "Gt": ">",
+            "GtEq": ">=",
+            "Eq": "=",
+            "Modulus": "%",
+            "And": "AND",
+            "Or": "OR",
+        }[op]
+    except KeyError:
+        raise NotImplementedError(op)
+
+
+def _escape_sql_identifier(identifier: str) -> str:
+    """
+    Escape SQL identifiers by doubling any double quotes and wrapping in double quotes.
+
+    Example:
+        >>> _escape_sql_identifier('column"name')
+        '"column""name"'
+    """
+    escaped = identifier.replace('"', '""')
+    return f'"{escaped}"'
+
+
+def _pl_tree_to_sql(tree: dict) -> str:
+    """
+    Recursively convert a Polars expression tree (as JSON) to a SQL string.
+    
+    Parameters:
+        tree (dict): JSON-deserialized expression tree from Polars
+    
+    Returns:
+        str: SQL expression string
+    
+    Example:
+        Input tree:
+        {
+            "BinaryExpr": {
+                "left": { "Column": "foo" },
+                "op": "Gt",
+                "right": { "Literal": { "Int": 5 } }
+            }
+        }
+        Output: "(foo > 5)"
+    """
+    [node_type] = tree.keys()
+    subtree = tree[node_type]
+
+    if node_type == "BinaryExpr":
+        # Binary expressions: left OP right
+        return (
+                "(" +
+                " ".join((
+                    _pl_tree_to_sql(subtree['left']),
+                    _pl_operation_to_sql(subtree['op']),
+                    _pl_tree_to_sql(subtree['right'])
+                )) +
+                ")"
+        )
+    if node_type == "Column":
+        # A reference to a column name
+        # Wrap in quotes to handle special characters
+        return _escape_sql_identifier(subtree)
+
+    if node_type in ("Literal", "Dyn"):
+        # Recursively process dynamic or literal values
+        return _pl_tree_to_sql(subtree)
+
+    if node_type == "Int":
+        # Direct integer literals
+        return str(subtree)
+
+    if node_type == "Function":
+        # Handle boolean functions like IsNull, IsNotNull
+        inputs = subtree["input"]
+        func_dict = subtree["function"]
+
+        if "Boolean" in func_dict:
+            func = func_dict["Boolean"]
+            arg_sql = _pl_tree_to_sql(inputs[0])
+
+            if func == "IsNull":
+                return f"({arg_sql} IS NULL)"
+            if func == "IsNotNull":
+                return f"({arg_sql} IS NOT NULL)"
+            raise NotImplementedError(f"Boolean function not supported: {func}")
+
+        raise NotImplementedError(f"Unsupported function type: {func_dict}")
+
+    if node_type == "Scalar":
+        # Detect format: old style (dtype/value) or new style (direct type key)
+        if "dtype" in subtree and "value" in subtree:
+            dtype = str(subtree["dtype"])
+            value = subtree["value"]
+        else:
+            # New style: dtype is the single key in the dict
+            dtype = next(iter(subtree.keys()))
+            value = subtree
+
+        # Decimal support
+        if dtype.startswith("{'Decimal'") or dtype == "Decimal":
+            decimal_value = value['Decimal']
+            decimal_value = Decimal(decimal_value[0]) / Decimal(10 ** decimal_value[1])
+            return str(decimal_value)
+
+        # Datetime with microseconds since epoch
+        if dtype.startswith("{'Datetime'") or dtype == "Datetime":
+            micros = value['Datetime'][0]
+            dt_timestamp = datetime.datetime.fromtimestamp(micros / 1_000_000, tz=datetime.UTC)
+            return f"'{str(dt_timestamp)}'::TIMESTAMP"
+
+        # Match simple numeric/boolean types
+        if dtype in ("Int8", "Int16", "Int32", "Int64",
+                     "UInt8", "UInt16", "UInt32", "UInt64",
+                     "Float32", "Float64", "Boolean"):
+            return str(value[dtype])
+
+        # Time type
+        if dtype == "Time":
+            nanoseconds = value["Time"]
+            seconds = nanoseconds // 1_000_000_000
+            microseconds = (nanoseconds % 1_000_000_000) // 1_000
+            dt_time = (datetime.datetime.min + datetime.timedelta(
+                seconds=seconds, microseconds=microseconds
+            )).time()
+            return f"'{dt_time}'::TIME"
+
+        # Date type
+        if dtype == "Date":
+            days_since_epoch = value["Date"]
+            date = datetime.date(1970, 1, 1) + datetime.timedelta(days=days_since_epoch)
+            return f"'{date}'::DATE"
+
+        # Binary type
+        if dtype == "Binary":
+            binary_data = bytes(value["Binary"])
+            escaped = ''.join(f'\\x{b:02x}' for b in binary_data)
+            return f"'{escaped}'::BLOB"
+
+        # String type
+        if dtype == "String" or dtype == "StringOwned":
+            # Some new formats may store directly under StringOwned
+            string_val = value.get("StringOwned", value.get("String", None))
+            return f"'{string_val}'"
+
+
+        raise NotImplementedError(f"Unsupported scalar type {str(dtype)}, with value {value}")
+
+    raise NotImplementedError(f"Node type: {node_type} is not implemented. {subtree}")
+
+def duckdb_source(relation: duckdb.DuckDBPyRelation, schema: pl.schema.Schema) -> pl.LazyFrame:
+    """
+    A polars IO plugin for DuckDB.
+    """
+    def source_generator(
+        with_columns: Optional[list[str]],
+        predicate: Optional[pl.Expr],
+        n_rows: Optional[int],
+        batch_size: Optional[int],
+    ) -> Iterator[pl.DataFrame]:
+        duck_predicate = None
+        relation_final = relation
+        if with_columns is not None:
+            cols = ",".join(map(_escape_sql_identifier, with_columns))
+            relation_final = relation_final.project(cols)
+        if n_rows is not None:
+            relation_final = relation_final.limit(n_rows)
+        if predicate is not None:
+            # We have a predicate, if possible, we push it down to DuckDB
+            duck_predicate = _predicate_to_expression(predicate)
+        # Try to pushdown filter, if one exists
+        if duck_predicate is not None:
+            relation_final = relation_final.filter(duck_predicate)
+        if batch_size is None:
+            results = relation_final.fetch_arrow_reader()
+        else:
+            results = relation_final.fetch_arrow_reader(batch_size)
+        while True:
+            try:
+                record_batch = results.read_next_batch()
+                if predicate is not None and duck_predicate is None:
+                    # We have a predicate, but did not manage to push it down, we fallback here
+                    yield pl.from_arrow(record_batch).filter(predicate)
+                else:
+                    yield pl.from_arrow(record_batch)
+            except StopIteration:
+                break
+
+    return register_io_source(source_generator, schema=schema)