Files
Alessandro MolinaandRaúl Cumplido e0ab40d379 GH-41692: [Python] Improve substrait extended expressions support (#41693)
Addresses some missing features and usability issues when using PyArrow
with Substrait ExtendedExpressions
* GitHub Issue: #41692

- [x] Allow passing `BoundExpressions` for `Scanner(columns=X)` instead
of a dict of expressions.
- [x] Allow passing `BoundExpressions` for `Scanner(filter=X)` so that
user doesn't have to distinguish between `Expression` and
`BoundExpressions` and can always just use
`pyarrow.substrait.deserialize_expressions`
- [x] Allow decoding `pyarrow.BoundExpressions` directly from
`protobuf.Message`, thus allowing to use substrait-python objects.
- [x] Return `memoryview` from methods encoding substrait, so that those
can be directly passed to substrait-python (or more in general other
python libraries) without a copy being involved.
- [x] Allow decoding messages from `memoryview` so that the output of
encoding functions can be sent back to dencoding functions.
- [x] Allow to encode and decode schemas from substrait 
- [x] When encoding schemas return the extension types required for a
substrait consumer to decode the schema
- [x] Handle arrow extension types when decoding a schema
- [x] Update docstrings and documentation

---------

Co-authored-by: Raúl Cumplido <raulcumplido@gmail.com>
2024-10-16 14:58:37 +02:00

101 lines
4.0 KiB
Cython

# 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.
# distutils: language = c++
from libcpp.vector cimport vector as std_vector
from pyarrow.includes.common cimport *
from pyarrow.includes.libarrow cimport *
from pyarrow.includes.libarrow_acero cimport *
ctypedef CResult[CDeclaration] CNamedTableProvider(const std_vector[c_string]&, const CSchema&)
cdef extern from "arrow/engine/substrait/options.h" namespace "arrow::engine" nogil:
cdef enum ConversionStrictness \
"arrow::engine::ConversionStrictness":
EXACT_ROUNDTRIP \
"arrow::engine::ConversionStrictness::EXACT_ROUNDTRIP"
PRESERVE_STRUCTURE \
"arrow::engine::ConversionStrictness::PRESERVE_STRUCTURE"
BEST_EFFORT \
"arrow::engine::ConversionStrictness::BEST_EFFORT"
cdef cppclass CConversionOptions \
"arrow::engine::ConversionOptions":
CConversionOptions()
ConversionStrictness strictness
function[CNamedTableProvider] named_table_provider
c_bool allow_arrow_extensions
cdef extern from "arrow/engine/substrait/extension_set.h" \
namespace "arrow::engine" nogil:
cdef struct CSubstraitId "arrow::engine::Id":
cpp_string_view uri
cpp_string_view name
cdef struct CExtensionSetTypeRecord "arrow::engine::ExtensionSet::TypeRecord":
CSubstraitId id
shared_ptr[CDataType] type
cdef cppclass CExtensionSet "arrow::engine::ExtensionSet":
CExtensionSet()
unordered_map[uint32_t, cpp_string_view]& uris()
CResult[uint32_t] EncodeType(const CDataType&)
CResult[CExtensionSetTypeRecord] DecodeType(uint32_t)
cdef cppclass ExtensionIdRegistry:
std_vector[c_string] GetSupportedSubstraitFunctions()
ExtensionIdRegistry* default_extension_id_registry()
cdef extern from "arrow/engine/substrait/relation.h" namespace "arrow::engine" nogil:
cdef cppclass CNamedExpression "arrow::engine::NamedExpression":
CExpression expression
c_string name
cdef cppclass CBoundExpressions "arrow::engine::BoundExpressions":
std_vector[CNamedExpression] named_expressions
shared_ptr[CSchema] schema
cdef extern from "arrow/engine/substrait/serde.h" namespace "arrow::engine" nogil:
CResult[shared_ptr[CBuffer]] SerializeExpressions(
const CBoundExpressions& bound_expressions, const CConversionOptions& conversion_options)
CResult[CBoundExpressions] DeserializeExpressions(
const CBuffer& serialized_expressions)
CResult[shared_ptr[CBuffer]] SerializeSchema(
const CSchema &schema, CExtensionSet* extension_set,
const CConversionOptions& conversion_options)
CResult[shared_ptr[CSchema]] DeserializeSchema(
const CBuffer& serialized_schema, const CExtensionSet& extension_set,
const CConversionOptions& conversion_options)
cdef extern from "arrow/engine/substrait/util.h" namespace "arrow::engine" nogil:
CResult[shared_ptr[CRecordBatchReader]] ExecuteSerializedPlan(
const CBuffer& substrait_buffer, const ExtensionIdRegistry* registry,
CFunctionRegistry* func_registry, const CConversionOptions& conversion_options,
c_bool use_threads)
CResult[shared_ptr[CBuffer]] SerializeJsonPlan(const c_string& substrait_json)