_substrait.pyx

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# cython: language_level = 3
from cython.operator cimport dereference as deref
from libcpp.vector cimport vector as std_vector

from pyarrow import Buffer, py_buffer
from pyarrow._compute cimport Expression
from pyarrow.lib import frombytes, tobytes
from pyarrow.lib cimport *
from pyarrow.includes.libarrow cimport *
from pyarrow.includes.libarrow_substrait cimport *

try:
    import substrait as py_substrait
except ImportError:
    py_substrait = None
else:
    import substrait.proto  # no-cython-lint


# TODO GH-37235: Fix exception handling
cdef CDeclaration _create_named_table_provider(
    dict named_args, const std_vector[c_string]& names, const CSchema& schema
) noexcept:
    cdef:
        c_string c_name
        shared_ptr[CTable] c_in_table
        shared_ptr[CTableSourceNodeOptions] c_tablesourceopts
        shared_ptr[CExecNodeOptions] c_input_node_opts
        vector[CDeclaration.Input] no_c_inputs

    py_names = []
    for i in range(names.size()):
        c_name = names[i]
        py_names.append(frombytes(c_name))
    py_schema = pyarrow_wrap_schema(make_shared[CSchema](schema))

    py_table = named_args["provider"](py_names, py_schema)
    c_in_table = pyarrow_unwrap_table(py_table)
    c_tablesourceopts = make_shared[CTableSourceNodeOptions](c_in_table)
    c_input_node_opts = static_pointer_cast[CExecNodeOptions, CTableSourceNodeOptions](
        c_tablesourceopts)
    return CDeclaration(tobytes("table_source"),
                        no_c_inputs, c_input_node_opts)


def run_query(plan, *, table_provider=None, use_threads=True):
    """
    Execute a Substrait plan and read the results as a RecordBatchReader.

    Parameters
    ----------
    plan : Union[Buffer, bytes]
        The serialized Substrait plan to execute.
    table_provider : object (optional)
        A function to resolve any NamedTable relation to a table.
        The function will receive two arguments which will be a list
        of strings representing the table name and a pyarrow.Schema representing
        the expected schema and should return a pyarrow.Table.
    use_threads : bool, default True
        If True then multiple threads will be used to run the query.  If False then
        all CPU intensive work will be done on the calling thread.

    Returns
    -------
    RecordBatchReader
        A reader containing the result of the executed query

    Examples
    --------
    >>> import pyarrow as pa
    >>> from pyarrow.lib import tobytes
    >>> import pyarrow.substrait as substrait
    >>> test_table_1 = pa.Table.from_pydict({"x": [1, 2, 3]})
    >>> test_table_2 = pa.Table.from_pydict({"x": [4, 5, 6]})
    >>> def table_provider(names, schema):
    ...     if not names:
    ...        raise Exception("No names provided")
    ...     elif names[0] == "t1":
    ...        return test_table_1
    ...     elif names[1] == "t2":
    ...        return test_table_2
    ...     else:
    ...        raise Exception("Unrecognized table name")
    ...
    >>> substrait_query = '''
    ...         {
    ...             "relations": [
    ...             {"rel": {
    ...                 "read": {
    ...                 "base_schema": {
    ...                     "struct": {
    ...                     "types": [
    ...                                 {"i64": {}}
    ...                             ]
    ...                     },
    ...                     "names": [
    ...                             "x"
    ...                             ]
    ...                 },
    ...                 "namedTable": {
    ...                         "names": ["t1"]
    ...                 }
    ...                 }
    ...             }}
    ...             ]
    ...         }
    ... '''
    >>> buf = pa._substrait._parse_json_plan(tobytes(substrait_query))
    >>> reader = pa.substrait.run_query(buf, table_provider=table_provider)
    >>> reader.read_all()
    pyarrow.Table
    x: int64
    ----
    x: [[1,2,3]]
    """

    cdef:
        CResult[shared_ptr[CRecordBatchReader]] c_res_reader
        shared_ptr[CRecordBatchReader] c_reader
        RecordBatchReader reader
        shared_ptr[CBuffer] c_buf_plan
        CConversionOptions c_conversion_options
        c_bool c_use_threads

    c_use_threads = use_threads
    if isinstance(plan, (bytes, memoryview)):
        c_buf_plan = pyarrow_unwrap_buffer(py_buffer(plan))
    elif isinstance(plan, Buffer):
        c_buf_plan = pyarrow_unwrap_buffer(plan)
    else:
        raise TypeError(
            f"Expected 'pyarrow.Buffer' or bytes, got '{type(plan)}'")

    if table_provider is not None:
        named_table_args = {
            "provider": table_provider
        }
        c_conversion_options.named_table_provider = BindFunction[CNamedTableProvider](
            &_create_named_table_provider, named_table_args)

    with nogil:
        c_res_reader = ExecuteSerializedPlan(
            deref(c_buf_plan), default_extension_id_registry(),
            GetFunctionRegistry(), c_conversion_options, c_use_threads)

    c_reader = GetResultValue(c_res_reader)

    reader = RecordBatchReader.__new__(RecordBatchReader)
    reader.reader = c_reader
    return reader


def _parse_json_plan(plan):
    """
    Parse a JSON plan into equivalent serialized Protobuf.

    Parameters
    ----------
    plan : bytes
        Substrait plan in JSON.

    Returns
    -------
    Buffer
        A buffer containing the serialized Protobuf plan.
    """

    cdef:
        CResult[shared_ptr[CBuffer]] c_res_buffer
        c_string c_str_plan
        shared_ptr[CBuffer] c_buf_plan

    c_str_plan = plan
    c_res_buffer = SerializeJsonPlan(c_str_plan)
    with nogil:
        c_buf_plan = GetResultValue(c_res_buffer)
    return pyarrow_wrap_buffer(c_buf_plan)


class SubstraitSchema:
    """A Schema encoded for Substrait usage.

    The SubstraitSchema contains a schema represented
    both as a substrait ``NamedStruct`` and as an
    ``ExtendedExpression``.

    The ``ExtendedExpression`` is available for cases where types
    used by the schema require extensions to decode them.
    In such case the schema will be the ``base_schema`` of the
    ``ExtendedExpression`` and all extensions will be provided.
    """

    def __init__(self, schema, expression):
        self.schema = schema
        self.expression = expression

    def to_pysubstrait(self):
        """Convert the schema to a substrait-python ExtendedExpression object."""
        if py_substrait is None:
            raise ImportError("The 'substrait' package is required.")
        return py_substrait.proto.ExtendedExpression.FromString(self.expression)


def serialize_schema(schema):
    """
    Serialize a schema into a SubstraitSchema object.

    Parameters
    ----------
    schema : Schema
        The schema to serialize

    Returns
    -------
    SubstraitSchema
        The schema stored in a SubstraitSchema object.
    """
    return SubstraitSchema(
        schema=_serialize_namedstruct_schema(schema),
        expression=serialize_expressions([], [], schema, allow_arrow_extensions=True)
    )


def _serialize_namedstruct_schema(schema):
    cdef:
        CResult[shared_ptr[CBuffer]] c_res_buffer
        shared_ptr[CBuffer] c_buffer
        CConversionOptions c_conversion_options
        CExtensionSet c_extensions

    with nogil:
        c_res_buffer = SerializeSchema(deref((<Schema> schema).sp_schema), &c_extensions, c_conversion_options)
        c_buffer = GetResultValue(c_res_buffer)

    return memoryview(pyarrow_wrap_buffer(c_buffer))


def deserialize_schema(buf):
    """
    Deserialize a ``NamedStruct`` Substrait message
    or a SubstraitSchema object into an Arrow Schema object

    Parameters
    ----------
    buf : Buffer or bytes or SubstraitSchema
        The message to deserialize

    Returns
    -------
    Schema
        The deserialized schema
    """
    cdef:
        shared_ptr[CBuffer] c_buffer
        CResult[shared_ptr[CSchema]] c_res_schema
        shared_ptr[CSchema] c_schema
        CConversionOptions c_conversion_options
        CExtensionSet c_extensions

    if isinstance(buf, SubstraitSchema):
        return deserialize_expressions(buf.expression).schema

    if isinstance(buf, (bytes, memoryview)):
        c_buffer = pyarrow_unwrap_buffer(py_buffer(buf))
    elif isinstance(buf, Buffer):
        c_buffer = pyarrow_unwrap_buffer(buf)
    else:
        raise TypeError(
            f"Expected 'pyarrow.Buffer' or bytes, got '{type(buf)}'")

    with nogil:
        c_res_schema = DeserializeSchema(
            deref(c_buffer), c_extensions, c_conversion_options)
        c_schema = GetResultValue(c_res_schema)

    return pyarrow_wrap_schema(c_schema)


def serialize_expressions(exprs, names, schema, *, allow_arrow_extensions=False):
    """
    Serialize a collection of expressions into Substrait

    Substrait expressions must be bound to a schema.  For example,
    the Substrait expression ``a:i32 + b:i32`` is different from the
    Substrait expression ``a:i64 + b:i64``.  Pyarrow expressions are
    typically unbound.  For example, both of the above expressions
    would be represented as ``a + b`` in pyarrow.

    This means a schema must be provided when serializing an expression.
    It also means that the serialization may fail if a matching function
    call cannot be found for the expression.

    Parameters
    ----------
    exprs : list of Expression
        The expressions to serialize
    names : list of str
        Names for the expressions
    schema : Schema
        The schema the expressions will be bound to
    allow_arrow_extensions : bool, default False
        If False then only functions that are part of the core Substrait function
        definitions will be allowed.  Set this to True to allow pyarrow-specific functions
        and user defined functions but the result may not be accepted by other
        compute libraries.

    Returns
    -------
    Buffer
        An ExtendedExpression message containing the serialized expressions
    """
    cdef:
        CResult[shared_ptr[CBuffer]] c_res_buffer
        shared_ptr[CBuffer] c_buffer
        CNamedExpression c_named_expr
        CBoundExpressions c_bound_exprs
        CConversionOptions c_conversion_options

    if len(exprs) != len(names):
        raise ValueError("exprs and names need to have the same length")
    for expr, name in zip(exprs, names):
        if not isinstance(expr, Expression):
            raise TypeError(f"Expected Expression, got '{type(expr)}' in exprs")
        if not isinstance(name, str):
            raise TypeError(f"Expected str, got '{type(name)}' in names")
        c_named_expr.expression = (<Expression> expr).unwrap()
        c_named_expr.name = tobytes(<str> name)
        c_bound_exprs.named_expressions.push_back(c_named_expr)

    c_bound_exprs.schema = (<Schema> schema).sp_schema

    c_conversion_options.allow_arrow_extensions = allow_arrow_extensions

    with nogil:
        c_res_buffer = SerializeExpressions(c_bound_exprs, c_conversion_options)
        c_buffer = GetResultValue(c_res_buffer)
    return memoryview(pyarrow_wrap_buffer(c_buffer))


cdef class BoundExpressions(_Weakrefable):
    """
    A collection of named expressions and the schema they are bound to

    This is equivalent to the Substrait ExtendedExpression message
    """

    cdef:
        CBoundExpressions c_bound_exprs

    def __init__(self):
        msg = 'BoundExpressions is an abstract class thus cannot be initialized.'
        raise TypeError(msg)

    cdef void init(self, CBoundExpressions bound_expressions):
        self.c_bound_exprs = bound_expressions

    @property
    def schema(self):
        """
        The common schema that all expressions are bound to
        """
        return pyarrow_wrap_schema(self.c_bound_exprs.schema)

    @property
    def expressions(self):
        """
        A dict from expression name to expression
        """
        expr_dict = {}
        for named_expr in self.c_bound_exprs.named_expressions:
            name = frombytes(named_expr.name)
            expr = Expression.wrap(named_expr.expression)
            expr_dict[name] = expr
        return expr_dict

    @staticmethod
    cdef wrap(const CBoundExpressions& bound_expressions):
        cdef BoundExpressions self = BoundExpressions.__new__(BoundExpressions)
        self.init(bound_expressions)
        return self

    @classmethod
    def from_substrait(cls, message):
        """
        Convert a Substrait message into a BoundExpressions object

        Parameters
        ----------
        message : Buffer or bytes or protobuf Message
            The message to convert to a BoundExpressions object

        Returns
        -------
        BoundExpressions
            The converted expressions, their names, and the bound schema
        """
        if isinstance(message, (bytes, memoryview)):
            return deserialize_expressions(message)
        elif isinstance(message, Buffer):
            return deserialize_expressions(message)
        else:
            try:
                return deserialize_expressions(message.SerializeToString())
            except AttributeError:
                raise TypeError(
                    f"Expected 'pyarrow.Buffer' or bytes or protobuf Message, got '{type(message)}'")


def deserialize_expressions(buf):
    """
    Deserialize an ExtendedExpression Substrait message into a BoundExpressions object

    Parameters
    ----------
    buf : Buffer or bytes
        The message to deserialize

    Returns
    -------
    BoundExpressions
        The deserialized expressions, their names, and the bound schema
    """
    cdef:
        shared_ptr[CBuffer] c_buffer
        CResult[CBoundExpressions] c_res_bound_exprs
        CBoundExpressions c_bound_exprs

    if isinstance(buf, (bytes, memoryview)):
        c_buffer = pyarrow_unwrap_buffer(py_buffer(buf))
    elif isinstance(buf, Buffer):
        c_buffer = pyarrow_unwrap_buffer(buf)
    else:
        raise TypeError(
            f"Expected 'pyarrow.Buffer' or bytes, got '{type(buf)}'")

    with nogil:
        c_res_bound_exprs = DeserializeExpressions(deref(c_buffer))
        c_bound_exprs = GetResultValue(c_res_bound_exprs)

    return BoundExpressions.wrap(c_bound_exprs)


def get_supported_functions():
    """
    Get a list of Substrait functions that the underlying
    engine currently supports.

    Returns
    -------
    list[str]
        A list of function ids encoded as '{uri}#{name}'
    """

    cdef:
        ExtensionIdRegistry* c_id_registry
        std_vector[c_string] c_ids

    c_id_registry = default_extension_id_registry()
    c_ids = c_id_registry.GetSupportedSubstraitFunctions()

    functions_list = []
    for c_id in c_ids:
        functions_list.append(frombytes(c_id))
    return functions_list