diff --git a/lib/python3.12/site-packages/cloudpickle/__init__.py b/lib/python3.12/site-packages/cloudpickle/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9faf0b6d81e958a7c41da4d920a5b77a2548b526 --- /dev/null +++ b/lib/python3.12/site-packages/cloudpickle/__init__.py @@ -0,0 +1,18 @@ +from . import cloudpickle +from .cloudpickle import * # noqa + +__doc__ = cloudpickle.__doc__ + +__version__ = "3.1.2" + +__all__ = [ # noqa + "__version__", + "Pickler", + "CloudPickler", + "dumps", + "loads", + "dump", + "load", + "register_pickle_by_value", + "unregister_pickle_by_value", +] diff --git a/lib/python3.12/site-packages/cloudpickle/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/cloudpickle/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8315ad786f53f23d782a92ea1150611fb3b50565 Binary files /dev/null and b/lib/python3.12/site-packages/cloudpickle/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc b/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bffae5d590512e485c1bca0eb8026c0cfa0bd579 Binary files /dev/null and b/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc b/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cca8aa2614c98b2da2cebfe3ca4dbad1c0bab86b Binary files /dev/null and b/lib/python3.12/site-packages/cloudpickle/__pycache__/cloudpickle_fast.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/cloudpickle/cloudpickle.py b/lib/python3.12/site-packages/cloudpickle/cloudpickle.py new file mode 100644 index 0000000000000000000000000000000000000000..e600b35f28422096173a783380ec912a66b453cf --- /dev/null +++ b/lib/python3.12/site-packages/cloudpickle/cloudpickle.py @@ -0,0 +1,1552 @@ +"""Pickler class to extend the standard pickle.Pickler functionality + +The main objective is to make it natural to perform distributed computing on +clusters (such as PySpark, Dask, Ray...) with interactively defined code +(functions, classes, ...) written in notebooks or console. + +In particular this pickler adds the following features: +- serialize interactively-defined or locally-defined functions, classes, + enums, typevars, lambdas and nested functions to compiled byte code; +- deal with some other non-serializable objects in an ad-hoc manner where + applicable. + +This pickler is therefore meant to be used for the communication between short +lived Python processes running the same version of Python and libraries. In +particular, it is not meant to be used for long term storage of Python objects. + +It does not include an unpickler, as standard Python unpickling suffices. + +This module was extracted from the `cloud` package, developed by `PiCloud, Inc. +`_. + +Copyright (c) 2012-now, CloudPickle developers and contributors. +Copyright (c) 2012, Regents of the University of California. +Copyright (c) 2009 `PiCloud, Inc. `_. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions +are met: + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + * Neither the name of the University of California, Berkeley nor the + names of its contributors may be used to endorse or promote + products derived from this software without specific prior written + permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF +LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING +NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +""" + +import _collections_abc +from collections import ChainMap, OrderedDict +import abc +import builtins +import copyreg +import dataclasses +import dis +from enum import Enum +import io +import itertools +import logging +import opcode +import pickle +from pickle import _getattribute as _pickle_getattribute +import platform +import struct +import sys +import threading +import types +import typing +import uuid +import warnings +import weakref + +# The following import is required to be imported in the cloudpickle +# namespace to be able to load pickle files generated with older versions of +# cloudpickle. See: tests/test_backward_compat.py +from types import CellType # noqa: F401 + + +# cloudpickle is meant for inter process communication: we expect all +# communicating processes to run the same Python version hence we favor +# communication speed over compatibility: +DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL + +# Names of modules whose resources should be treated as dynamic. +_PICKLE_BY_VALUE_MODULES = set() + +# Track the provenance of reconstructed dynamic classes to make it possible to +# reconstruct instances from the matching singleton class definition when +# appropriate and preserve the usual "isinstance" semantics of Python objects. +_DYNAMIC_CLASS_TRACKER_BY_CLASS = weakref.WeakKeyDictionary() +_DYNAMIC_CLASS_TRACKER_BY_ID = weakref.WeakValueDictionary() +_DYNAMIC_CLASS_TRACKER_LOCK = threading.Lock() + +PYPY = platform.python_implementation() == "PyPy" + +builtin_code_type = None +if PYPY: + # builtin-code objects only exist in pypy + builtin_code_type = type(float.__new__.__code__) + +_extract_code_globals_cache = weakref.WeakKeyDictionary() + + +def _get_or_create_tracker_id(class_def): + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_tracker_id = _DYNAMIC_CLASS_TRACKER_BY_CLASS.get(class_def) + if class_tracker_id is None: + class_tracker_id = uuid.uuid4().hex + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + _DYNAMIC_CLASS_TRACKER_BY_ID[class_tracker_id] = class_def + return class_tracker_id + + +def _lookup_class_or_track(class_tracker_id, class_def): + if class_tracker_id is not None: + with _DYNAMIC_CLASS_TRACKER_LOCK: + class_def = _DYNAMIC_CLASS_TRACKER_BY_ID.setdefault( + class_tracker_id, class_def + ) + _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id + return class_def + + +def register_pickle_by_value(module): + """Register a module to make its functions and classes picklable by value. + + By default, functions and classes that are attributes of an importable + module are to be pickled by reference, that is relying on re-importing + the attribute from the module at load time. + + If `register_pickle_by_value(module)` is called, all its functions and + classes are subsequently to be pickled by value, meaning that they can + be loaded in Python processes where the module is not importable. + + This is especially useful when developing a module in a distributed + execution environment: restarting the client Python process with the new + source code is enough: there is no need to re-install the new version + of the module on all the worker nodes nor to restart the workers. + + Note: this feature is considered experimental. See the cloudpickle + README.md file for more details and limitations. + """ + if not isinstance(module, types.ModuleType): + raise ValueError(f"Input should be a module object, got {str(module)} instead") + # In the future, cloudpickle may need a way to access any module registered + # for pickling by value in order to introspect relative imports inside + # functions pickled by value. (see + # https://github.com/cloudpipe/cloudpickle/pull/417#issuecomment-873684633). + # This access can be ensured by checking that module is present in + # sys.modules at registering time and assuming that it will still be in + # there when accessed during pickling. Another alternative would be to + # store a weakref to the module. Even though cloudpickle does not implement + # this introspection yet, in order to avoid a possible breaking change + # later, we still enforce the presence of module inside sys.modules. + if module.__name__ not in sys.modules: + raise ValueError( + f"{module} was not imported correctly, have you used an " + "`import` statement to access it?" + ) + _PICKLE_BY_VALUE_MODULES.add(module.__name__) + + +def unregister_pickle_by_value(module): + """Unregister that the input module should be pickled by value.""" + if not isinstance(module, types.ModuleType): + raise ValueError(f"Input should be a module object, got {str(module)} instead") + if module.__name__ not in _PICKLE_BY_VALUE_MODULES: + raise ValueError(f"{module} is not registered for pickle by value") + else: + _PICKLE_BY_VALUE_MODULES.remove(module.__name__) + + +def list_registry_pickle_by_value(): + return _PICKLE_BY_VALUE_MODULES.copy() + + +def _is_registered_pickle_by_value(module): + module_name = module.__name__ + if module_name in _PICKLE_BY_VALUE_MODULES: + return True + while True: + parent_name = module_name.rsplit(".", 1)[0] + if parent_name == module_name: + break + if parent_name in _PICKLE_BY_VALUE_MODULES: + return True + module_name = parent_name + return False + + +if sys.version_info >= (3, 14): + def _getattribute(obj, name): + return _pickle_getattribute(obj, name.split('.')) +else: + def _getattribute(obj, name): + return _pickle_getattribute(obj, name)[0] + + +def _whichmodule(obj, name): + """Find the module an object belongs to. + + This function differs from ``pickle.whichmodule`` in two ways: + - it does not mangle the cases where obj's module is __main__ and obj was + not found in any module. + - Errors arising during module introspection are ignored, as those errors + are considered unwanted side effects. + """ + module_name = getattr(obj, "__module__", None) + + if module_name is not None: + return module_name + # Protect the iteration by using a copy of sys.modules against dynamic + # modules that trigger imports of other modules upon calls to getattr or + # other threads importing at the same time. + for module_name, module in sys.modules.copy().items(): + # Some modules such as coverage can inject non-module objects inside + # sys.modules + if ( + module_name == "__main__" + or module_name == "__mp_main__" + or module is None + or not isinstance(module, types.ModuleType) + ): + continue + try: + if _getattribute(module, name) is obj: + return module_name + except Exception: + pass + return None + + +def _should_pickle_by_reference(obj, name=None): + """Test whether an function or a class should be pickled by reference + + Pickling by reference means by that the object (typically a function or a + class) is an attribute of a module that is assumed to be importable in the + target Python environment. Loading will therefore rely on importing the + module and then calling `getattr` on it to access the function or class. + + Pickling by reference is the only option to pickle functions and classes + in the standard library. In cloudpickle the alternative option is to + pickle by value (for instance for interactively or locally defined + functions and classes or for attributes of modules that have been + explicitly registered to be pickled by value. + """ + if isinstance(obj, types.FunctionType) or issubclass(type(obj), type): + module_and_name = _lookup_module_and_qualname(obj, name=name) + if module_and_name is None: + return False + module, name = module_and_name + return not _is_registered_pickle_by_value(module) + + elif isinstance(obj, types.ModuleType): + # We assume that sys.modules is primarily used as a cache mechanism for + # the Python import machinery. Checking if a module has been added in + # is sys.modules therefore a cheap and simple heuristic to tell us + # whether we can assume that a given module could be imported by name + # in another Python process. + if _is_registered_pickle_by_value(obj): + return False + return obj.__name__ in sys.modules + else: + raise TypeError( + "cannot check importability of {} instances".format(type(obj).__name__) + ) + + +def _lookup_module_and_qualname(obj, name=None): + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: # pragma: no cover + # This used to be needed for Python 2.7 support but is probably not + # needed anymore. However we keep the __name__ introspection in case + # users of cloudpickle rely on this old behavior for unknown reasons. + name = getattr(obj, "__name__", None) + + module_name = _whichmodule(obj, name) + + if module_name is None: + # In this case, obj.__module__ is None AND obj was not found in any + # imported module. obj is thus treated as dynamic. + return None + + if module_name == "__main__": + return None + + # Note: if module_name is in sys.modules, the corresponding module is + # assumed importable at unpickling time. See #357 + module = sys.modules.get(module_name, None) + if module is None: + # The main reason why obj's module would not be imported is that this + # module has been dynamically created, using for example + # types.ModuleType. The other possibility is that module was removed + # from sys.modules after obj was created/imported. But this case is not + # supported, as the standard pickle does not support it either. + return None + + try: + obj2 = _getattribute(module, name) + except AttributeError: + # obj was not found inside the module it points to + return None + if obj2 is not obj: + return None + return module, name + + +def _extract_code_globals(co): + """Find all globals names read or written to by codeblock co.""" + out_names = _extract_code_globals_cache.get(co) + if out_names is None: + # We use a dict with None values instead of a set to get a + # deterministic order and avoid introducing non-deterministic pickle + # bytes as a results. + out_names = {name: None for name in _walk_global_ops(co)} + + # Declaring a function inside another one using the "def ..." syntax + # generates a constant code object corresponding to the one of the + # nested function's As the nested function may itself need global + # variables, we need to introspect its code, extract its globals, (look + # for code object in it's co_consts attribute..) and add the result to + # code_globals + if co.co_consts: + for const in co.co_consts: + if isinstance(const, types.CodeType): + out_names.update(_extract_code_globals(const)) + + _extract_code_globals_cache[co] = out_names + + return out_names + + +def _find_imported_submodules(code, top_level_dependencies): + """Find currently imported submodules used by a function. + + Submodules used by a function need to be detected and referenced for the + function to work correctly at depickling time. Because submodules can be + referenced as attribute of their parent package (``package.submodule``), we + need a special introspection technique that does not rely on GLOBAL-related + opcodes to find references of them in a code object. + + Example: + ``` + import concurrent.futures + import cloudpickle + def func(): + x = concurrent.futures.ThreadPoolExecutor + if __name__ == '__main__': + cloudpickle.dumps(func) + ``` + The globals extracted by cloudpickle in the function's state include the + concurrent package, but not its submodule (here, concurrent.futures), which + is the module used by func. Find_imported_submodules will detect the usage + of concurrent.futures. Saving this module alongside with func will ensure + that calling func once depickled does not fail due to concurrent.futures + not being imported + """ + + subimports = [] + # check if any known dependency is an imported package + for x in top_level_dependencies: + if ( + isinstance(x, types.ModuleType) + and hasattr(x, "__package__") + and x.__package__ + ): + # check if the package has any currently loaded sub-imports + prefix = x.__name__ + "." + # A concurrent thread could mutate sys.modules, + # make sure we iterate over a copy to avoid exceptions + for name in list(sys.modules): + # Older versions of pytest will add a "None" module to + # sys.modules. + if name is not None and name.startswith(prefix): + # check whether the function can address the sub-module + tokens = set(name[len(prefix) :].split(".")) + if not tokens - set(code.co_names): + subimports.append(sys.modules[name]) + return subimports + + +# relevant opcodes +STORE_GLOBAL = opcode.opmap["STORE_GLOBAL"] +DELETE_GLOBAL = opcode.opmap["DELETE_GLOBAL"] +LOAD_GLOBAL = opcode.opmap["LOAD_GLOBAL"] +GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL) +HAVE_ARGUMENT = dis.HAVE_ARGUMENT +EXTENDED_ARG = dis.EXTENDED_ARG + + +_BUILTIN_TYPE_NAMES = {} +for k, v in types.__dict__.items(): + if type(v) is type: + _BUILTIN_TYPE_NAMES[v] = k + + +def _builtin_type(name): + if name == "ClassType": # pragma: no cover + # Backward compat to load pickle files generated with cloudpickle + # < 1.3 even if loading pickle files from older versions is not + # officially supported. + return type + return getattr(types, name) + + +def _walk_global_ops(code): + """Yield referenced name for global-referencing instructions in code.""" + for instr in dis.get_instructions(code): + op = instr.opcode + if op in GLOBAL_OPS: + yield instr.argval + + +def _extract_class_dict(cls): + """Retrieve a copy of the dict of a class without the inherited method.""" + # Hack to circumvent non-predictable memoization caused by string interning. + # See the inline comment in _class_setstate for details. + clsdict = {"".join(k): cls.__dict__[k] for k in sorted(cls.__dict__)} + + if len(cls.__bases__) == 1: + inherited_dict = cls.__bases__[0].__dict__ + else: + inherited_dict = {} + for base in reversed(cls.__bases__): + inherited_dict.update(base.__dict__) + to_remove = [] + for name, value in clsdict.items(): + try: + base_value = inherited_dict[name] + if value is base_value: + to_remove.append(name) + except KeyError: + pass + for name in to_remove: + clsdict.pop(name) + return clsdict + + +def is_tornado_coroutine(func): + """Return whether `func` is a Tornado coroutine function. + + Running coroutines are not supported. + """ + warnings.warn( + "is_tornado_coroutine is deprecated in cloudpickle 3.0 and will be " + "removed in cloudpickle 4.0. Use tornado.gen.is_coroutine_function " + "directly instead.", + category=DeprecationWarning, + ) + if "tornado.gen" not in sys.modules: + return False + gen = sys.modules["tornado.gen"] + if not hasattr(gen, "is_coroutine_function"): + # Tornado version is too old + return False + return gen.is_coroutine_function(func) + + +def subimport(name): + # We cannot do simply: `return __import__(name)`: Indeed, if ``name`` is + # the name of a submodule, __import__ will return the top-level root module + # of this submodule. For instance, __import__('os.path') returns the `os` + # module. + __import__(name) + return sys.modules[name] + + +def dynamic_subimport(name, vars): + mod = types.ModuleType(name) + mod.__dict__.update(vars) + mod.__dict__["__builtins__"] = builtins.__dict__ + return mod + + +def _get_cell_contents(cell): + try: + return cell.cell_contents + except ValueError: + # Handle empty cells explicitly with a sentinel value. + return _empty_cell_value + + +def instance(cls): + """Create a new instance of a class. + + Parameters + ---------- + cls : type + The class to create an instance of. + + Returns + ------- + instance : cls + A new instance of ``cls``. + """ + return cls() + + +@instance +class _empty_cell_value: + """Sentinel for empty closures.""" + + @classmethod + def __reduce__(cls): + return cls.__name__ + + +def _make_function(code, globals, name, argdefs, closure): + # Setting __builtins__ in globals is needed for nogil CPython. + globals["__builtins__"] = __builtins__ + return types.FunctionType(code, globals, name, argdefs, closure) + + +def _make_empty_cell(): + if False: + # trick the compiler into creating an empty cell in our lambda + cell = None + raise AssertionError("this route should not be executed") + + return (lambda: cell).__closure__[0] + + +def _make_cell(value=_empty_cell_value): + cell = _make_empty_cell() + if value is not _empty_cell_value: + cell.cell_contents = value + return cell + + +def _make_skeleton_class( + type_constructor, name, bases, type_kwargs, class_tracker_id, extra +): + """Build dynamic class with an empty __dict__ to be filled once memoized + + If class_tracker_id is not None, try to lookup an existing class definition + matching that id. If none is found, track a newly reconstructed class + definition under that id so that other instances stemming from the same + class id will also reuse this class definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + # We need to intern the keys of the type_kwargs dict to avoid having + # different pickles for the same dynamic class depending on whether it was + # dynamically created or reconstructed from a pickled stream. + type_kwargs = {sys.intern(k): v for k, v in type_kwargs.items()} + + skeleton_class = types.new_class( + name, bases, {"metaclass": type_constructor}, lambda ns: ns.update(type_kwargs) + ) + + return _lookup_class_or_track(class_tracker_id, skeleton_class) + + +def _make_skeleton_enum( + bases, name, qualname, members, module, class_tracker_id, extra +): + """Build dynamic enum with an empty __dict__ to be filled once memoized + + The creation of the enum class is inspired by the code of + EnumMeta._create_. + + If class_tracker_id is not None, try to lookup an existing enum definition + matching that id. If none is found, track a newly reconstructed enum + definition under that id so that other instances stemming from the same + class id will also reuse this enum definition. + + The "extra" variable is meant to be a dict (or None) that can be used for + forward compatibility shall the need arise. + """ + # enums always inherit from their base Enum class at the last position in + # the list of base classes: + enum_base = bases[-1] + metacls = enum_base.__class__ + classdict = metacls.__prepare__(name, bases) + + for member_name, member_value in members.items(): + classdict[member_name] = member_value + enum_class = metacls.__new__(metacls, name, bases, classdict) + enum_class.__module__ = module + enum_class.__qualname__ = qualname + + return _lookup_class_or_track(class_tracker_id, enum_class) + + +def _make_typevar(name, bound, constraints, covariant, contravariant, class_tracker_id): + tv = typing.TypeVar( + name, + *constraints, + bound=bound, + covariant=covariant, + contravariant=contravariant, + ) + return _lookup_class_or_track(class_tracker_id, tv) + + +def _decompose_typevar(obj): + return ( + obj.__name__, + obj.__bound__, + obj.__constraints__, + obj.__covariant__, + obj.__contravariant__, + _get_or_create_tracker_id(obj), + ) + + +def _typevar_reduce(obj): + # TypeVar instances require the module information hence why we + # are not using the _should_pickle_by_reference directly + module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__) + + if module_and_name is None: + return (_make_typevar, _decompose_typevar(obj)) + elif _is_registered_pickle_by_value(module_and_name[0]): + return (_make_typevar, _decompose_typevar(obj)) + + return (getattr, module_and_name) + + +def _get_bases(typ): + if "__orig_bases__" in getattr(typ, "__dict__", {}): + # For generic types (see PEP 560) + # Note that simply checking `hasattr(typ, '__orig_bases__')` is not + # correct. Subclasses of a fully-parameterized generic class does not + # have `__orig_bases__` defined, but `hasattr(typ, '__orig_bases__')` + # will return True because it's defined in the base class. + bases_attr = "__orig_bases__" + else: + # For regular class objects + bases_attr = "__bases__" + return getattr(typ, bases_attr) + + +def _make_dict_keys(obj, is_ordered=False): + if is_ordered: + return OrderedDict.fromkeys(obj).keys() + else: + return dict.fromkeys(obj).keys() + + +def _make_dict_values(obj, is_ordered=False): + if is_ordered: + return OrderedDict((i, _) for i, _ in enumerate(obj)).values() + else: + return {i: _ for i, _ in enumerate(obj)}.values() + + +def _make_dict_items(obj, is_ordered=False): + if is_ordered: + return OrderedDict(obj).items() + else: + return obj.items() + + +# COLLECTION OF OBJECTS __getnewargs__-LIKE METHODS +# ------------------------------------------------- + + +def _class_getnewargs(obj): + type_kwargs = {} + if "__module__" in obj.__dict__: + type_kwargs["__module__"] = obj.__module__ + + __dict__ = obj.__dict__.get("__dict__", None) + if isinstance(__dict__, property): + type_kwargs["__dict__"] = __dict__ + + return ( + type(obj), + obj.__name__, + _get_bases(obj), + type_kwargs, + _get_or_create_tracker_id(obj), + None, + ) + + +def _enum_getnewargs(obj): + members = {e.name: e.value for e in obj} + return ( + obj.__bases__, + obj.__name__, + obj.__qualname__, + members, + obj.__module__, + _get_or_create_tracker_id(obj), + None, + ) + + +# COLLECTION OF OBJECTS RECONSTRUCTORS +# ------------------------------------ +def _file_reconstructor(retval): + return retval + + +# COLLECTION OF OBJECTS STATE GETTERS +# ----------------------------------- + + +def _function_getstate(func): + # - Put func's dynamic attributes (stored in func.__dict__) in state. These + # attributes will be restored at unpickling time using + # f.__dict__.update(state) + # - Put func's members into slotstate. Such attributes will be restored at + # unpickling time by iterating over slotstate and calling setattr(func, + # slotname, slotvalue) + slotstate = { + # Hack to circumvent non-predictable memoization caused by string interning. + # See the inline comment in _class_setstate for details. + "__name__": "".join(func.__name__), + "__qualname__": "".join(func.__qualname__), + "__annotations__": func.__annotations__, + "__kwdefaults__": func.__kwdefaults__, + "__defaults__": func.__defaults__, + "__module__": func.__module__, + "__doc__": func.__doc__, + "__closure__": func.__closure__, + } + + f_globals_ref = _extract_code_globals(func.__code__) + f_globals = {k: func.__globals__[k] for k in f_globals_ref if k in func.__globals__} + + if func.__closure__ is not None: + closure_values = list(map(_get_cell_contents, func.__closure__)) + else: + closure_values = () + + # Extract currently-imported submodules used by func. Storing these modules + # in a smoke _cloudpickle_subimports attribute of the object's state will + # trigger the side effect of importing these modules at unpickling time + # (which is necessary for func to work correctly once depickled) + slotstate["_cloudpickle_submodules"] = _find_imported_submodules( + func.__code__, itertools.chain(f_globals.values(), closure_values) + ) + slotstate["__globals__"] = f_globals + + # Hack to circumvent non-predictable memoization caused by string interning. + # See the inline comment in _class_setstate for details. + state = {"".join(k): v for k, v in func.__dict__.items()} + return state, slotstate + + +def _class_getstate(obj): + clsdict = _extract_class_dict(obj) + clsdict.pop("__weakref__", None) + + if issubclass(type(obj), abc.ABCMeta): + # If obj is an instance of an ABCMeta subclass, don't pickle the + # cache/negative caches populated during isinstance/issubclass + # checks, but pickle the list of registered subclasses of obj. + clsdict.pop("_abc_cache", None) + clsdict.pop("_abc_negative_cache", None) + clsdict.pop("_abc_negative_cache_version", None) + registry = clsdict.pop("_abc_registry", None) + if registry is None: + # The abc caches and registered subclasses of a + # class are bundled into the single _abc_impl attribute + clsdict.pop("_abc_impl", None) + (registry, _, _, _) = abc._get_dump(obj) + + clsdict["_abc_impl"] = [subclass_weakref() for subclass_weakref in registry] + else: + # In the above if clause, registry is a set of weakrefs -- in + # this case, registry is a WeakSet + clsdict["_abc_impl"] = [type_ for type_ in registry] + + if "__slots__" in clsdict: + # pickle string length optimization: member descriptors of obj are + # created automatically from obj's __slots__ attribute, no need to + # save them in obj's state + if isinstance(obj.__slots__, str): + clsdict.pop(obj.__slots__) + else: + for k in obj.__slots__: + clsdict.pop(k, None) + + clsdict.pop("__dict__", None) # unpicklable property object + + if sys.version_info >= (3, 14): + # PEP-649/749: __annotate_func__ contains a closure that references the class + # dict. We need to exclude it from pickling. Python will recreate it when + # __annotations__ is accessed at unpickling time. + clsdict.pop("__annotate_func__", None) + + return (clsdict, {}) + + +def _enum_getstate(obj): + clsdict, slotstate = _class_getstate(obj) + + members = {e.name: e.value for e in obj} + # Cleanup the clsdict that will be passed to _make_skeleton_enum: + # Those attributes are already handled by the metaclass. + for attrname in [ + "_generate_next_value_", + "_member_names_", + "_member_map_", + "_member_type_", + "_value2member_map_", + ]: + clsdict.pop(attrname, None) + for member in members: + clsdict.pop(member) + # Special handling of Enum subclasses + return clsdict, slotstate + + +# COLLECTIONS OF OBJECTS REDUCERS +# ------------------------------- +# A reducer is a function taking a single argument (obj), and that returns a +# tuple with all the necessary data to re-construct obj. Apart from a few +# exceptions (list, dict, bytes, int, etc.), a reducer is necessary to +# correctly pickle an object. +# While many built-in objects (Exceptions objects, instances of the "object" +# class, etc), are shipped with their own built-in reducer (invoked using +# obj.__reduce__), some do not. The following methods were created to "fill +# these holes". + + +def _code_reduce(obj): + """code object reducer.""" + # If you are not sure about the order of arguments, take a look at help + # of the specific type from types, for example: + # >>> from types import CodeType + # >>> help(CodeType) + + # Hack to circumvent non-predictable memoization caused by string interning. + # See the inline comment in _class_setstate for details. + co_name = "".join(obj.co_name) + + # Create shallow copies of these tuple to make cloudpickle payload deterministic. + # When creating a code object during load, copies of these four tuples are + # created, while in the main process, these tuples can be shared. + # By always creating copies, we make sure the resulting payload is deterministic. + co_names = tuple(name for name in obj.co_names) + co_varnames = tuple(name for name in obj.co_varnames) + co_freevars = tuple(name for name in obj.co_freevars) + co_cellvars = tuple(name for name in obj.co_cellvars) + if hasattr(obj, "co_exceptiontable"): + # Python 3.11 and later: there are some new attributes + # related to the enhanced exceptions. + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + co_names, + co_varnames, + obj.co_filename, + co_name, + obj.co_qualname, + obj.co_firstlineno, + obj.co_linetable, + obj.co_exceptiontable, + co_freevars, + co_cellvars, + ) + elif hasattr(obj, "co_linetable"): + # Python 3.10 and later: obj.co_lnotab is deprecated and constructor + # expects obj.co_linetable instead. + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + co_names, + co_varnames, + obj.co_filename, + co_name, + obj.co_firstlineno, + obj.co_linetable, + co_freevars, + co_cellvars, + ) + elif hasattr(obj, "co_nmeta"): # pragma: no cover + # "nogil" Python: modified attributes from 3.9 + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_framesize, + obj.co_ndefaultargs, + obj.co_nmeta, + obj.co_flags, + obj.co_code, + obj.co_consts, + co_varnames, + obj.co_filename, + co_name, + obj.co_firstlineno, + obj.co_lnotab, + obj.co_exc_handlers, + obj.co_jump_table, + co_freevars, + co_cellvars, + obj.co_free2reg, + obj.co_cell2reg, + ) + else: + # Backward compat for 3.8 and 3.9 + args = ( + obj.co_argcount, + obj.co_posonlyargcount, + obj.co_kwonlyargcount, + obj.co_nlocals, + obj.co_stacksize, + obj.co_flags, + obj.co_code, + obj.co_consts, + co_names, + co_varnames, + obj.co_filename, + co_name, + obj.co_firstlineno, + obj.co_lnotab, + co_freevars, + co_cellvars, + ) + return types.CodeType, args + + +def _cell_reduce(obj): + """Cell (containing values of a function's free variables) reducer.""" + try: + obj.cell_contents + except ValueError: # cell is empty + return _make_empty_cell, () + else: + return _make_cell, (obj.cell_contents,) + + +def _classmethod_reduce(obj): + orig_func = obj.__func__ + return type(obj), (orig_func,) + + +def _file_reduce(obj): + """Save a file.""" + import io + + if not hasattr(obj, "name") or not hasattr(obj, "mode"): + raise pickle.PicklingError( + "Cannot pickle files that do not map to an actual file" + ) + if obj is sys.stdout: + return getattr, (sys, "stdout") + if obj is sys.stderr: + return getattr, (sys, "stderr") + if obj is sys.stdin: + raise pickle.PicklingError("Cannot pickle standard input") + if obj.closed: + raise pickle.PicklingError("Cannot pickle closed files") + if hasattr(obj, "isatty") and obj.isatty(): + raise pickle.PicklingError("Cannot pickle files that map to tty objects") + if "r" not in obj.mode and "+" not in obj.mode: + raise pickle.PicklingError( + "Cannot pickle files that are not opened for reading: %s" % obj.mode + ) + + name = obj.name + + retval = io.StringIO() + + try: + # Read the whole file + curloc = obj.tell() + obj.seek(0) + contents = obj.read() + obj.seek(curloc) + except OSError as e: + raise pickle.PicklingError( + "Cannot pickle file %s as it cannot be read" % name + ) from e + retval.write(contents) + retval.seek(curloc) + + retval.name = name + return _file_reconstructor, (retval,) + + +def _getset_descriptor_reduce(obj): + return getattr, (obj.__objclass__, obj.__name__) + + +def _mappingproxy_reduce(obj): + return types.MappingProxyType, (dict(obj),) + + +def _memoryview_reduce(obj): + return bytes, (obj.tobytes(),) + + +def _module_reduce(obj): + if _should_pickle_by_reference(obj): + return subimport, (obj.__name__,) + else: + # Some external libraries can populate the "__builtins__" entry of a + # module's `__dict__` with unpicklable objects (see #316). For that + # reason, we do not attempt to pickle the "__builtins__" entry, and + # restore a default value for it at unpickling time. + state = obj.__dict__.copy() + state.pop("__builtins__", None) + return dynamic_subimport, (obj.__name__, state) + + +def _method_reduce(obj): + return (types.MethodType, (obj.__func__, obj.__self__)) + + +def _logger_reduce(obj): + return logging.getLogger, (obj.name,) + + +def _root_logger_reduce(obj): + return logging.getLogger, () + + +def _property_reduce(obj): + return property, (obj.fget, obj.fset, obj.fdel, obj.__doc__) + + +def _weakset_reduce(obj): + return weakref.WeakSet, (list(obj),) + + +def _dynamic_class_reduce(obj): + """Save a class that can't be referenced as a module attribute. + + This method is used to serialize classes that are defined inside + functions, or that otherwise can't be serialized as attribute lookups + from importable modules. + """ + if Enum is not None and issubclass(obj, Enum): + return ( + _make_skeleton_enum, + _enum_getnewargs(obj), + _enum_getstate(obj), + None, + None, + _class_setstate, + ) + else: + return ( + _make_skeleton_class, + _class_getnewargs(obj), + _class_getstate(obj), + None, + None, + _class_setstate, + ) + + +def _class_reduce(obj): + """Select the reducer depending on the dynamic nature of the class obj.""" + if obj is type(None): # noqa + return type, (None,) + elif obj is type(Ellipsis): + return type, (Ellipsis,) + elif obj is type(NotImplemented): + return type, (NotImplemented,) + elif obj in _BUILTIN_TYPE_NAMES: + return _builtin_type, (_BUILTIN_TYPE_NAMES[obj],) + elif not _should_pickle_by_reference(obj): + return _dynamic_class_reduce(obj) + return NotImplemented + + +def _dict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj),) + + +def _dict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj),) + + +def _dict_items_reduce(obj): + return _make_dict_items, (dict(obj),) + + +def _odict_keys_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_keys, (list(obj), True) + + +def _odict_values_reduce(obj): + # Safer not to ship the full dict as sending the rest might + # be unintended and could potentially cause leaking of + # sensitive information + return _make_dict_values, (list(obj), True) + + +def _odict_items_reduce(obj): + return _make_dict_items, (dict(obj), True) + + +def _dataclass_field_base_reduce(obj): + return _get_dataclass_field_type_sentinel, (obj.name,) + + +# COLLECTIONS OF OBJECTS STATE SETTERS +# ------------------------------------ +# state setters are called at unpickling time, once the object is created and +# it has to be updated to how it was at unpickling time. + + +def _function_setstate(obj, state): + """Update the state of a dynamic function. + + As __closure__ and __globals__ are readonly attributes of a function, we + cannot rely on the native setstate routine of pickle.load_build, that calls + setattr on items of the slotstate. Instead, we have to modify them inplace. + """ + state, slotstate = state + obj.__dict__.update(state) + + obj_globals = slotstate.pop("__globals__") + obj_closure = slotstate.pop("__closure__") + # _cloudpickle_subimports is a set of submodules that must be loaded for + # the pickled function to work correctly at unpickling time. Now that these + # submodules are depickled (hence imported), they can be removed from the + # object's state (the object state only served as a reference holder to + # these submodules) + slotstate.pop("_cloudpickle_submodules") + + obj.__globals__.update(obj_globals) + obj.__globals__["__builtins__"] = __builtins__ + + if obj_closure is not None: + for i, cell in enumerate(obj_closure): + try: + value = cell.cell_contents + except ValueError: # cell is empty + continue + obj.__closure__[i].cell_contents = value + + for k, v in slotstate.items(): + setattr(obj, k, v) + + +def _class_setstate(obj, state): + state, slotstate = state + registry = None + for attrname, attr in state.items(): + if attrname == "_abc_impl": + registry = attr + else: + # Note: setting attribute names on a class automatically triggers their + # interning in CPython: + # https://github.com/python/cpython/blob/v3.12.0/Objects/object.c#L957 + # + # This means that to get deterministic pickling for a dynamic class that + # was initially defined in a different Python process, the pickler + # needs to ensure that dynamic class and function attribute names are + # systematically copied into a non-interned version to avoid + # unpredictable pickle payloads. + # + # Indeed the Pickler's memoizer relies on physical object identity to break + # cycles in the reference graph of the object being serialized. + setattr(obj, attrname, attr) + + if sys.version_info >= (3, 13) and "__firstlineno__" in state: + # Set the Python 3.13+ only __firstlineno__ attribute one more time, as it + # will be automatically deleted by the `setattr(obj, attrname, attr)` call + # above when `attrname` is "__firstlineno__". We assume that preserving this + # information might be important for some users and that it not stale in the + # context of cloudpickle usage, hence legitimate to propagate. Furthermore it + # is necessary to do so to keep deterministic chained pickling as tested in + # test_deterministic_str_interning_for_chained_dynamic_class_pickling. + obj.__firstlineno__ = state["__firstlineno__"] + + if registry is not None: + for subclass in registry: + obj.register(subclass) + + # PEP-649/749: During pickling, we excluded the __annotate_func__ attribute but it + # will be created by Python. Subsequently, annotations will be recreated when + # __annotations__ is accessed. + + return obj + + +# COLLECTION OF DATACLASS UTILITIES +# --------------------------------- +# There are some internal sentinel values whose identity must be preserved when +# unpickling dataclass fields. Each sentinel value has a unique name that we can +# use to retrieve its identity at unpickling time. + + +_DATACLASSE_FIELD_TYPE_SENTINELS = { + dataclasses._FIELD.name: dataclasses._FIELD, + dataclasses._FIELD_CLASSVAR.name: dataclasses._FIELD_CLASSVAR, + dataclasses._FIELD_INITVAR.name: dataclasses._FIELD_INITVAR, +} + + +def _get_dataclass_field_type_sentinel(name): + return _DATACLASSE_FIELD_TYPE_SENTINELS[name] + + +class Pickler(pickle.Pickler): + # set of reducers defined and used by cloudpickle (private) + _dispatch_table = {} + _dispatch_table[classmethod] = _classmethod_reduce + _dispatch_table[io.TextIOWrapper] = _file_reduce + _dispatch_table[logging.Logger] = _logger_reduce + _dispatch_table[logging.RootLogger] = _root_logger_reduce + _dispatch_table[memoryview] = _memoryview_reduce + _dispatch_table[property] = _property_reduce + _dispatch_table[staticmethod] = _classmethod_reduce + _dispatch_table[CellType] = _cell_reduce + _dispatch_table[types.CodeType] = _code_reduce + _dispatch_table[types.GetSetDescriptorType] = _getset_descriptor_reduce + _dispatch_table[types.ModuleType] = _module_reduce + _dispatch_table[types.MethodType] = _method_reduce + _dispatch_table[types.MappingProxyType] = _mappingproxy_reduce + _dispatch_table[weakref.WeakSet] = _weakset_reduce + _dispatch_table[typing.TypeVar] = _typevar_reduce + _dispatch_table[_collections_abc.dict_keys] = _dict_keys_reduce + _dispatch_table[_collections_abc.dict_values] = _dict_values_reduce + _dispatch_table[_collections_abc.dict_items] = _dict_items_reduce + _dispatch_table[type(OrderedDict().keys())] = _odict_keys_reduce + _dispatch_table[type(OrderedDict().values())] = _odict_values_reduce + _dispatch_table[type(OrderedDict().items())] = _odict_items_reduce + _dispatch_table[abc.abstractmethod] = _classmethod_reduce + _dispatch_table[abc.abstractclassmethod] = _classmethod_reduce + _dispatch_table[abc.abstractstaticmethod] = _classmethod_reduce + _dispatch_table[abc.abstractproperty] = _property_reduce + _dispatch_table[dataclasses._FIELD_BASE] = _dataclass_field_base_reduce + + dispatch_table = ChainMap(_dispatch_table, copyreg.dispatch_table) + + # function reducers are defined as instance methods of cloudpickle.Pickler + # objects, as they rely on a cloudpickle.Pickler attribute (globals_ref) + def _dynamic_function_reduce(self, func): + """Reduce a function that is not pickleable via attribute lookup.""" + newargs = self._function_getnewargs(func) + state = _function_getstate(func) + return (_make_function, newargs, state, None, None, _function_setstate) + + def _function_reduce(self, obj): + """Reducer for function objects. + + If obj is a top-level attribute of a file-backed module, this reducer + returns NotImplemented, making the cloudpickle.Pickler fall back to + traditional pickle.Pickler routines to save obj. Otherwise, it reduces + obj using a custom cloudpickle reducer designed specifically to handle + dynamic functions. + """ + if _should_pickle_by_reference(obj): + return NotImplemented + else: + return self._dynamic_function_reduce(obj) + + def _function_getnewargs(self, func): + code = func.__code__ + + # base_globals represents the future global namespace of func at + # unpickling time. Looking it up and storing it in + # cloudpickle.Pickler.globals_ref allow functions sharing the same + # globals at pickling time to also share them once unpickled, at one + # condition: since globals_ref is an attribute of a cloudpickle.Pickler + # instance, and that a new cloudpickle.Pickler is created each time + # cloudpickle.dump or cloudpickle.dumps is called, functions also need + # to be saved within the same invocation of + # cloudpickle.dump/cloudpickle.dumps (for example: + # cloudpickle.dumps([f1, f2])). There is no such limitation when using + # cloudpickle.Pickler.dump, as long as the multiple invocations are + # bound to the same cloudpickle.Pickler instance. + base_globals = self.globals_ref.setdefault(id(func.__globals__), {}) + + if base_globals == {}: + # Add module attributes used to resolve relative imports + # instructions inside func. + for k in ["__package__", "__name__", "__path__", "__file__"]: + if k in func.__globals__: + base_globals[k] = func.__globals__[k] + + # Do not bind the free variables before the function is created to + # avoid infinite recursion. + if func.__closure__ is None: + closure = None + else: + closure = tuple(_make_empty_cell() for _ in range(len(code.co_freevars))) + + return code, base_globals, None, None, closure + + def dump(self, obj): + try: + return super().dump(obj) + except RecursionError as e: + msg = "Could not pickle object as excessively deep recursion required." + raise pickle.PicklingError(msg) from e + + def __init__(self, file, protocol=None, buffer_callback=None): + if protocol is None: + protocol = DEFAULT_PROTOCOL + super().__init__(file, protocol=protocol, buffer_callback=buffer_callback) + # map functions __globals__ attribute ids, to ensure that functions + # sharing the same global namespace at pickling time also share + # their global namespace at unpickling time. + self.globals_ref = {} + self.proto = int(protocol) + + if not PYPY: + # pickle.Pickler is the C implementation of the CPython pickler and + # therefore we rely on reduce_override method to customize the pickler + # behavior. + + # `cloudpickle.Pickler.dispatch` is only left for backward + # compatibility - note that when using protocol 5, + # `cloudpickle.Pickler.dispatch` is not an extension of + # `pickle._Pickler.dispatch` dictionary, because `cloudpickle.Pickler` + # subclasses the C-implemented `pickle.Pickler`, which does not expose + # a `dispatch` attribute. Earlier versions of `cloudpickle.Pickler` + # used `cloudpickle.Pickler.dispatch` as a class-level attribute + # storing all reducers implemented by cloudpickle, but the attribute + # name was not a great choice given because it would collide with a + # similarly named attribute in the pure-Python `pickle._Pickler` + # implementation in the standard library. + dispatch = dispatch_table + + # Implementation of the reducer_override callback, in order to + # efficiently serialize dynamic functions and classes by subclassing + # the C-implemented `pickle.Pickler`. + # TODO: decorrelate reducer_override (which is tied to CPython's + # implementation - would it make sense to backport it to pypy? - and + # pickle's protocol 5 which is implementation agnostic. Currently, the + # availability of both notions coincide on CPython's pickle, but it may + # not be the case anymore when pypy implements protocol 5. + + def reducer_override(self, obj): + """Type-agnostic reducing callback for function and classes. + + For performance reasons, subclasses of the C `pickle.Pickler` class + cannot register custom reducers for functions and classes in the + dispatch_table attribute. Reducers for such types must instead + implemented via the special `reducer_override` method. + + Note that this method will be called for any object except a few + builtin-types (int, lists, dicts etc.), which differs from reducers + in the Pickler's dispatch_table, each of them being invoked for + objects of a specific type only. + + This property comes in handy for classes: although most classes are + instances of the ``type`` metaclass, some of them can be instances + of other custom metaclasses (such as enum.EnumMeta for example). In + particular, the metaclass will likely not be known in advance, and + thus cannot be special-cased using an entry in the dispatch_table. + reducer_override, among other things, allows us to register a + reducer that will be called for any class, independently of its + type. + + Notes: + + * reducer_override has the priority over dispatch_table-registered + reducers. + * reducer_override can be used to fix other limitations of + cloudpickle for other types that suffered from type-specific + reducers, such as Exceptions. See + https://github.com/cloudpipe/cloudpickle/issues/248 + """ + t = type(obj) + try: + is_anyclass = issubclass(t, type) + except TypeError: # t is not a class (old Boost; see SF #502085) + is_anyclass = False + + if is_anyclass: + return _class_reduce(obj) + elif isinstance(obj, types.FunctionType): + return self._function_reduce(obj) + else: + # fallback to save_global, including the Pickler's + # dispatch_table + return NotImplemented + + else: + # When reducer_override is not available, hack the pure-Python + # Pickler's types.FunctionType and type savers. Note: the type saver + # must override Pickler.save_global, because pickle.py contains a + # hard-coded call to save_global when pickling meta-classes. + dispatch = pickle.Pickler.dispatch.copy() + + def _save_reduce_pickle5( + self, + func, + args, + state=None, + listitems=None, + dictitems=None, + state_setter=None, + obj=None, + ): + save = self.save + write = self.write + self.save_reduce( + func, + args, + state=None, + listitems=listitems, + dictitems=dictitems, + obj=obj, + ) + # backport of the Python 3.8 state_setter pickle operations + save(state_setter) + save(obj) # simple BINGET opcode as obj is already memoized. + save(state) + write(pickle.TUPLE2) + # Trigger a state_setter(obj, state) function call. + write(pickle.REDUCE) + # The purpose of state_setter is to carry-out an + # inplace modification of obj. We do not care about what the + # method might return, so its output is eventually removed from + # the stack. + write(pickle.POP) + + def save_global(self, obj, name=None, pack=struct.pack): + """Main dispatch method. + + The name of this method is somewhat misleading: all types get + dispatched here. + """ + if obj is type(None): # noqa + return self.save_reduce(type, (None,), obj=obj) + elif obj is type(Ellipsis): + return self.save_reduce(type, (Ellipsis,), obj=obj) + elif obj is type(NotImplemented): + return self.save_reduce(type, (NotImplemented,), obj=obj) + elif obj in _BUILTIN_TYPE_NAMES: + return self.save_reduce( + _builtin_type, (_BUILTIN_TYPE_NAMES[obj],), obj=obj + ) + + if name is not None: + super().save_global(obj, name=name) + elif not _should_pickle_by_reference(obj, name=name): + self._save_reduce_pickle5(*_dynamic_class_reduce(obj), obj=obj) + else: + super().save_global(obj, name=name) + + dispatch[type] = save_global + + def save_function(self, obj, name=None): + """Registered with the dispatch to handle all function types. + + Determines what kind of function obj is (e.g. lambda, defined at + interactive prompt, etc) and handles the pickling appropriately. + """ + if _should_pickle_by_reference(obj, name=name): + return super().save_global(obj, name=name) + elif PYPY and isinstance(obj.__code__, builtin_code_type): + return self.save_pypy_builtin_func(obj) + else: + return self._save_reduce_pickle5( + *self._dynamic_function_reduce(obj), obj=obj + ) + + def save_pypy_builtin_func(self, obj): + """Save pypy equivalent of builtin functions. + + PyPy does not have the concept of builtin-functions. Instead, + builtin-functions are simple function instances, but with a + builtin-code attribute. + Most of the time, builtin functions should be pickled by attribute. + But PyPy has flaky support for __qualname__, so some builtin + functions such as float.__new__ will be classified as dynamic. For + this reason only, we created this special routine. Because + builtin-functions are not expected to have closure or globals, + there is no additional hack (compared the one already implemented + in pickle) to protect ourselves from reference cycles. A simple + (reconstructor, newargs, obj.__dict__) tuple is save_reduced. Note + also that PyPy improved their support for __qualname__ in v3.6, so + this routing should be removed when cloudpickle supports only PyPy + 3.6 and later. + """ + rv = ( + types.FunctionType, + (obj.__code__, {}, obj.__name__, obj.__defaults__, obj.__closure__), + obj.__dict__, + ) + self.save_reduce(*rv, obj=obj) + + dispatch[types.FunctionType] = save_function + + +# Shorthands similar to pickle.dump/pickle.dumps + + +def dump(obj, file, protocol=None, buffer_callback=None): + """Serialize obj as bytes streamed into file + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python (although this is not always + guaranteed to work because cloudpickle relies on some internal + implementation details that can change from one Python version to the + next). + """ + Pickler(file, protocol=protocol, buffer_callback=buffer_callback).dump(obj) + + +def dumps(obj, protocol=None, buffer_callback=None): + """Serialize obj as a string of bytes allocated in memory + + protocol defaults to cloudpickle.DEFAULT_PROTOCOL which is an alias to + pickle.HIGHEST_PROTOCOL. This setting favors maximum communication + speed between processes running the same Python version. + + Set protocol=pickle.DEFAULT_PROTOCOL instead if you need to ensure + compatibility with older versions of Python (although this is not always + guaranteed to work because cloudpickle relies on some internal + implementation details that can change from one Python version to the + next). + """ + with io.BytesIO() as file: + cp = Pickler(file, protocol=protocol, buffer_callback=buffer_callback) + cp.dump(obj) + return file.getvalue() + + +# Include pickles unloading functions in this namespace for convenience. +load, loads = pickle.load, pickle.loads + +# Backward compat alias. +CloudPickler = Pickler diff --git a/lib/python3.12/site-packages/cloudpickle/cloudpickle_fast.py b/lib/python3.12/site-packages/cloudpickle/cloudpickle_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..20280f0ca354a691861ab6f17821bbeb04632003 --- /dev/null +++ b/lib/python3.12/site-packages/cloudpickle/cloudpickle_fast.py @@ -0,0 +1,14 @@ +"""Compatibility module. + +It can be necessary to load files generated by previous versions of cloudpickle +that rely on symbols being defined under the `cloudpickle.cloudpickle_fast` +namespace. + +See: tests/test_backward_compat.py +""" + +from . import cloudpickle + + +def __getattr__(name): + return getattr(cloudpickle, name) diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/INSTALLER b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/METADATA b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..b2ef543b933135817ff56c1e5354c1c90e29f59a --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/METADATA @@ -0,0 +1,102 @@ +Metadata-Version: 2.1 +Name: ninja +Version: 1.13.0 +Summary: Ninja is a small build system with a focus on speed +Keywords: build,c++,cross-compilation,cross-platform,fortran,ninja +Author-Email: Jean-Christophe Fillion-Robin , Henry Schreiner +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: Apache Software License +Classifier: License :: OSI Approved :: BSD License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: C +Classifier: Programming Language :: C++ +Classifier: Programming Language :: Fortran +Classifier: Programming Language :: Python +Classifier: Topic :: Software Development :: Build Tools +Classifier: Typing :: Typed +Project-URL: Bug Tracker, https://github.com/scikit-build/ninja-python-distributions/issues +Project-URL: Documentation, https://github.com/scikit-build/ninja-python-distributions#readme +Project-URL: Download, https://github.com/ninja-build/ninja/releases +Project-URL: Homepage, http://ninja-build.org/ +Project-URL: Mailing list, https://groups.google.com/forum/#!forum/scikit-build +Project-URL: Source Code, https://github.com/scikit-build/ninja-python-distributions +Requires-Python: >=3.8 +Description-Content-Type: text/x-rst + +========================== +Ninja Python Distributions +========================== + +`Ninja `_ is a small build system with a focus on speed. + +The latest Ninja python wheels provide `ninja 1.13.0.gd74ef.kitware.jobserver-pipe-1 `_ executable +and `ninja_syntax.py` for generating `.ninja` files. + +.. image:: https://raw.githubusercontent.com/scikit-build/ninja-python-distributions/master/ninja-python-distributions-logo.png + +Latest Release +-------------- + +.. table:: + + +----------------------------------------------------------------------+---------------------------------------------------------------------------+ + | Versions | Downloads | + +======================================================================+===========================================================================+ + | .. image:: https://img.shields.io/pypi/v/ninja.svg | .. image:: https://img.shields.io/badge/downloads-2535k%20total-green.svg | + | :target: https://pypi.python.org/pypi/ninja | :target: https://pypi.python.org/pypi/ninja | + +----------------------------------------------------------------------+---------------------------------------------------------------------------+ + +Build Status +------------ + +.. table:: + + +---------------+-------------------------------------------------------------------------------------------------------------+ + | | GitHub Actions (Windows, macOS, Linux) | + +===============+=============================================================================================================+ + | PyPI | .. image:: https://github.com/scikit-build/ninja-python-distributions/actions/workflows/build.yml/badge.svg | + | | :target: https://github.com/scikit-build/ninja-python-distributions/actions/workflows/build.yml | + +---------------+-------------------------------------------------------------------------------------------------------------+ + +Maintainers +----------- + +* `How to update ninja version ? `_ + +* `How to make a release ? `_ + + +Miscellaneous +------------- + +* Documentation: https://github.com/scikit-build/ninja-python-distributions#readme +* Source code: https://github.com/scikit-build/ninja-python-distributions +* Mailing list: https://groups.google.com/forum/#!forum/scikit-build + +Python Version Support +---------------------- + +Versions after 1.11.1.1 no longer support Python 2-3.6, and require manylinux2010+ on linux. +Versions after 1.13 no longer support Python 3.7, and require manylinux2014+/musllinux_1_2+ on linux. + +License +------- + +This project is maintained by Jean-Christophe Fillion-Robin from Kitware Inc. +It is covered by the `Apache License, Version 2.0 `_. + +Ninja is also distributed under the `Apache License, Version 2.0 `_. +For more information about Ninja, visit https://ninja-build.org + +Logo was originally created by Libby Rose from Kitware Inc. +It is covered by `CC BY 4.0 `_. + + +History +------- + +ninja-python-distributions was initially developed in November 2016 by +Jean-Christophe Fillion-Robin to facilitate the distribution of project using +`scikit-build `_ and depending on CMake +and Ninja. diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/RECORD b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..ebd353f47262eaaebbadc74eb6b1aec8b26697a9 --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/RECORD @@ -0,0 +1,18 @@ +../../../bin/ninja,sha256=aW-WKKednOUDFM-VVtfNGh0exSuP1Sgo9vnbFxlWW2c,372384 +ninja-1.13.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +ninja-1.13.0.dist-info/METADATA,sha256=jXb8Tjgs7c0ivK6qdIrzmz-nkxtyRCOW0VyxgWWdCk4,5148 +ninja-1.13.0.dist-info/RECORD,, +ninja-1.13.0.dist-info/WHEEL,sha256=zzOwTeuxXsHOT9QV_vcZk1hX9KJAUdqJD7e-2jGJxLA,150 +ninja-1.13.0.dist-info/licenses/AUTHORS.rst,sha256=bGE1t_Lhm2ir8S7n_jbLDohP84fpJ5sNCuxvDVsKNQg,142 +ninja-1.13.0.dist-info/licenses/LICENSE_Apache_20,sha256=c7p036pSC0mkAbXSFFmoUjoUbzt1GKgz7qXvqFEwv2g,10273 +ninja/__init__.py,sha256=taDHI20kpmjxOT72PiSyE42Pvz3KLXlwMfTZvvnRMKM,1533 +ninja/__main__.py,sha256=6iPLwHHAc2TMbojFVcUzERrzN0RvIsywuZc4KpJCg_4,100 +ninja/__pycache__/__init__.cpython-312.pyc,, +ninja/__pycache__/__main__.cpython-312.pyc,, +ninja/__pycache__/_version.cpython-312.pyc,, +ninja/__pycache__/ninja_syntax.cpython-312.pyc,, +ninja/_version.py,sha256=M85oP8JJdZ4yZHcp9qfGYLKUYvnN3kTyQosVcYPCPow,19 +ninja/_version.pyi,sha256=j5kbzfm6lOn8BzASXWjGIA1yT0OlHTWqlbyZ8Si_o0E,118 +ninja/ninja_syntax.py,sha256=RCXZ6Roda3lwnbhtQw6I6bJcpE98mXUZ9jocnvC4IQY,8148 +ninja/ninja_syntax.pyi,sha256=5IHH7N9CTKnMDDetFKHOsPK0SOSK_7Q4YtgzhIMNLE0,1576 +ninja/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/WHEEL b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..0a6a9a310a0d2edd46721b7b00182339807a1e25 --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: scikit-build-core 0.11.5 +Root-Is-Purelib: false +Tag: py3-none-manylinux_2_17_x86_64 +Tag: py3-none-manylinux2014_x86_64 + diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/AUTHORS.rst b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/AUTHORS.rst new file mode 100644 index 0000000000000000000000000000000000000000..b1b06c180f71d5a97709335823145b230576f2c5 --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/AUTHORS.rst @@ -0,0 +1,5 @@ +======= +Credits +======= + +Please see the GitHub project page at https://github.com/scikit-build/ninja-python-distributions/graphs/contributors diff --git a/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/LICENSE_Apache_20 b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/LICENSE_Apache_20 new file mode 100644 index 0000000000000000000000000000000000000000..37ec93a14fdcd0d6e525d97c0cfa6b314eaa98d8 --- /dev/null +++ b/lib/python3.12/site-packages/ninja-1.13.0.dist-info/licenses/LICENSE_Apache_20 @@ -0,0 +1,191 @@ +Apache License +Version 2.0, January 2004 +http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + +"License" shall mean the terms and conditions for use, reproduction, and +distribution as defined by Sections 1 through 9 of this document. + +"Licensor" shall mean the copyright owner or entity authorized by the copyright +owner that is granting the License. + +"Legal Entity" shall mean the union of the acting entity and all other entities +that control, are controlled by, or are under common control with that entity. +For the purposes of this definition, "control" means (i) the power, direct or +indirect, to cause the direction or management of such entity, whether by +contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the +outstanding shares, or (iii) beneficial ownership of such entity. + +"You" (or "Your") shall mean an individual or Legal Entity exercising +permissions granted by this License. + +"Source" form shall mean the preferred form for making modifications, including +but not limited to software source code, documentation source, and configuration +files. + +"Object" form shall mean any form resulting from mechanical transformation or +translation of a Source form, including but not limited to compiled object code, +generated documentation, and conversions to other media types. + +"Work" shall mean the work of authorship, whether in Source or Object form, made +available under the License, as indicated by a copyright notice that is included +in or attached to the work (an example is provided in the Appendix below). + +"Derivative Works" shall mean any work, whether in Source or Object form, that +is based on (or derived from) the Work and for which the editorial revisions, +annotations, elaborations, or other modifications represent, as a whole, an +original work of authorship. For the purposes of this License, Derivative Works +shall not include works that remain separable from, or merely link (or bind by +name) to the interfaces of, the Work and Derivative Works thereof. + +"Contribution" shall mean any work of authorship, including the original version +of the Work and any modifications or additions to that Work or Derivative Works +thereof, that is intentionally submitted to Licensor for inclusion in the Work +by the copyright owner or by an individual or Legal Entity authorized to submit +on behalf of the copyright owner. For the purposes of this definition, +"submitted" means any form of electronic, verbal, or written communication sent +to the Licensor or its representatives, including but not limited to +communication on electronic mailing lists, source code control systems, and +issue tracking systems that are managed by, or on behalf of, the Licensor for +the purpose of discussing and improving the Work, but excluding communication +that is conspicuously marked or otherwise designated in writing by the copyright +owner as "Not a Contribution." + +"Contributor" shall mean Licensor and any individual or Legal Entity on behalf +of whom a Contribution has been received by Licensor and subsequently +incorporated within the Work. + +2. Grant of Copyright License. + +Subject to the terms and conditions of this License, each Contributor hereby +grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, +irrevocable copyright license to reproduce, prepare Derivative Works of, +publicly display, publicly perform, sublicense, and distribute the Work and such +Derivative Works in Source or Object form. + +3. Grant of Patent License. + +Subject to the terms and conditions of this License, each Contributor hereby +grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, +irrevocable (except as stated in this section) patent license to make, have +made, use, offer to sell, sell, import, and otherwise transfer the Work, where +such license applies only to those patent claims licensable by such Contributor +that are necessarily infringed by their Contribution(s) alone or by combination +of their Contribution(s) with the Work to which such Contribution(s) was +submitted. If You institute patent litigation against any entity (including a +cross-claim or counterclaim in a lawsuit) alleging that the Work or a +Contribution incorporated within the Work constitutes direct or contributory +patent infringement, then any patent licenses granted to You under this License +for that Work shall terminate as of the date such litigation is filed. + +4. Redistribution. + +You may reproduce and distribute copies of the Work or Derivative Works thereof +in any medium, with or without modifications, and in Source or Object form, +provided that You meet the following conditions: + +You must give any other recipients of the Work or Derivative Works a copy of +this License; and +You must cause any modified files to carry prominent notices stating that You +changed the files; and +You must retain, in the Source form of any Derivative Works that You distribute, +all copyright, patent, trademark, and attribution notices from the Source form +of the Work, excluding those notices that do not pertain to any part of the +Derivative Works; and +If the Work includes a "NOTICE" text file as part of its distribution, then any +Derivative Works that You distribute must include a readable copy of the +attribution notices contained within such NOTICE file, excluding those notices +that do not pertain to any part of the Derivative Works, in at least one of the +following places: within a NOTICE text file distributed as part of the +Derivative Works; within the Source form or documentation, if provided along +with the Derivative Works; or, within a display generated by the Derivative +Works, if and wherever such third-party notices normally appear. The contents of +the NOTICE file are for informational purposes only and do not modify the +License. You may add Your own attribution notices within Derivative Works that +You distribute, alongside or as an addendum to the NOTICE text from the Work, +provided that such additional attribution notices cannot be construed as +modifying the License. +You may add Your own copyright statement to Your modifications and may provide +additional or different license terms and conditions for use, reproduction, or +distribution of Your modifications, or for any such Derivative Works as a whole, +provided Your use, reproduction, and distribution of the Work otherwise complies +with the conditions stated in this License. + +5. Submission of Contributions. + +Unless You explicitly state otherwise, any Contribution intentionally submitted +for inclusion in the Work by You to the Licensor shall be under the terms and +conditions of this License, without any additional terms or conditions. +Notwithstanding the above, nothing herein shall supersede or modify the terms of +any separate license agreement you may have executed with Licensor regarding +such Contributions. + +6. Trademarks. + +This License does not grant permission to use the trade names, trademarks, +service marks, or product names of the Licensor, except as required for +reasonable and customary use in describing the origin of the Work and +reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. + +Unless required by applicable law or agreed to in writing, Licensor provides the +Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, +including, without limitation, any warranties or conditions of TITLE, +NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are +solely responsible for determining the appropriateness of using or +redistributing the Work and assume any risks associated with Your exercise of +permissions under this License. + +8. Limitation of Liability. + +In no event and under no legal theory, whether in tort (including negligence), +contract, or otherwise, unless required by applicable law (such as deliberate +and grossly negligent acts) or agreed to in writing, shall any Contributor be +liable to You for damages, including any direct, indirect, special, incidental, +or consequential damages of any character arising as a result of this License or +out of the use or inability to use the Work (including but not limited to +damages for loss of goodwill, work stoppage, computer failure or malfunction, or +any and all other commercial damages or losses), even if such Contributor has +been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. + +While redistributing the Work or Derivative Works thereof, You may choose to +offer, and charge a fee for, acceptance of support, warranty, indemnity, or +other liability obligations and/or rights consistent with this License. However, +in accepting such obligations, You may act only on Your own behalf and on Your +sole responsibility, not on behalf of any other Contributor, and only if You +agree to indemnify, defend, and hold each Contributor harmless for any liability +incurred by, or claims asserted against, such Contributor by reason of your +accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work + +To apply the Apache License to your work, attach the following boilerplate +notice, with the fields enclosed by brackets "[]" replaced with your own +identifying information. (Don't include the brackets!) The text should be +enclosed in the appropriate comment syntax for the file format. We also +recommend that a file or class name and description of purpose be included on +the same "printed page" as the copyright notice for easier identification within +third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed 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. diff --git a/lib/python3.12/site-packages/numpy/__config__.py b/lib/python3.12/site-packages/numpy/__config__.py new file mode 100644 index 0000000000000000000000000000000000000000..361cf053ddf1cf04a73f117f6bcffc7f928f6349 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/__config__.py @@ -0,0 +1,162 @@ +# This file is generated by numpy's build process +# It contains system_info results at the time of building this package. +from enum import Enum +from numpy.core._multiarray_umath import ( + __cpu_features__, + __cpu_baseline__, + __cpu_dispatch__, +) + +__all__ = ["show"] +_built_with_meson = True + + +class DisplayModes(Enum): + stdout = "stdout" + dicts = "dicts" + + +def _cleanup(d): + """ + Removes empty values in a `dict` recursively + This ensures we remove values that Meson could not provide to CONFIG + """ + if isinstance(d, dict): + return {k: _cleanup(v) for k, v in d.items() if v and _cleanup(v)} + else: + return d + + +CONFIG = _cleanup( + { + "Compilers": { + "c": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.2.1", + "commands": r"cc", + "args": r"-fno-strict-aliasing", + "linker args": r"-Wl,--strip-debug, -fno-strict-aliasing", + }, + "cython": { + "name": "cython", + "linker": r"cython", + "version": "3.0.8", + "commands": r"cython", + "args": r"", + "linker args": r"", + }, + "c++": { + "name": "gcc", + "linker": r"ld.bfd", + "version": "10.2.1", + "commands": r"c++", + "args": r"", + "linker args": r"-Wl,--strip-debug", + }, + }, + "Machine Information": { + "host": { + "cpu": "x86_64", + "family": "x86_64", + "endian": "little", + "system": "linux", + }, + "build": { + "cpu": "x86_64", + "family": "x86_64", + "endian": "little", + "system": "linux", + }, + "cross-compiled": bool("False".lower().replace("false", "")), + }, + "Build Dependencies": { + "blas": { + "name": "openblas64", + "found": bool("True".lower().replace("false", "")), + "version": "0.3.23.dev", + "detection method": "pkgconfig", + "include directory": r"/usr/local/include", + "lib directory": r"/usr/local/lib", + "openblas configuration": r"USE_64BITINT=1 DYNAMIC_ARCH=1 DYNAMIC_OLDER= NO_CBLAS= NO_LAPACK= NO_LAPACKE= NO_AFFINITY=1 USE_OPENMP= HASWELL MAX_THREADS=2", + "pc file directory": r"/usr/local/lib/pkgconfig", + }, + "lapack": { + "name": "dep140551260102944", + "found": bool("True".lower().replace("false", "")), + "version": "1.26.4", + "detection method": "internal", + "include directory": r"unknown", + "lib directory": r"unknown", + "openblas configuration": r"unknown", + "pc file directory": r"unknown", + }, + }, + "Python Information": { + "path": r"/opt/python/cp312-cp312/bin/python", + "version": "3.12", + }, + "SIMD Extensions": { + "baseline": __cpu_baseline__, + "found": [ + feature for feature in __cpu_dispatch__ if __cpu_features__[feature] + ], + "not found": [ + feature for feature in __cpu_dispatch__ if not __cpu_features__[feature] + ], + }, + } +) + + +def _check_pyyaml(): + import yaml + + return yaml + + +def show(mode=DisplayModes.stdout.value): + """ + Show libraries and system information on which NumPy was built + and is being used + + Parameters + ---------- + mode : {`'stdout'`, `'dicts'`}, optional. + Indicates how to display the config information. + `'stdout'` prints to console, `'dicts'` returns a dictionary + of the configuration. + + Returns + ------- + out : {`dict`, `None`} + If mode is `'dicts'`, a dict is returned, else None + + See Also + -------- + get_include : Returns the directory containing NumPy C + header files. + + Notes + ----- + 1. The `'stdout'` mode will give more readable + output if ``pyyaml`` is installed + + """ + if mode == DisplayModes.stdout.value: + try: # Non-standard library, check import + yaml = _check_pyyaml() + + print(yaml.dump(CONFIG)) + except ModuleNotFoundError: + import warnings + import json + + warnings.warn("Install `pyyaml` for better output", stacklevel=1) + print(json.dumps(CONFIG, indent=2)) + elif mode == DisplayModes.dicts.value: + return CONFIG + else: + raise AttributeError( + f"Invalid `mode`, use one of: {', '.join([e.value for e in DisplayModes])}" + ) diff --git a/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd b/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd new file mode 100644 index 0000000000000000000000000000000000000000..1409514f7a845501a7787f6acb3a5570502d330d --- /dev/null +++ b/lib/python3.12/site-packages/numpy/__init__.cython-30.pxd @@ -0,0 +1,1050 @@ +# NumPy static imports for Cython >= 3.0 +# +# If any of the PyArray_* functions are called, import_array must be +# called first. This is done automatically by Cython 3.0+ if a call +# is not detected inside of the module. +# +# Author: Dag Sverre Seljebotn +# + +from cpython.ref cimport Py_INCREF +from cpython.object cimport PyObject, PyTypeObject, PyObject_TypeCheck +cimport libc.stdio as stdio + + +cdef extern from *: + # Leave a marker that the NumPy declarations came from NumPy itself and not from Cython. + # See https://github.com/cython/cython/issues/3573 + """ + /* Using NumPy API declarations from "numpy/__init__.cython-30.pxd" */ + """ + + +cdef extern from "Python.h": + ctypedef int Py_intptr_t + +cdef extern from "numpy/arrayobject.h": + ctypedef Py_intptr_t npy_intp + ctypedef size_t npy_uintp + + cdef enum NPY_TYPES: + NPY_BOOL + NPY_BYTE + NPY_UBYTE + NPY_SHORT + NPY_USHORT + NPY_INT + NPY_UINT + NPY_LONG + NPY_ULONG + NPY_LONGLONG + NPY_ULONGLONG + NPY_FLOAT + NPY_DOUBLE + NPY_LONGDOUBLE + NPY_CFLOAT + NPY_CDOUBLE + NPY_CLONGDOUBLE + NPY_OBJECT + NPY_STRING + NPY_UNICODE + NPY_VOID + NPY_DATETIME + NPY_TIMEDELTA + NPY_NTYPES + NPY_NOTYPE + + NPY_INT8 + NPY_INT16 + NPY_INT32 + NPY_INT64 + NPY_INT128 + NPY_INT256 + NPY_UINT8 + NPY_UINT16 + NPY_UINT32 + NPY_UINT64 + NPY_UINT128 + NPY_UINT256 + NPY_FLOAT16 + NPY_FLOAT32 + NPY_FLOAT64 + NPY_FLOAT80 + NPY_FLOAT96 + NPY_FLOAT128 + NPY_FLOAT256 + NPY_COMPLEX32 + NPY_COMPLEX64 + NPY_COMPLEX128 + NPY_COMPLEX160 + NPY_COMPLEX192 + NPY_COMPLEX256 + NPY_COMPLEX512 + + NPY_INTP + + ctypedef enum NPY_ORDER: + NPY_ANYORDER + NPY_CORDER + NPY_FORTRANORDER + NPY_KEEPORDER + + ctypedef enum NPY_CASTING: + NPY_NO_CASTING + NPY_EQUIV_CASTING + NPY_SAFE_CASTING + NPY_SAME_KIND_CASTING + NPY_UNSAFE_CASTING + + ctypedef enum NPY_CLIPMODE: + NPY_CLIP + NPY_WRAP + NPY_RAISE + + ctypedef enum NPY_SCALARKIND: + NPY_NOSCALAR, + NPY_BOOL_SCALAR, + NPY_INTPOS_SCALAR, + NPY_INTNEG_SCALAR, + NPY_FLOAT_SCALAR, + NPY_COMPLEX_SCALAR, + NPY_OBJECT_SCALAR + + ctypedef enum NPY_SORTKIND: + NPY_QUICKSORT + NPY_HEAPSORT + NPY_MERGESORT + + ctypedef enum NPY_SEARCHSIDE: + NPY_SEARCHLEFT + NPY_SEARCHRIGHT + + enum: + # DEPRECATED since NumPy 1.7 ! Do not use in new code! + NPY_C_CONTIGUOUS + NPY_F_CONTIGUOUS + NPY_CONTIGUOUS + NPY_FORTRAN + NPY_OWNDATA + NPY_FORCECAST + NPY_ENSURECOPY + NPY_ENSUREARRAY + NPY_ELEMENTSTRIDES + NPY_ALIGNED + NPY_NOTSWAPPED + NPY_WRITEABLE + NPY_ARR_HAS_DESCR + + NPY_BEHAVED + NPY_BEHAVED_NS + NPY_CARRAY + NPY_CARRAY_RO + NPY_FARRAY + NPY_FARRAY_RO + NPY_DEFAULT + + NPY_IN_ARRAY + NPY_OUT_ARRAY + NPY_INOUT_ARRAY + NPY_IN_FARRAY + NPY_OUT_FARRAY + NPY_INOUT_FARRAY + + NPY_UPDATE_ALL + + enum: + # Added in NumPy 1.7 to replace the deprecated enums above. + NPY_ARRAY_C_CONTIGUOUS + NPY_ARRAY_F_CONTIGUOUS + NPY_ARRAY_OWNDATA + NPY_ARRAY_FORCECAST + NPY_ARRAY_ENSURECOPY + NPY_ARRAY_ENSUREARRAY + NPY_ARRAY_ELEMENTSTRIDES + NPY_ARRAY_ALIGNED + NPY_ARRAY_NOTSWAPPED + NPY_ARRAY_WRITEABLE + NPY_ARRAY_WRITEBACKIFCOPY + + NPY_ARRAY_BEHAVED + NPY_ARRAY_BEHAVED_NS + NPY_ARRAY_CARRAY + NPY_ARRAY_CARRAY_RO + NPY_ARRAY_FARRAY + NPY_ARRAY_FARRAY_RO + NPY_ARRAY_DEFAULT + + NPY_ARRAY_IN_ARRAY + NPY_ARRAY_OUT_ARRAY + NPY_ARRAY_INOUT_ARRAY + NPY_ARRAY_IN_FARRAY + NPY_ARRAY_OUT_FARRAY + NPY_ARRAY_INOUT_FARRAY + + NPY_ARRAY_UPDATE_ALL + + cdef enum: + NPY_MAXDIMS + + npy_intp NPY_MAX_ELSIZE + + ctypedef void (*PyArray_VectorUnaryFunc)(void *, void *, npy_intp, void *, void *) + + ctypedef struct PyArray_ArrayDescr: + # shape is a tuple, but Cython doesn't support "tuple shape" + # inside a non-PyObject declaration, so we have to declare it + # as just a PyObject*. + PyObject* shape + + ctypedef struct PyArray_Descr: + pass + + ctypedef class numpy.dtype [object PyArray_Descr, check_size ignore]: + # Use PyDataType_* macros when possible, however there are no macros + # for accessing some of the fields, so some are defined. + cdef PyTypeObject* typeobj + cdef char kind + cdef char type + # Numpy sometimes mutates this without warning (e.g. it'll + # sometimes change "|" to "<" in shared dtype objects on + # little-endian machines). If this matters to you, use + # PyArray_IsNativeByteOrder(dtype.byteorder) instead of + # directly accessing this field. + cdef char byteorder + cdef char flags + cdef int type_num + cdef int itemsize "elsize" + cdef int alignment + cdef object fields + cdef tuple names + # Use PyDataType_HASSUBARRAY to test whether this field is + # valid (the pointer can be NULL). Most users should access + # this field via the inline helper method PyDataType_SHAPE. + cdef PyArray_ArrayDescr* subarray + + ctypedef class numpy.flatiter [object PyArrayIterObject, check_size ignore]: + # Use through macros + pass + + ctypedef class numpy.broadcast [object PyArrayMultiIterObject, check_size ignore]: + # Use through macros + pass + + ctypedef struct PyArrayObject: + # For use in situations where ndarray can't replace PyArrayObject*, + # like PyArrayObject**. + pass + + ctypedef class numpy.ndarray [object PyArrayObject, check_size ignore]: + cdef __cythonbufferdefaults__ = {"mode": "strided"} + + # NOTE: no field declarations since direct access is deprecated since NumPy 1.7 + # Instead, we use properties that map to the corresponding C-API functions. + + @property + cdef inline PyObject* base(self) nogil: + """Returns a borrowed reference to the object owning the data/memory. + """ + return PyArray_BASE(self) + + @property + cdef inline dtype descr(self): + """Returns an owned reference to the dtype of the array. + """ + return PyArray_DESCR(self) + + @property + cdef inline int ndim(self) nogil: + """Returns the number of dimensions in the array. + """ + return PyArray_NDIM(self) + + @property + cdef inline npy_intp *shape(self) nogil: + """Returns a pointer to the dimensions/shape of the array. + The number of elements matches the number of dimensions of the array (ndim). + Can return NULL for 0-dimensional arrays. + """ + return PyArray_DIMS(self) + + @property + cdef inline npy_intp *strides(self) nogil: + """Returns a pointer to the strides of the array. + The number of elements matches the number of dimensions of the array (ndim). + """ + return PyArray_STRIDES(self) + + @property + cdef inline npy_intp size(self) nogil: + """Returns the total size (in number of elements) of the array. + """ + return PyArray_SIZE(self) + + @property + cdef inline char* data(self) nogil: + """The pointer to the data buffer as a char*. + This is provided for legacy reasons to avoid direct struct field access. + For new code that needs this access, you probably want to cast the result + of `PyArray_DATA()` instead, which returns a 'void*'. + """ + return PyArray_BYTES(self) + + ctypedef unsigned char npy_bool + + ctypedef signed char npy_byte + ctypedef signed short npy_short + ctypedef signed int npy_int + ctypedef signed long npy_long + ctypedef signed long long npy_longlong + + ctypedef unsigned char npy_ubyte + ctypedef unsigned short npy_ushort + ctypedef unsigned int npy_uint + ctypedef unsigned long npy_ulong + ctypedef unsigned long long npy_ulonglong + + ctypedef float npy_float + ctypedef double npy_double + ctypedef long double npy_longdouble + + ctypedef signed char npy_int8 + ctypedef signed short npy_int16 + ctypedef signed int npy_int32 + ctypedef signed long long npy_int64 + ctypedef signed long long npy_int96 + ctypedef signed long long npy_int128 + + ctypedef unsigned char npy_uint8 + ctypedef unsigned short npy_uint16 + ctypedef unsigned int npy_uint32 + ctypedef unsigned long long npy_uint64 + ctypedef unsigned long long npy_uint96 + ctypedef unsigned long long npy_uint128 + + ctypedef float npy_float32 + ctypedef double npy_float64 + ctypedef long double npy_float80 + ctypedef long double npy_float96 + ctypedef long double npy_float128 + + ctypedef struct npy_cfloat: + float real + float imag + + ctypedef struct npy_cdouble: + double real + double imag + + ctypedef struct npy_clongdouble: + long double real + long double imag + + ctypedef struct npy_complex64: + float real + float imag + + ctypedef struct npy_complex128: + double real + double imag + + ctypedef struct npy_complex160: + long double real + long double imag + + ctypedef struct npy_complex192: + long double real + long double imag + + ctypedef struct npy_complex256: + long double real + long double imag + + ctypedef struct PyArray_Dims: + npy_intp *ptr + int len + + int _import_array() except -1 + # A second definition so _import_array isn't marked as used when we use it here. + # Do not use - subject to change any time. + int __pyx_import_array "_import_array"() except -1 + + # + # Macros from ndarrayobject.h + # + bint PyArray_CHKFLAGS(ndarray m, int flags) nogil + bint PyArray_IS_C_CONTIGUOUS(ndarray arr) nogil + bint PyArray_IS_F_CONTIGUOUS(ndarray arr) nogil + bint PyArray_ISCONTIGUOUS(ndarray m) nogil + bint PyArray_ISWRITEABLE(ndarray m) nogil + bint PyArray_ISALIGNED(ndarray m) nogil + + int PyArray_NDIM(ndarray) nogil + bint PyArray_ISONESEGMENT(ndarray) nogil + bint PyArray_ISFORTRAN(ndarray) nogil + int PyArray_FORTRANIF(ndarray) nogil + + void* PyArray_DATA(ndarray) nogil + char* PyArray_BYTES(ndarray) nogil + + npy_intp* PyArray_DIMS(ndarray) nogil + npy_intp* PyArray_STRIDES(ndarray) nogil + npy_intp PyArray_DIM(ndarray, size_t) nogil + npy_intp PyArray_STRIDE(ndarray, size_t) nogil + + PyObject *PyArray_BASE(ndarray) nogil # returns borrowed reference! + PyArray_Descr *PyArray_DESCR(ndarray) nogil # returns borrowed reference to dtype! + PyArray_Descr *PyArray_DTYPE(ndarray) nogil # returns borrowed reference to dtype! NP 1.7+ alias for descr. + int PyArray_FLAGS(ndarray) nogil + void PyArray_CLEARFLAGS(ndarray, int flags) nogil # Added in NumPy 1.7 + void PyArray_ENABLEFLAGS(ndarray, int flags) nogil # Added in NumPy 1.7 + npy_intp PyArray_ITEMSIZE(ndarray) nogil + int PyArray_TYPE(ndarray arr) nogil + + object PyArray_GETITEM(ndarray arr, void *itemptr) + int PyArray_SETITEM(ndarray arr, void *itemptr, object obj) except -1 + + bint PyTypeNum_ISBOOL(int) nogil + bint PyTypeNum_ISUNSIGNED(int) nogil + bint PyTypeNum_ISSIGNED(int) nogil + bint PyTypeNum_ISINTEGER(int) nogil + bint PyTypeNum_ISFLOAT(int) nogil + bint PyTypeNum_ISNUMBER(int) nogil + bint PyTypeNum_ISSTRING(int) nogil + bint PyTypeNum_ISCOMPLEX(int) nogil + bint PyTypeNum_ISPYTHON(int) nogil + bint PyTypeNum_ISFLEXIBLE(int) nogil + bint PyTypeNum_ISUSERDEF(int) nogil + bint PyTypeNum_ISEXTENDED(int) nogil + bint PyTypeNum_ISOBJECT(int) nogil + + bint PyDataType_ISBOOL(dtype) nogil + bint PyDataType_ISUNSIGNED(dtype) nogil + bint PyDataType_ISSIGNED(dtype) nogil + bint PyDataType_ISINTEGER(dtype) nogil + bint PyDataType_ISFLOAT(dtype) nogil + bint PyDataType_ISNUMBER(dtype) nogil + bint PyDataType_ISSTRING(dtype) nogil + bint PyDataType_ISCOMPLEX(dtype) nogil + bint PyDataType_ISPYTHON(dtype) nogil + bint PyDataType_ISFLEXIBLE(dtype) nogil + bint PyDataType_ISUSERDEF(dtype) nogil + bint PyDataType_ISEXTENDED(dtype) nogil + bint PyDataType_ISOBJECT(dtype) nogil + bint PyDataType_HASFIELDS(dtype) nogil + bint PyDataType_HASSUBARRAY(dtype) nogil + + bint PyArray_ISBOOL(ndarray) nogil + bint PyArray_ISUNSIGNED(ndarray) nogil + bint PyArray_ISSIGNED(ndarray) nogil + bint PyArray_ISINTEGER(ndarray) nogil + bint PyArray_ISFLOAT(ndarray) nogil + bint PyArray_ISNUMBER(ndarray) nogil + bint PyArray_ISSTRING(ndarray) nogil + bint PyArray_ISCOMPLEX(ndarray) nogil + bint PyArray_ISPYTHON(ndarray) nogil + bint PyArray_ISFLEXIBLE(ndarray) nogil + bint PyArray_ISUSERDEF(ndarray) nogil + bint PyArray_ISEXTENDED(ndarray) nogil + bint PyArray_ISOBJECT(ndarray) nogil + bint PyArray_HASFIELDS(ndarray) nogil + + bint PyArray_ISVARIABLE(ndarray) nogil + + bint PyArray_SAFEALIGNEDCOPY(ndarray) nogil + bint PyArray_ISNBO(char) nogil # works on ndarray.byteorder + bint PyArray_IsNativeByteOrder(char) nogil # works on ndarray.byteorder + bint PyArray_ISNOTSWAPPED(ndarray) nogil + bint PyArray_ISBYTESWAPPED(ndarray) nogil + + bint PyArray_FLAGSWAP(ndarray, int) nogil + + bint PyArray_ISCARRAY(ndarray) nogil + bint PyArray_ISCARRAY_RO(ndarray) nogil + bint PyArray_ISFARRAY(ndarray) nogil + bint PyArray_ISFARRAY_RO(ndarray) nogil + bint PyArray_ISBEHAVED(ndarray) nogil + bint PyArray_ISBEHAVED_RO(ndarray) nogil + + + bint PyDataType_ISNOTSWAPPED(dtype) nogil + bint PyDataType_ISBYTESWAPPED(dtype) nogil + + bint PyArray_DescrCheck(object) + + bint PyArray_Check(object) + bint PyArray_CheckExact(object) + + # Cannot be supported due to out arg: + # bint PyArray_HasArrayInterfaceType(object, dtype, object, object&) + # bint PyArray_HasArrayInterface(op, out) + + + bint PyArray_IsZeroDim(object) + # Cannot be supported due to ## ## in macro: + # bint PyArray_IsScalar(object, verbatim work) + bint PyArray_CheckScalar(object) + bint PyArray_IsPythonNumber(object) + bint PyArray_IsPythonScalar(object) + bint PyArray_IsAnyScalar(object) + bint PyArray_CheckAnyScalar(object) + + ndarray PyArray_GETCONTIGUOUS(ndarray) + bint PyArray_SAMESHAPE(ndarray, ndarray) nogil + npy_intp PyArray_SIZE(ndarray) nogil + npy_intp PyArray_NBYTES(ndarray) nogil + + object PyArray_FROM_O(object) + object PyArray_FROM_OF(object m, int flags) + object PyArray_FROM_OT(object m, int type) + object PyArray_FROM_OTF(object m, int type, int flags) + object PyArray_FROMANY(object m, int type, int min, int max, int flags) + object PyArray_ZEROS(int nd, npy_intp* dims, int type, int fortran) + object PyArray_EMPTY(int nd, npy_intp* dims, int type, int fortran) + void PyArray_FILLWBYTE(object, int val) + npy_intp PyArray_REFCOUNT(object) + object PyArray_ContiguousFromAny(op, int, int min_depth, int max_depth) + unsigned char PyArray_EquivArrTypes(ndarray a1, ndarray a2) + bint PyArray_EquivByteorders(int b1, int b2) nogil + object PyArray_SimpleNew(int nd, npy_intp* dims, int typenum) + object PyArray_SimpleNewFromData(int nd, npy_intp* dims, int typenum, void* data) + #object PyArray_SimpleNewFromDescr(int nd, npy_intp* dims, dtype descr) + object PyArray_ToScalar(void* data, ndarray arr) + + void* PyArray_GETPTR1(ndarray m, npy_intp i) nogil + void* PyArray_GETPTR2(ndarray m, npy_intp i, npy_intp j) nogil + void* PyArray_GETPTR3(ndarray m, npy_intp i, npy_intp j, npy_intp k) nogil + void* PyArray_GETPTR4(ndarray m, npy_intp i, npy_intp j, npy_intp k, npy_intp l) nogil + + # Cannot be supported due to out arg + # void PyArray_DESCR_REPLACE(descr) + + + object PyArray_Copy(ndarray) + object PyArray_FromObject(object op, int type, int min_depth, int max_depth) + object PyArray_ContiguousFromObject(object op, int type, int min_depth, int max_depth) + object PyArray_CopyFromObject(object op, int type, int min_depth, int max_depth) + + object PyArray_Cast(ndarray mp, int type_num) + object PyArray_Take(ndarray ap, object items, int axis) + object PyArray_Put(ndarray ap, object items, object values) + + void PyArray_ITER_RESET(flatiter it) nogil + void PyArray_ITER_NEXT(flatiter it) nogil + void PyArray_ITER_GOTO(flatiter it, npy_intp* destination) nogil + void PyArray_ITER_GOTO1D(flatiter it, npy_intp ind) nogil + void* PyArray_ITER_DATA(flatiter it) nogil + bint PyArray_ITER_NOTDONE(flatiter it) nogil + + void PyArray_MultiIter_RESET(broadcast multi) nogil + void PyArray_MultiIter_NEXT(broadcast multi) nogil + void PyArray_MultiIter_GOTO(broadcast multi, npy_intp dest) nogil + void PyArray_MultiIter_GOTO1D(broadcast multi, npy_intp ind) nogil + void* PyArray_MultiIter_DATA(broadcast multi, npy_intp i) nogil + void PyArray_MultiIter_NEXTi(broadcast multi, npy_intp i) nogil + bint PyArray_MultiIter_NOTDONE(broadcast multi) nogil + + # Functions from __multiarray_api.h + + # Functions taking dtype and returning object/ndarray are disabled + # for now as they steal dtype references. I'm conservative and disable + # more than is probably needed until it can be checked further. + int PyArray_SetNumericOps (object) except -1 + object PyArray_GetNumericOps () + int PyArray_INCREF (ndarray) except * # uses PyArray_Item_INCREF... + int PyArray_XDECREF (ndarray) except * # uses PyArray_Item_DECREF... + void PyArray_SetStringFunction (object, int) + dtype PyArray_DescrFromType (int) + object PyArray_TypeObjectFromType (int) + char * PyArray_Zero (ndarray) + char * PyArray_One (ndarray) + #object PyArray_CastToType (ndarray, dtype, int) + int PyArray_CastTo (ndarray, ndarray) except -1 + int PyArray_CastAnyTo (ndarray, ndarray) except -1 + int PyArray_CanCastSafely (int, int) # writes errors + npy_bool PyArray_CanCastTo (dtype, dtype) # writes errors + int PyArray_ObjectType (object, int) except 0 + dtype PyArray_DescrFromObject (object, dtype) + #ndarray* PyArray_ConvertToCommonType (object, int *) + dtype PyArray_DescrFromScalar (object) + dtype PyArray_DescrFromTypeObject (object) + npy_intp PyArray_Size (object) + #object PyArray_Scalar (void *, dtype, object) + #object PyArray_FromScalar (object, dtype) + void PyArray_ScalarAsCtype (object, void *) + #int PyArray_CastScalarToCtype (object, void *, dtype) + #int PyArray_CastScalarDirect (object, dtype, void *, int) + object PyArray_ScalarFromObject (object) + #PyArray_VectorUnaryFunc * PyArray_GetCastFunc (dtype, int) + object PyArray_FromDims (int, int *, int) + #object PyArray_FromDimsAndDataAndDescr (int, int *, dtype, char *) + #object PyArray_FromAny (object, dtype, int, int, int, object) + object PyArray_EnsureArray (object) + object PyArray_EnsureAnyArray (object) + #object PyArray_FromFile (stdio.FILE *, dtype, npy_intp, char *) + #object PyArray_FromString (char *, npy_intp, dtype, npy_intp, char *) + #object PyArray_FromBuffer (object, dtype, npy_intp, npy_intp) + #object PyArray_FromIter (object, dtype, npy_intp) + object PyArray_Return (ndarray) + #object PyArray_GetField (ndarray, dtype, int) + #int PyArray_SetField (ndarray, dtype, int, object) except -1 + object PyArray_Byteswap (ndarray, npy_bool) + object PyArray_Resize (ndarray, PyArray_Dims *, int, NPY_ORDER) + int PyArray_MoveInto (ndarray, ndarray) except -1 + int PyArray_CopyInto (ndarray, ndarray) except -1 + int PyArray_CopyAnyInto (ndarray, ndarray) except -1 + int PyArray_CopyObject (ndarray, object) except -1 + object PyArray_NewCopy (ndarray, NPY_ORDER) + object PyArray_ToList (ndarray) + object PyArray_ToString (ndarray, NPY_ORDER) + int PyArray_ToFile (ndarray, stdio.FILE *, char *, char *) except -1 + int PyArray_Dump (object, object, int) except -1 + object PyArray_Dumps (object, int) + int PyArray_ValidType (int) # Cannot error + void PyArray_UpdateFlags (ndarray, int) + object PyArray_New (type, int, npy_intp *, int, npy_intp *, void *, int, int, object) + #object PyArray_NewFromDescr (type, dtype, int, npy_intp *, npy_intp *, void *, int, object) + #dtype PyArray_DescrNew (dtype) + dtype PyArray_DescrNewFromType (int) + double PyArray_GetPriority (object, double) # clears errors as of 1.25 + object PyArray_IterNew (object) + object PyArray_MultiIterNew (int, ...) + + int PyArray_PyIntAsInt (object) except? -1 + npy_intp PyArray_PyIntAsIntp (object) + int PyArray_Broadcast (broadcast) except -1 + void PyArray_FillObjectArray (ndarray, object) except * + int PyArray_FillWithScalar (ndarray, object) except -1 + npy_bool PyArray_CheckStrides (int, int, npy_intp, npy_intp, npy_intp *, npy_intp *) + dtype PyArray_DescrNewByteorder (dtype, char) + object PyArray_IterAllButAxis (object, int *) + #object PyArray_CheckFromAny (object, dtype, int, int, int, object) + #object PyArray_FromArray (ndarray, dtype, int) + object PyArray_FromInterface (object) + object PyArray_FromStructInterface (object) + #object PyArray_FromArrayAttr (object, dtype, object) + #NPY_SCALARKIND PyArray_ScalarKind (int, ndarray*) + int PyArray_CanCoerceScalar (int, int, NPY_SCALARKIND) + object PyArray_NewFlagsObject (object) + npy_bool PyArray_CanCastScalar (type, type) + #int PyArray_CompareUCS4 (npy_ucs4 *, npy_ucs4 *, register size_t) + int PyArray_RemoveSmallest (broadcast) except -1 + int PyArray_ElementStrides (object) + void PyArray_Item_INCREF (char *, dtype) except * + void PyArray_Item_XDECREF (char *, dtype) except * + object PyArray_FieldNames (object) + object PyArray_Transpose (ndarray, PyArray_Dims *) + object PyArray_TakeFrom (ndarray, object, int, ndarray, NPY_CLIPMODE) + object PyArray_PutTo (ndarray, object, object, NPY_CLIPMODE) + object PyArray_PutMask (ndarray, object, object) + object PyArray_Repeat (ndarray, object, int) + object PyArray_Choose (ndarray, object, ndarray, NPY_CLIPMODE) + int PyArray_Sort (ndarray, int, NPY_SORTKIND) except -1 + object PyArray_ArgSort (ndarray, int, NPY_SORTKIND) + object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *) + object PyArray_ArgMax (ndarray, int, ndarray) + object PyArray_ArgMin (ndarray, int, ndarray) + object PyArray_Reshape (ndarray, object) + object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER) + object PyArray_Squeeze (ndarray) + #object PyArray_View (ndarray, dtype, type) + object PyArray_SwapAxes (ndarray, int, int) + object PyArray_Max (ndarray, int, ndarray) + object PyArray_Min (ndarray, int, ndarray) + object PyArray_Ptp (ndarray, int, ndarray) + object PyArray_Mean (ndarray, int, int, ndarray) + object PyArray_Trace (ndarray, int, int, int, int, ndarray) + object PyArray_Diagonal (ndarray, int, int, int) + object PyArray_Clip (ndarray, object, object, ndarray) + object PyArray_Conjugate (ndarray, ndarray) + object PyArray_Nonzero (ndarray) + object PyArray_Std (ndarray, int, int, ndarray, int) + object PyArray_Sum (ndarray, int, int, ndarray) + object PyArray_CumSum (ndarray, int, int, ndarray) + object PyArray_Prod (ndarray, int, int, ndarray) + object PyArray_CumProd (ndarray, int, int, ndarray) + object PyArray_All (ndarray, int, ndarray) + object PyArray_Any (ndarray, int, ndarray) + object PyArray_Compress (ndarray, object, int, ndarray) + object PyArray_Flatten (ndarray, NPY_ORDER) + object PyArray_Ravel (ndarray, NPY_ORDER) + npy_intp PyArray_MultiplyList (npy_intp *, int) + int PyArray_MultiplyIntList (int *, int) + void * PyArray_GetPtr (ndarray, npy_intp*) + int PyArray_CompareLists (npy_intp *, npy_intp *, int) + #int PyArray_AsCArray (object*, void *, npy_intp *, int, dtype) + #int PyArray_As1D (object*, char **, int *, int) + #int PyArray_As2D (object*, char ***, int *, int *, int) + int PyArray_Free (object, void *) + #int PyArray_Converter (object, object*) + int PyArray_IntpFromSequence (object, npy_intp *, int) except -1 + object PyArray_Concatenate (object, int) + object PyArray_InnerProduct (object, object) + object PyArray_MatrixProduct (object, object) + object PyArray_CopyAndTranspose (object) + object PyArray_Correlate (object, object, int) + int PyArray_TypestrConvert (int, int) + #int PyArray_DescrConverter (object, dtype*) except 0 + #int PyArray_DescrConverter2 (object, dtype*) except 0 + int PyArray_IntpConverter (object, PyArray_Dims *) except 0 + #int PyArray_BufferConverter (object, chunk) except 0 + int PyArray_AxisConverter (object, int *) except 0 + int PyArray_BoolConverter (object, npy_bool *) except 0 + int PyArray_ByteorderConverter (object, char *) except 0 + int PyArray_OrderConverter (object, NPY_ORDER *) except 0 + unsigned char PyArray_EquivTypes (dtype, dtype) # clears errors + #object PyArray_Zeros (int, npy_intp *, dtype, int) + #object PyArray_Empty (int, npy_intp *, dtype, int) + object PyArray_Where (object, object, object) + object PyArray_Arange (double, double, double, int) + #object PyArray_ArangeObj (object, object, object, dtype) + int PyArray_SortkindConverter (object, NPY_SORTKIND *) except 0 + object PyArray_LexSort (object, int) + object PyArray_Round (ndarray, int, ndarray) + unsigned char PyArray_EquivTypenums (int, int) + int PyArray_RegisterDataType (dtype) except -1 + int PyArray_RegisterCastFunc (dtype, int, PyArray_VectorUnaryFunc *) except -1 + int PyArray_RegisterCanCast (dtype, int, NPY_SCALARKIND) except -1 + #void PyArray_InitArrFuncs (PyArray_ArrFuncs *) + object PyArray_IntTupleFromIntp (int, npy_intp *) + int PyArray_TypeNumFromName (char *) + int PyArray_ClipmodeConverter (object, NPY_CLIPMODE *) except 0 + #int PyArray_OutputConverter (object, ndarray*) except 0 + object PyArray_BroadcastToShape (object, npy_intp *, int) + void _PyArray_SigintHandler (int) + void* _PyArray_GetSigintBuf () + #int PyArray_DescrAlignConverter (object, dtype*) except 0 + #int PyArray_DescrAlignConverter2 (object, dtype*) except 0 + int PyArray_SearchsideConverter (object, void *) except 0 + object PyArray_CheckAxis (ndarray, int *, int) + npy_intp PyArray_OverflowMultiplyList (npy_intp *, int) + int PyArray_CompareString (char *, char *, size_t) + int PyArray_SetBaseObject(ndarray, base) except -1 # NOTE: steals a reference to base! Use "set_array_base()" instead. + + +# Typedefs that matches the runtime dtype objects in +# the numpy module. + +# The ones that are commented out needs an IFDEF function +# in Cython to enable them only on the right systems. + +ctypedef npy_int8 int8_t +ctypedef npy_int16 int16_t +ctypedef npy_int32 int32_t +ctypedef npy_int64 int64_t +#ctypedef npy_int96 int96_t +#ctypedef npy_int128 int128_t + +ctypedef npy_uint8 uint8_t +ctypedef npy_uint16 uint16_t +ctypedef npy_uint32 uint32_t +ctypedef npy_uint64 uint64_t +#ctypedef npy_uint96 uint96_t +#ctypedef npy_uint128 uint128_t + +ctypedef npy_float32 float32_t +ctypedef npy_float64 float64_t +#ctypedef npy_float80 float80_t +#ctypedef npy_float128 float128_t + +ctypedef float complex complex64_t +ctypedef double complex complex128_t + +# The int types are mapped a bit surprising -- +# numpy.int corresponds to 'l' and numpy.long to 'q' +ctypedef npy_long int_t +ctypedef npy_longlong longlong_t + +ctypedef npy_ulong uint_t +ctypedef npy_ulonglong ulonglong_t + +ctypedef npy_intp intp_t +ctypedef npy_uintp uintp_t + +ctypedef npy_double float_t +ctypedef npy_double double_t +ctypedef npy_longdouble longdouble_t + +ctypedef npy_cfloat cfloat_t +ctypedef npy_cdouble cdouble_t +ctypedef npy_clongdouble clongdouble_t + +ctypedef npy_cdouble complex_t + +cdef inline object PyArray_MultiIterNew1(a): + return PyArray_MultiIterNew(1, a) + +cdef inline object PyArray_MultiIterNew2(a, b): + return PyArray_MultiIterNew(2, a, b) + +cdef inline object PyArray_MultiIterNew3(a, b, c): + return PyArray_MultiIterNew(3, a, b, c) + +cdef inline object PyArray_MultiIterNew4(a, b, c, d): + return PyArray_MultiIterNew(4, a, b, c, d) + +cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): + return PyArray_MultiIterNew(5, a, b, c, d, e) + +cdef inline tuple PyDataType_SHAPE(dtype d): + if PyDataType_HASSUBARRAY(d): + return d.subarray.shape + else: + return () + + +cdef extern from "numpy/ndarrayobject.h": + PyTypeObject PyTimedeltaArrType_Type + PyTypeObject PyDatetimeArrType_Type + ctypedef int64_t npy_timedelta + ctypedef int64_t npy_datetime + +cdef extern from "numpy/ndarraytypes.h": + ctypedef struct PyArray_DatetimeMetaData: + NPY_DATETIMEUNIT base + int64_t num + +cdef extern from "numpy/arrayscalars.h": + + # abstract types + ctypedef class numpy.generic [object PyObject]: + pass + ctypedef class numpy.number [object PyObject]: + pass + ctypedef class numpy.integer [object PyObject]: + pass + ctypedef class numpy.signedinteger [object PyObject]: + pass + ctypedef class numpy.unsignedinteger [object PyObject]: + pass + ctypedef class numpy.inexact [object PyObject]: + pass + ctypedef class numpy.floating [object PyObject]: + pass + ctypedef class numpy.complexfloating [object PyObject]: + pass + ctypedef class numpy.flexible [object PyObject]: + pass + ctypedef class numpy.character [object PyObject]: + pass + + ctypedef struct PyDatetimeScalarObject: + # PyObject_HEAD + npy_datetime obval + PyArray_DatetimeMetaData obmeta + + ctypedef struct PyTimedeltaScalarObject: + # PyObject_HEAD + npy_timedelta obval + PyArray_DatetimeMetaData obmeta + + ctypedef enum NPY_DATETIMEUNIT: + NPY_FR_Y + NPY_FR_M + NPY_FR_W + NPY_FR_D + NPY_FR_B + NPY_FR_h + NPY_FR_m + NPY_FR_s + NPY_FR_ms + NPY_FR_us + NPY_FR_ns + NPY_FR_ps + NPY_FR_fs + NPY_FR_as + NPY_FR_GENERIC + + +# +# ufunc API +# + +cdef extern from "numpy/ufuncobject.h": + + ctypedef void (*PyUFuncGenericFunction) (char **, npy_intp *, npy_intp *, void *) + + ctypedef class numpy.ufunc [object PyUFuncObject, check_size ignore]: + cdef: + int nin, nout, nargs + int identity + PyUFuncGenericFunction *functions + void **data + int ntypes + int check_return + char *name + char *types + char *doc + void *ptr + PyObject *obj + PyObject *userloops + + cdef enum: + PyUFunc_Zero + PyUFunc_One + PyUFunc_None + UFUNC_ERR_IGNORE + UFUNC_ERR_WARN + UFUNC_ERR_RAISE + UFUNC_ERR_CALL + UFUNC_ERR_PRINT + UFUNC_ERR_LOG + UFUNC_MASK_DIVIDEBYZERO + UFUNC_MASK_OVERFLOW + UFUNC_MASK_UNDERFLOW + UFUNC_MASK_INVALID + UFUNC_SHIFT_DIVIDEBYZERO + UFUNC_SHIFT_OVERFLOW + UFUNC_SHIFT_UNDERFLOW + UFUNC_SHIFT_INVALID + UFUNC_FPE_DIVIDEBYZERO + UFUNC_FPE_OVERFLOW + UFUNC_FPE_UNDERFLOW + UFUNC_FPE_INVALID + UFUNC_ERR_DEFAULT + UFUNC_ERR_DEFAULT2 + + object PyUFunc_FromFuncAndData(PyUFuncGenericFunction *, + void **, char *, int, int, int, int, char *, char *, int) + int PyUFunc_RegisterLoopForType(ufunc, int, + PyUFuncGenericFunction, int *, void *) except -1 + void PyUFunc_f_f_As_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_f_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_g_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F_As_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_G_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f_As_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_gg_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F_As_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_GG_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_On_Om \ + (char **, npy_intp *, npy_intp *, void *) + int PyUFunc_GetPyValues \ + (char *, int *, int *, PyObject **) + int PyUFunc_checkfperr \ + (int, PyObject *, int *) + void PyUFunc_clearfperr() + int PyUFunc_getfperr() + int PyUFunc_handlefperr \ + (int, PyObject *, int, int *) except -1 + int PyUFunc_ReplaceLoopBySignature \ + (ufunc, PyUFuncGenericFunction, int *, PyUFuncGenericFunction *) + object PyUFunc_FromFuncAndDataAndSignature \ + (PyUFuncGenericFunction *, void **, char *, int, int, int, + int, char *, char *, int, char *) + + int _import_umath() except -1 + +cdef inline void set_array_base(ndarray arr, object base): + Py_INCREF(base) # important to do this before stealing the reference below! + PyArray_SetBaseObject(arr, base) + +cdef inline object get_array_base(ndarray arr): + base = PyArray_BASE(arr) + if base is NULL: + return None + return base + +# Versions of the import_* functions which are more suitable for +# Cython code. +cdef inline int import_array() except -1: + try: + __pyx_import_array() + except Exception: + raise ImportError("numpy.core.multiarray failed to import") + +cdef inline int import_umath() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + +cdef inline int import_ufunc() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + + +cdef inline bint is_timedelta64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.timedelta64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyTimedeltaArrType_Type) + + +cdef inline bint is_datetime64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.datetime64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyDatetimeArrType_Type) + + +cdef inline npy_datetime get_datetime64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy datetime64 object + + Note that to interpret this as a datetime, the corresponding unit is + also needed. That can be found using `get_datetime64_unit`. + """ + return (obj).obval + + +cdef inline npy_timedelta get_timedelta64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy timedelta64 object + """ + return (obj).obval + + +cdef inline NPY_DATETIMEUNIT get_datetime64_unit(object obj) nogil: + """ + returns the unit part of the dtype for a numpy datetime64 object. + """ + return (obj).obmeta.base diff --git a/lib/python3.12/site-packages/numpy/__init__.pxd b/lib/python3.12/site-packages/numpy/__init__.pxd new file mode 100644 index 0000000000000000000000000000000000000000..ca0a3a6c5288ac5ee384636fff3574ed257e49e4 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/__init__.pxd @@ -0,0 +1,1015 @@ +# NumPy static imports for Cython < 3.0 +# +# If any of the PyArray_* functions are called, import_array must be +# called first. +# +# Author: Dag Sverre Seljebotn +# + +DEF _buffer_format_string_len = 255 + +cimport cpython.buffer as pybuf +from cpython.ref cimport Py_INCREF +from cpython.mem cimport PyObject_Malloc, PyObject_Free +from cpython.object cimport PyObject, PyTypeObject +from cpython.buffer cimport PyObject_GetBuffer +from cpython.type cimport type +cimport libc.stdio as stdio + +cdef extern from "Python.h": + ctypedef int Py_intptr_t + bint PyObject_TypeCheck(object obj, PyTypeObject* type) + +cdef extern from "numpy/arrayobject.h": + ctypedef Py_intptr_t npy_intp + ctypedef size_t npy_uintp + + cdef enum NPY_TYPES: + NPY_BOOL + NPY_BYTE + NPY_UBYTE + NPY_SHORT + NPY_USHORT + NPY_INT + NPY_UINT + NPY_LONG + NPY_ULONG + NPY_LONGLONG + NPY_ULONGLONG + NPY_FLOAT + NPY_DOUBLE + NPY_LONGDOUBLE + NPY_CFLOAT + NPY_CDOUBLE + NPY_CLONGDOUBLE + NPY_OBJECT + NPY_STRING + NPY_UNICODE + NPY_VOID + NPY_DATETIME + NPY_TIMEDELTA + NPY_NTYPES + NPY_NOTYPE + + NPY_INT8 + NPY_INT16 + NPY_INT32 + NPY_INT64 + NPY_INT128 + NPY_INT256 + NPY_UINT8 + NPY_UINT16 + NPY_UINT32 + NPY_UINT64 + NPY_UINT128 + NPY_UINT256 + NPY_FLOAT16 + NPY_FLOAT32 + NPY_FLOAT64 + NPY_FLOAT80 + NPY_FLOAT96 + NPY_FLOAT128 + NPY_FLOAT256 + NPY_COMPLEX32 + NPY_COMPLEX64 + NPY_COMPLEX128 + NPY_COMPLEX160 + NPY_COMPLEX192 + NPY_COMPLEX256 + NPY_COMPLEX512 + + NPY_INTP + + ctypedef enum NPY_ORDER: + NPY_ANYORDER + NPY_CORDER + NPY_FORTRANORDER + NPY_KEEPORDER + + ctypedef enum NPY_CASTING: + NPY_NO_CASTING + NPY_EQUIV_CASTING + NPY_SAFE_CASTING + NPY_SAME_KIND_CASTING + NPY_UNSAFE_CASTING + + ctypedef enum NPY_CLIPMODE: + NPY_CLIP + NPY_WRAP + NPY_RAISE + + ctypedef enum NPY_SCALARKIND: + NPY_NOSCALAR, + NPY_BOOL_SCALAR, + NPY_INTPOS_SCALAR, + NPY_INTNEG_SCALAR, + NPY_FLOAT_SCALAR, + NPY_COMPLEX_SCALAR, + NPY_OBJECT_SCALAR + + ctypedef enum NPY_SORTKIND: + NPY_QUICKSORT + NPY_HEAPSORT + NPY_MERGESORT + + ctypedef enum NPY_SEARCHSIDE: + NPY_SEARCHLEFT + NPY_SEARCHRIGHT + + enum: + # DEPRECATED since NumPy 1.7 ! Do not use in new code! + NPY_C_CONTIGUOUS + NPY_F_CONTIGUOUS + NPY_CONTIGUOUS + NPY_FORTRAN + NPY_OWNDATA + NPY_FORCECAST + NPY_ENSURECOPY + NPY_ENSUREARRAY + NPY_ELEMENTSTRIDES + NPY_ALIGNED + NPY_NOTSWAPPED + NPY_WRITEABLE + NPY_ARR_HAS_DESCR + + NPY_BEHAVED + NPY_BEHAVED_NS + NPY_CARRAY + NPY_CARRAY_RO + NPY_FARRAY + NPY_FARRAY_RO + NPY_DEFAULT + + NPY_IN_ARRAY + NPY_OUT_ARRAY + NPY_INOUT_ARRAY + NPY_IN_FARRAY + NPY_OUT_FARRAY + NPY_INOUT_FARRAY + + NPY_UPDATE_ALL + + enum: + # Added in NumPy 1.7 to replace the deprecated enums above. + NPY_ARRAY_C_CONTIGUOUS + NPY_ARRAY_F_CONTIGUOUS + NPY_ARRAY_OWNDATA + NPY_ARRAY_FORCECAST + NPY_ARRAY_ENSURECOPY + NPY_ARRAY_ENSUREARRAY + NPY_ARRAY_ELEMENTSTRIDES + NPY_ARRAY_ALIGNED + NPY_ARRAY_NOTSWAPPED + NPY_ARRAY_WRITEABLE + NPY_ARRAY_WRITEBACKIFCOPY + + NPY_ARRAY_BEHAVED + NPY_ARRAY_BEHAVED_NS + NPY_ARRAY_CARRAY + NPY_ARRAY_CARRAY_RO + NPY_ARRAY_FARRAY + NPY_ARRAY_FARRAY_RO + NPY_ARRAY_DEFAULT + + NPY_ARRAY_IN_ARRAY + NPY_ARRAY_OUT_ARRAY + NPY_ARRAY_INOUT_ARRAY + NPY_ARRAY_IN_FARRAY + NPY_ARRAY_OUT_FARRAY + NPY_ARRAY_INOUT_FARRAY + + NPY_ARRAY_UPDATE_ALL + + cdef enum: + NPY_MAXDIMS + + npy_intp NPY_MAX_ELSIZE + + ctypedef void (*PyArray_VectorUnaryFunc)(void *, void *, npy_intp, void *, void *) + + ctypedef struct PyArray_ArrayDescr: + # shape is a tuple, but Cython doesn't support "tuple shape" + # inside a non-PyObject declaration, so we have to declare it + # as just a PyObject*. + PyObject* shape + + ctypedef struct PyArray_Descr: + pass + + ctypedef class numpy.dtype [object PyArray_Descr, check_size ignore]: + # Use PyDataType_* macros when possible, however there are no macros + # for accessing some of the fields, so some are defined. + cdef PyTypeObject* typeobj + cdef char kind + cdef char type + # Numpy sometimes mutates this without warning (e.g. it'll + # sometimes change "|" to "<" in shared dtype objects on + # little-endian machines). If this matters to you, use + # PyArray_IsNativeByteOrder(dtype.byteorder) instead of + # directly accessing this field. + cdef char byteorder + cdef char flags + cdef int type_num + cdef int itemsize "elsize" + cdef int alignment + cdef object fields + cdef tuple names + # Use PyDataType_HASSUBARRAY to test whether this field is + # valid (the pointer can be NULL). Most users should access + # this field via the inline helper method PyDataType_SHAPE. + cdef PyArray_ArrayDescr* subarray + + ctypedef class numpy.flatiter [object PyArrayIterObject, check_size ignore]: + # Use through macros + pass + + ctypedef class numpy.broadcast [object PyArrayMultiIterObject, check_size ignore]: + cdef int numiter + cdef npy_intp size, index + cdef int nd + cdef npy_intp *dimensions + cdef void **iters + + ctypedef struct PyArrayObject: + # For use in situations where ndarray can't replace PyArrayObject*, + # like PyArrayObject**. + pass + + ctypedef class numpy.ndarray [object PyArrayObject, check_size ignore]: + cdef __cythonbufferdefaults__ = {"mode": "strided"} + + cdef: + # Only taking a few of the most commonly used and stable fields. + # One should use PyArray_* macros instead to access the C fields. + char *data + int ndim "nd" + npy_intp *shape "dimensions" + npy_intp *strides + dtype descr # deprecated since NumPy 1.7 ! + PyObject* base # NOT PUBLIC, DO NOT USE ! + + + + ctypedef unsigned char npy_bool + + ctypedef signed char npy_byte + ctypedef signed short npy_short + ctypedef signed int npy_int + ctypedef signed long npy_long + ctypedef signed long long npy_longlong + + ctypedef unsigned char npy_ubyte + ctypedef unsigned short npy_ushort + ctypedef unsigned int npy_uint + ctypedef unsigned long npy_ulong + ctypedef unsigned long long npy_ulonglong + + ctypedef float npy_float + ctypedef double npy_double + ctypedef long double npy_longdouble + + ctypedef signed char npy_int8 + ctypedef signed short npy_int16 + ctypedef signed int npy_int32 + ctypedef signed long long npy_int64 + ctypedef signed long long npy_int96 + ctypedef signed long long npy_int128 + + ctypedef unsigned char npy_uint8 + ctypedef unsigned short npy_uint16 + ctypedef unsigned int npy_uint32 + ctypedef unsigned long long npy_uint64 + ctypedef unsigned long long npy_uint96 + ctypedef unsigned long long npy_uint128 + + ctypedef float npy_float32 + ctypedef double npy_float64 + ctypedef long double npy_float80 + ctypedef long double npy_float96 + ctypedef long double npy_float128 + + ctypedef struct npy_cfloat: + float real + float imag + + ctypedef struct npy_cdouble: + double real + double imag + + ctypedef struct npy_clongdouble: + long double real + long double imag + + ctypedef struct npy_complex64: + float real + float imag + + ctypedef struct npy_complex128: + double real + double imag + + ctypedef struct npy_complex160: + long double real + long double imag + + ctypedef struct npy_complex192: + long double real + long double imag + + ctypedef struct npy_complex256: + long double real + long double imag + + ctypedef struct PyArray_Dims: + npy_intp *ptr + int len + + int _import_array() except -1 + # A second definition so _import_array isn't marked as used when we use it here. + # Do not use - subject to change any time. + int __pyx_import_array "_import_array"() except -1 + + # + # Macros from ndarrayobject.h + # + bint PyArray_CHKFLAGS(ndarray m, int flags) nogil + bint PyArray_IS_C_CONTIGUOUS(ndarray arr) nogil + bint PyArray_IS_F_CONTIGUOUS(ndarray arr) nogil + bint PyArray_ISCONTIGUOUS(ndarray m) nogil + bint PyArray_ISWRITEABLE(ndarray m) nogil + bint PyArray_ISALIGNED(ndarray m) nogil + + int PyArray_NDIM(ndarray) nogil + bint PyArray_ISONESEGMENT(ndarray) nogil + bint PyArray_ISFORTRAN(ndarray) nogil + int PyArray_FORTRANIF(ndarray) nogil + + void* PyArray_DATA(ndarray) nogil + char* PyArray_BYTES(ndarray) nogil + + npy_intp* PyArray_DIMS(ndarray) nogil + npy_intp* PyArray_STRIDES(ndarray) nogil + npy_intp PyArray_DIM(ndarray, size_t) nogil + npy_intp PyArray_STRIDE(ndarray, size_t) nogil + + PyObject *PyArray_BASE(ndarray) nogil # returns borrowed reference! + PyArray_Descr *PyArray_DESCR(ndarray) nogil # returns borrowed reference to dtype! + int PyArray_FLAGS(ndarray) nogil + npy_intp PyArray_ITEMSIZE(ndarray) nogil + int PyArray_TYPE(ndarray arr) nogil + + object PyArray_GETITEM(ndarray arr, void *itemptr) + int PyArray_SETITEM(ndarray arr, void *itemptr, object obj) except -1 + + bint PyTypeNum_ISBOOL(int) nogil + bint PyTypeNum_ISUNSIGNED(int) nogil + bint PyTypeNum_ISSIGNED(int) nogil + bint PyTypeNum_ISINTEGER(int) nogil + bint PyTypeNum_ISFLOAT(int) nogil + bint PyTypeNum_ISNUMBER(int) nogil + bint PyTypeNum_ISSTRING(int) nogil + bint PyTypeNum_ISCOMPLEX(int) nogil + bint PyTypeNum_ISPYTHON(int) nogil + bint PyTypeNum_ISFLEXIBLE(int) nogil + bint PyTypeNum_ISUSERDEF(int) nogil + bint PyTypeNum_ISEXTENDED(int) nogil + bint PyTypeNum_ISOBJECT(int) nogil + + bint PyDataType_ISBOOL(dtype) nogil + bint PyDataType_ISUNSIGNED(dtype) nogil + bint PyDataType_ISSIGNED(dtype) nogil + bint PyDataType_ISINTEGER(dtype) nogil + bint PyDataType_ISFLOAT(dtype) nogil + bint PyDataType_ISNUMBER(dtype) nogil + bint PyDataType_ISSTRING(dtype) nogil + bint PyDataType_ISCOMPLEX(dtype) nogil + bint PyDataType_ISPYTHON(dtype) nogil + bint PyDataType_ISFLEXIBLE(dtype) nogil + bint PyDataType_ISUSERDEF(dtype) nogil + bint PyDataType_ISEXTENDED(dtype) nogil + bint PyDataType_ISOBJECT(dtype) nogil + bint PyDataType_HASFIELDS(dtype) nogil + bint PyDataType_HASSUBARRAY(dtype) nogil + + bint PyArray_ISBOOL(ndarray) nogil + bint PyArray_ISUNSIGNED(ndarray) nogil + bint PyArray_ISSIGNED(ndarray) nogil + bint PyArray_ISINTEGER(ndarray) nogil + bint PyArray_ISFLOAT(ndarray) nogil + bint PyArray_ISNUMBER(ndarray) nogil + bint PyArray_ISSTRING(ndarray) nogil + bint PyArray_ISCOMPLEX(ndarray) nogil + bint PyArray_ISPYTHON(ndarray) nogil + bint PyArray_ISFLEXIBLE(ndarray) nogil + bint PyArray_ISUSERDEF(ndarray) nogil + bint PyArray_ISEXTENDED(ndarray) nogil + bint PyArray_ISOBJECT(ndarray) nogil + bint PyArray_HASFIELDS(ndarray) nogil + + bint PyArray_ISVARIABLE(ndarray) nogil + + bint PyArray_SAFEALIGNEDCOPY(ndarray) nogil + bint PyArray_ISNBO(char) nogil # works on ndarray.byteorder + bint PyArray_IsNativeByteOrder(char) nogil # works on ndarray.byteorder + bint PyArray_ISNOTSWAPPED(ndarray) nogil + bint PyArray_ISBYTESWAPPED(ndarray) nogil + + bint PyArray_FLAGSWAP(ndarray, int) nogil + + bint PyArray_ISCARRAY(ndarray) nogil + bint PyArray_ISCARRAY_RO(ndarray) nogil + bint PyArray_ISFARRAY(ndarray) nogil + bint PyArray_ISFARRAY_RO(ndarray) nogil + bint PyArray_ISBEHAVED(ndarray) nogil + bint PyArray_ISBEHAVED_RO(ndarray) nogil + + + bint PyDataType_ISNOTSWAPPED(dtype) nogil + bint PyDataType_ISBYTESWAPPED(dtype) nogil + + bint PyArray_DescrCheck(object) + + bint PyArray_Check(object) + bint PyArray_CheckExact(object) + + # Cannot be supported due to out arg: + # bint PyArray_HasArrayInterfaceType(object, dtype, object, object&) + # bint PyArray_HasArrayInterface(op, out) + + + bint PyArray_IsZeroDim(object) + # Cannot be supported due to ## ## in macro: + # bint PyArray_IsScalar(object, verbatim work) + bint PyArray_CheckScalar(object) + bint PyArray_IsPythonNumber(object) + bint PyArray_IsPythonScalar(object) + bint PyArray_IsAnyScalar(object) + bint PyArray_CheckAnyScalar(object) + + ndarray PyArray_GETCONTIGUOUS(ndarray) + bint PyArray_SAMESHAPE(ndarray, ndarray) nogil + npy_intp PyArray_SIZE(ndarray) nogil + npy_intp PyArray_NBYTES(ndarray) nogil + + object PyArray_FROM_O(object) + object PyArray_FROM_OF(object m, int flags) + object PyArray_FROM_OT(object m, int type) + object PyArray_FROM_OTF(object m, int type, int flags) + object PyArray_FROMANY(object m, int type, int min, int max, int flags) + object PyArray_ZEROS(int nd, npy_intp* dims, int type, int fortran) + object PyArray_EMPTY(int nd, npy_intp* dims, int type, int fortran) + void PyArray_FILLWBYTE(object, int val) + npy_intp PyArray_REFCOUNT(object) + object PyArray_ContiguousFromAny(op, int, int min_depth, int max_depth) + unsigned char PyArray_EquivArrTypes(ndarray a1, ndarray a2) + bint PyArray_EquivByteorders(int b1, int b2) nogil + object PyArray_SimpleNew(int nd, npy_intp* dims, int typenum) + object PyArray_SimpleNewFromData(int nd, npy_intp* dims, int typenum, void* data) + #object PyArray_SimpleNewFromDescr(int nd, npy_intp* dims, dtype descr) + object PyArray_ToScalar(void* data, ndarray arr) + + void* PyArray_GETPTR1(ndarray m, npy_intp i) nogil + void* PyArray_GETPTR2(ndarray m, npy_intp i, npy_intp j) nogil + void* PyArray_GETPTR3(ndarray m, npy_intp i, npy_intp j, npy_intp k) nogil + void* PyArray_GETPTR4(ndarray m, npy_intp i, npy_intp j, npy_intp k, npy_intp l) nogil + + # Cannot be supported due to out arg + # void PyArray_DESCR_REPLACE(descr) + + + object PyArray_Copy(ndarray) + object PyArray_FromObject(object op, int type, int min_depth, int max_depth) + object PyArray_ContiguousFromObject(object op, int type, int min_depth, int max_depth) + object PyArray_CopyFromObject(object op, int type, int min_depth, int max_depth) + + object PyArray_Cast(ndarray mp, int type_num) + object PyArray_Take(ndarray ap, object items, int axis) + object PyArray_Put(ndarray ap, object items, object values) + + void PyArray_ITER_RESET(flatiter it) nogil + void PyArray_ITER_NEXT(flatiter it) nogil + void PyArray_ITER_GOTO(flatiter it, npy_intp* destination) nogil + void PyArray_ITER_GOTO1D(flatiter it, npy_intp ind) nogil + void* PyArray_ITER_DATA(flatiter it) nogil + bint PyArray_ITER_NOTDONE(flatiter it) nogil + + void PyArray_MultiIter_RESET(broadcast multi) nogil + void PyArray_MultiIter_NEXT(broadcast multi) nogil + void PyArray_MultiIter_GOTO(broadcast multi, npy_intp dest) nogil + void PyArray_MultiIter_GOTO1D(broadcast multi, npy_intp ind) nogil + void* PyArray_MultiIter_DATA(broadcast multi, npy_intp i) nogil + void PyArray_MultiIter_NEXTi(broadcast multi, npy_intp i) nogil + bint PyArray_MultiIter_NOTDONE(broadcast multi) nogil + + # Functions from __multiarray_api.h + + # Functions taking dtype and returning object/ndarray are disabled + # for now as they steal dtype references. I'm conservative and disable + # more than is probably needed until it can be checked further. + int PyArray_SetNumericOps (object) except -1 + object PyArray_GetNumericOps () + int PyArray_INCREF (ndarray) except * # uses PyArray_Item_INCREF... + int PyArray_XDECREF (ndarray) except * # uses PyArray_Item_DECREF... + void PyArray_SetStringFunction (object, int) + dtype PyArray_DescrFromType (int) + object PyArray_TypeObjectFromType (int) + char * PyArray_Zero (ndarray) + char * PyArray_One (ndarray) + #object PyArray_CastToType (ndarray, dtype, int) + int PyArray_CastTo (ndarray, ndarray) except -1 + int PyArray_CastAnyTo (ndarray, ndarray) except -1 + int PyArray_CanCastSafely (int, int) # writes errors + npy_bool PyArray_CanCastTo (dtype, dtype) # writes errors + int PyArray_ObjectType (object, int) except 0 + dtype PyArray_DescrFromObject (object, dtype) + #ndarray* PyArray_ConvertToCommonType (object, int *) + dtype PyArray_DescrFromScalar (object) + dtype PyArray_DescrFromTypeObject (object) + npy_intp PyArray_Size (object) + #object PyArray_Scalar (void *, dtype, object) + #object PyArray_FromScalar (object, dtype) + void PyArray_ScalarAsCtype (object, void *) + #int PyArray_CastScalarToCtype (object, void *, dtype) + #int PyArray_CastScalarDirect (object, dtype, void *, int) + object PyArray_ScalarFromObject (object) + #PyArray_VectorUnaryFunc * PyArray_GetCastFunc (dtype, int) + object PyArray_FromDims (int, int *, int) + #object PyArray_FromDimsAndDataAndDescr (int, int *, dtype, char *) + #object PyArray_FromAny (object, dtype, int, int, int, object) + object PyArray_EnsureArray (object) + object PyArray_EnsureAnyArray (object) + #object PyArray_FromFile (stdio.FILE *, dtype, npy_intp, char *) + #object PyArray_FromString (char *, npy_intp, dtype, npy_intp, char *) + #object PyArray_FromBuffer (object, dtype, npy_intp, npy_intp) + #object PyArray_FromIter (object, dtype, npy_intp) + object PyArray_Return (ndarray) + #object PyArray_GetField (ndarray, dtype, int) + #int PyArray_SetField (ndarray, dtype, int, object) except -1 + object PyArray_Byteswap (ndarray, npy_bool) + object PyArray_Resize (ndarray, PyArray_Dims *, int, NPY_ORDER) + int PyArray_MoveInto (ndarray, ndarray) except -1 + int PyArray_CopyInto (ndarray, ndarray) except -1 + int PyArray_CopyAnyInto (ndarray, ndarray) except -1 + int PyArray_CopyObject (ndarray, object) except -1 + object PyArray_NewCopy (ndarray, NPY_ORDER) + object PyArray_ToList (ndarray) + object PyArray_ToString (ndarray, NPY_ORDER) + int PyArray_ToFile (ndarray, stdio.FILE *, char *, char *) except -1 + int PyArray_Dump (object, object, int) except -1 + object PyArray_Dumps (object, int) + int PyArray_ValidType (int) # Cannot error + void PyArray_UpdateFlags (ndarray, int) + object PyArray_New (type, int, npy_intp *, int, npy_intp *, void *, int, int, object) + #object PyArray_NewFromDescr (type, dtype, int, npy_intp *, npy_intp *, void *, int, object) + #dtype PyArray_DescrNew (dtype) + dtype PyArray_DescrNewFromType (int) + double PyArray_GetPriority (object, double) # clears errors as of 1.25 + object PyArray_IterNew (object) + object PyArray_MultiIterNew (int, ...) + + int PyArray_PyIntAsInt (object) except? -1 + npy_intp PyArray_PyIntAsIntp (object) + int PyArray_Broadcast (broadcast) except -1 + void PyArray_FillObjectArray (ndarray, object) except * + int PyArray_FillWithScalar (ndarray, object) except -1 + npy_bool PyArray_CheckStrides (int, int, npy_intp, npy_intp, npy_intp *, npy_intp *) + dtype PyArray_DescrNewByteorder (dtype, char) + object PyArray_IterAllButAxis (object, int *) + #object PyArray_CheckFromAny (object, dtype, int, int, int, object) + #object PyArray_FromArray (ndarray, dtype, int) + object PyArray_FromInterface (object) + object PyArray_FromStructInterface (object) + #object PyArray_FromArrayAttr (object, dtype, object) + #NPY_SCALARKIND PyArray_ScalarKind (int, ndarray*) + int PyArray_CanCoerceScalar (int, int, NPY_SCALARKIND) + object PyArray_NewFlagsObject (object) + npy_bool PyArray_CanCastScalar (type, type) + #int PyArray_CompareUCS4 (npy_ucs4 *, npy_ucs4 *, register size_t) + int PyArray_RemoveSmallest (broadcast) except -1 + int PyArray_ElementStrides (object) + void PyArray_Item_INCREF (char *, dtype) except * + void PyArray_Item_XDECREF (char *, dtype) except * + object PyArray_FieldNames (object) + object PyArray_Transpose (ndarray, PyArray_Dims *) + object PyArray_TakeFrom (ndarray, object, int, ndarray, NPY_CLIPMODE) + object PyArray_PutTo (ndarray, object, object, NPY_CLIPMODE) + object PyArray_PutMask (ndarray, object, object) + object PyArray_Repeat (ndarray, object, int) + object PyArray_Choose (ndarray, object, ndarray, NPY_CLIPMODE) + int PyArray_Sort (ndarray, int, NPY_SORTKIND) except -1 + object PyArray_ArgSort (ndarray, int, NPY_SORTKIND) + object PyArray_SearchSorted (ndarray, object, NPY_SEARCHSIDE, PyObject *) + object PyArray_ArgMax (ndarray, int, ndarray) + object PyArray_ArgMin (ndarray, int, ndarray) + object PyArray_Reshape (ndarray, object) + object PyArray_Newshape (ndarray, PyArray_Dims *, NPY_ORDER) + object PyArray_Squeeze (ndarray) + #object PyArray_View (ndarray, dtype, type) + object PyArray_SwapAxes (ndarray, int, int) + object PyArray_Max (ndarray, int, ndarray) + object PyArray_Min (ndarray, int, ndarray) + object PyArray_Ptp (ndarray, int, ndarray) + object PyArray_Mean (ndarray, int, int, ndarray) + object PyArray_Trace (ndarray, int, int, int, int, ndarray) + object PyArray_Diagonal (ndarray, int, int, int) + object PyArray_Clip (ndarray, object, object, ndarray) + object PyArray_Conjugate (ndarray, ndarray) + object PyArray_Nonzero (ndarray) + object PyArray_Std (ndarray, int, int, ndarray, int) + object PyArray_Sum (ndarray, int, int, ndarray) + object PyArray_CumSum (ndarray, int, int, ndarray) + object PyArray_Prod (ndarray, int, int, ndarray) + object PyArray_CumProd (ndarray, int, int, ndarray) + object PyArray_All (ndarray, int, ndarray) + object PyArray_Any (ndarray, int, ndarray) + object PyArray_Compress (ndarray, object, int, ndarray) + object PyArray_Flatten (ndarray, NPY_ORDER) + object PyArray_Ravel (ndarray, NPY_ORDER) + npy_intp PyArray_MultiplyList (npy_intp *, int) + int PyArray_MultiplyIntList (int *, int) + void * PyArray_GetPtr (ndarray, npy_intp*) + int PyArray_CompareLists (npy_intp *, npy_intp *, int) + #int PyArray_AsCArray (object*, void *, npy_intp *, int, dtype) + #int PyArray_As1D (object*, char **, int *, int) + #int PyArray_As2D (object*, char ***, int *, int *, int) + int PyArray_Free (object, void *) + #int PyArray_Converter (object, object*) + int PyArray_IntpFromSequence (object, npy_intp *, int) except -1 + object PyArray_Concatenate (object, int) + object PyArray_InnerProduct (object, object) + object PyArray_MatrixProduct (object, object) + object PyArray_CopyAndTranspose (object) + object PyArray_Correlate (object, object, int) + int PyArray_TypestrConvert (int, int) + #int PyArray_DescrConverter (object, dtype*) except 0 + #int PyArray_DescrConverter2 (object, dtype*) except 0 + int PyArray_IntpConverter (object, PyArray_Dims *) except 0 + #int PyArray_BufferConverter (object, chunk) except 0 + int PyArray_AxisConverter (object, int *) except 0 + int PyArray_BoolConverter (object, npy_bool *) except 0 + int PyArray_ByteorderConverter (object, char *) except 0 + int PyArray_OrderConverter (object, NPY_ORDER *) except 0 + unsigned char PyArray_EquivTypes (dtype, dtype) # clears errors + #object PyArray_Zeros (int, npy_intp *, dtype, int) + #object PyArray_Empty (int, npy_intp *, dtype, int) + object PyArray_Where (object, object, object) + object PyArray_Arange (double, double, double, int) + #object PyArray_ArangeObj (object, object, object, dtype) + int PyArray_SortkindConverter (object, NPY_SORTKIND *) except 0 + object PyArray_LexSort (object, int) + object PyArray_Round (ndarray, int, ndarray) + unsigned char PyArray_EquivTypenums (int, int) + int PyArray_RegisterDataType (dtype) except -1 + int PyArray_RegisterCastFunc (dtype, int, PyArray_VectorUnaryFunc *) except -1 + int PyArray_RegisterCanCast (dtype, int, NPY_SCALARKIND) except -1 + #void PyArray_InitArrFuncs (PyArray_ArrFuncs *) + object PyArray_IntTupleFromIntp (int, npy_intp *) + int PyArray_TypeNumFromName (char *) + int PyArray_ClipmodeConverter (object, NPY_CLIPMODE *) except 0 + #int PyArray_OutputConverter (object, ndarray*) except 0 + object PyArray_BroadcastToShape (object, npy_intp *, int) + void _PyArray_SigintHandler (int) + void* _PyArray_GetSigintBuf () + #int PyArray_DescrAlignConverter (object, dtype*) except 0 + #int PyArray_DescrAlignConverter2 (object, dtype*) except 0 + int PyArray_SearchsideConverter (object, void *) except 0 + object PyArray_CheckAxis (ndarray, int *, int) + npy_intp PyArray_OverflowMultiplyList (npy_intp *, int) + int PyArray_CompareString (char *, char *, size_t) + int PyArray_SetBaseObject(ndarray, base) except -1 # NOTE: steals a reference to base! Use "set_array_base()" instead. + + +# Typedefs that matches the runtime dtype objects in +# the numpy module. + +# The ones that are commented out needs an IFDEF function +# in Cython to enable them only on the right systems. + +ctypedef npy_int8 int8_t +ctypedef npy_int16 int16_t +ctypedef npy_int32 int32_t +ctypedef npy_int64 int64_t +#ctypedef npy_int96 int96_t +#ctypedef npy_int128 int128_t + +ctypedef npy_uint8 uint8_t +ctypedef npy_uint16 uint16_t +ctypedef npy_uint32 uint32_t +ctypedef npy_uint64 uint64_t +#ctypedef npy_uint96 uint96_t +#ctypedef npy_uint128 uint128_t + +ctypedef npy_float32 float32_t +ctypedef npy_float64 float64_t +#ctypedef npy_float80 float80_t +#ctypedef npy_float128 float128_t + +ctypedef float complex complex64_t +ctypedef double complex complex128_t + +# The int types are mapped a bit surprising -- +# numpy.int corresponds to 'l' and numpy.long to 'q' +ctypedef npy_long int_t +ctypedef npy_longlong longlong_t + +ctypedef npy_ulong uint_t +ctypedef npy_ulonglong ulonglong_t + +ctypedef npy_intp intp_t +ctypedef npy_uintp uintp_t + +ctypedef npy_double float_t +ctypedef npy_double double_t +ctypedef npy_longdouble longdouble_t + +ctypedef npy_cfloat cfloat_t +ctypedef npy_cdouble cdouble_t +ctypedef npy_clongdouble clongdouble_t + +ctypedef npy_cdouble complex_t + +cdef inline object PyArray_MultiIterNew1(a): + return PyArray_MultiIterNew(1, a) + +cdef inline object PyArray_MultiIterNew2(a, b): + return PyArray_MultiIterNew(2, a, b) + +cdef inline object PyArray_MultiIterNew3(a, b, c): + return PyArray_MultiIterNew(3, a, b, c) + +cdef inline object PyArray_MultiIterNew4(a, b, c, d): + return PyArray_MultiIterNew(4, a, b, c, d) + +cdef inline object PyArray_MultiIterNew5(a, b, c, d, e): + return PyArray_MultiIterNew(5, a, b, c, d, e) + +cdef inline tuple PyDataType_SHAPE(dtype d): + if PyDataType_HASSUBARRAY(d): + return d.subarray.shape + else: + return () + + +cdef extern from "numpy/ndarrayobject.h": + PyTypeObject PyTimedeltaArrType_Type + PyTypeObject PyDatetimeArrType_Type + ctypedef int64_t npy_timedelta + ctypedef int64_t npy_datetime + +cdef extern from "numpy/ndarraytypes.h": + ctypedef struct PyArray_DatetimeMetaData: + NPY_DATETIMEUNIT base + int64_t num + +cdef extern from "numpy/arrayscalars.h": + + # abstract types + ctypedef class numpy.generic [object PyObject]: + pass + ctypedef class numpy.number [object PyObject]: + pass + ctypedef class numpy.integer [object PyObject]: + pass + ctypedef class numpy.signedinteger [object PyObject]: + pass + ctypedef class numpy.unsignedinteger [object PyObject]: + pass + ctypedef class numpy.inexact [object PyObject]: + pass + ctypedef class numpy.floating [object PyObject]: + pass + ctypedef class numpy.complexfloating [object PyObject]: + pass + ctypedef class numpy.flexible [object PyObject]: + pass + ctypedef class numpy.character [object PyObject]: + pass + + ctypedef struct PyDatetimeScalarObject: + # PyObject_HEAD + npy_datetime obval + PyArray_DatetimeMetaData obmeta + + ctypedef struct PyTimedeltaScalarObject: + # PyObject_HEAD + npy_timedelta obval + PyArray_DatetimeMetaData obmeta + + ctypedef enum NPY_DATETIMEUNIT: + NPY_FR_Y + NPY_FR_M + NPY_FR_W + NPY_FR_D + NPY_FR_B + NPY_FR_h + NPY_FR_m + NPY_FR_s + NPY_FR_ms + NPY_FR_us + NPY_FR_ns + NPY_FR_ps + NPY_FR_fs + NPY_FR_as + NPY_FR_GENERIC + + +# +# ufunc API +# + +cdef extern from "numpy/ufuncobject.h": + + ctypedef void (*PyUFuncGenericFunction) (char **, npy_intp *, npy_intp *, void *) + + ctypedef class numpy.ufunc [object PyUFuncObject, check_size ignore]: + cdef: + int nin, nout, nargs + int identity + PyUFuncGenericFunction *functions + void **data + int ntypes + int check_return + char *name + char *types + char *doc + void *ptr + PyObject *obj + PyObject *userloops + + cdef enum: + PyUFunc_Zero + PyUFunc_One + PyUFunc_None + UFUNC_ERR_IGNORE + UFUNC_ERR_WARN + UFUNC_ERR_RAISE + UFUNC_ERR_CALL + UFUNC_ERR_PRINT + UFUNC_ERR_LOG + UFUNC_MASK_DIVIDEBYZERO + UFUNC_MASK_OVERFLOW + UFUNC_MASK_UNDERFLOW + UFUNC_MASK_INVALID + UFUNC_SHIFT_DIVIDEBYZERO + UFUNC_SHIFT_OVERFLOW + UFUNC_SHIFT_UNDERFLOW + UFUNC_SHIFT_INVALID + UFUNC_FPE_DIVIDEBYZERO + UFUNC_FPE_OVERFLOW + UFUNC_FPE_UNDERFLOW + UFUNC_FPE_INVALID + UFUNC_ERR_DEFAULT + UFUNC_ERR_DEFAULT2 + + object PyUFunc_FromFuncAndData(PyUFuncGenericFunction *, + void **, char *, int, int, int, int, char *, char *, int) + int PyUFunc_RegisterLoopForType(ufunc, int, + PyUFuncGenericFunction, int *, void *) except -1 + void PyUFunc_f_f_As_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_d_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_f_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_g_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F_As_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_F_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_D_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_G_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f_As_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_ff_f \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_dd_d \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_gg_g \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F_As_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_DD_D \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_FF_F \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_GG_G \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_O_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_OO_O_method \ + (char **, npy_intp *, npy_intp *, void *) + void PyUFunc_On_Om \ + (char **, npy_intp *, npy_intp *, void *) + int PyUFunc_GetPyValues \ + (char *, int *, int *, PyObject **) + int PyUFunc_checkfperr \ + (int, PyObject *, int *) + void PyUFunc_clearfperr() + int PyUFunc_getfperr() + int PyUFunc_handlefperr \ + (int, PyObject *, int, int *) except -1 + int PyUFunc_ReplaceLoopBySignature \ + (ufunc, PyUFuncGenericFunction, int *, PyUFuncGenericFunction *) + object PyUFunc_FromFuncAndDataAndSignature \ + (PyUFuncGenericFunction *, void **, char *, int, int, int, + int, char *, char *, int, char *) + + int _import_umath() except -1 + +cdef inline void set_array_base(ndarray arr, object base): + Py_INCREF(base) # important to do this before stealing the reference below! + PyArray_SetBaseObject(arr, base) + +cdef inline object get_array_base(ndarray arr): + base = PyArray_BASE(arr) + if base is NULL: + return None + return base + +# Versions of the import_* functions which are more suitable for +# Cython code. +cdef inline int import_array() except -1: + try: + __pyx_import_array() + except Exception: + raise ImportError("numpy.core.multiarray failed to import") + +cdef inline int import_umath() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + +cdef inline int import_ufunc() except -1: + try: + _import_umath() + except Exception: + raise ImportError("numpy.core.umath failed to import") + +cdef extern from *: + # Leave a marker that the NumPy declarations came from this file + # See https://github.com/cython/cython/issues/3573 + """ + /* NumPy API declarations from "numpy/__init__.pxd" */ + """ + + +cdef inline bint is_timedelta64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.timedelta64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyTimedeltaArrType_Type) + + +cdef inline bint is_datetime64_object(object obj): + """ + Cython equivalent of `isinstance(obj, np.datetime64)` + + Parameters + ---------- + obj : object + + Returns + ------- + bool + """ + return PyObject_TypeCheck(obj, &PyDatetimeArrType_Type) + + +cdef inline npy_datetime get_datetime64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy datetime64 object + + Note that to interpret this as a datetime, the corresponding unit is + also needed. That can be found using `get_datetime64_unit`. + """ + return (obj).obval + + +cdef inline npy_timedelta get_timedelta64_value(object obj) nogil: + """ + returns the int64 value underlying scalar numpy timedelta64 object + """ + return (obj).obval + + +cdef inline NPY_DATETIMEUNIT get_datetime64_unit(object obj) nogil: + """ + returns the unit part of the dtype for a numpy datetime64 object. + """ + return (obj).obmeta.base diff --git a/lib/python3.12/site-packages/numpy/__init__.py b/lib/python3.12/site-packages/numpy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..91da496a95271f8e3eb4f4ea2cbaec925325a1f5 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/__init__.py @@ -0,0 +1,461 @@ +""" +NumPy +===== + +Provides + 1. An array object of arbitrary homogeneous items + 2. Fast mathematical operations over arrays + 3. Linear Algebra, Fourier Transforms, Random Number Generation + +How to use the documentation +---------------------------- +Documentation is available in two forms: docstrings provided +with the code, and a loose standing reference guide, available from +`the NumPy homepage `_. + +We recommend exploring the docstrings using +`IPython `_, an advanced Python shell with +TAB-completion and introspection capabilities. See below for further +instructions. + +The docstring examples assume that `numpy` has been imported as ``np``:: + + >>> import numpy as np + +Code snippets are indicated by three greater-than signs:: + + >>> x = 42 + >>> x = x + 1 + +Use the built-in ``help`` function to view a function's docstring:: + + >>> help(np.sort) + ... # doctest: +SKIP + +For some objects, ``np.info(obj)`` may provide additional help. This is +particularly true if you see the line "Help on ufunc object:" at the top +of the help() page. Ufuncs are implemented in C, not Python, for speed. +The native Python help() does not know how to view their help, but our +np.info() function does. + +To search for documents containing a keyword, do:: + + >>> np.lookfor('keyword') + ... # doctest: +SKIP + +General-purpose documents like a glossary and help on the basic concepts +of numpy are available under the ``doc`` sub-module:: + + >>> from numpy import doc + >>> help(doc) + ... # doctest: +SKIP + +Available subpackages +--------------------- +lib + Basic functions used by several sub-packages. +random + Core Random Tools +linalg + Core Linear Algebra Tools +fft + Core FFT routines +polynomial + Polynomial tools +testing + NumPy testing tools +distutils + Enhancements to distutils with support for + Fortran compilers support and more (for Python <= 3.11). + +Utilities +--------- +test + Run numpy unittests +show_config + Show numpy build configuration +matlib + Make everything matrices. +__version__ + NumPy version string + +Viewing documentation using IPython +----------------------------------- + +Start IPython and import `numpy` usually under the alias ``np``: `import +numpy as np`. Then, directly past or use the ``%cpaste`` magic to paste +examples into the shell. To see which functions are available in `numpy`, +type ``np.`` (where ```` refers to the TAB key), or use +``np.*cos*?`` (where ```` refers to the ENTER key) to narrow +down the list. To view the docstring for a function, use +``np.cos?`` (to view the docstring) and ``np.cos??`` (to view +the source code). + +Copies vs. in-place operation +----------------------------- +Most of the functions in `numpy` return a copy of the array argument +(e.g., `np.sort`). In-place versions of these functions are often +available as array methods, i.e. ``x = np.array([1,2,3]); x.sort()``. +Exceptions to this rule are documented. + +""" +import sys +import warnings + +from ._globals import _NoValue, _CopyMode +# These exceptions were moved in 1.25 and are hidden from __dir__() +from .exceptions import ( + ComplexWarning, ModuleDeprecationWarning, VisibleDeprecationWarning, + TooHardError, AxisError) + + +# If a version with git hash was stored, use that instead +from . import version +from .version import __version__ + +# We first need to detect if we're being called as part of the numpy setup +# procedure itself in a reliable manner. +try: + __NUMPY_SETUP__ +except NameError: + __NUMPY_SETUP__ = False + +if __NUMPY_SETUP__: + sys.stderr.write('Running from numpy source directory.\n') +else: + # Allow distributors to run custom init code before importing numpy.core + from . import _distributor_init + + try: + from numpy.__config__ import show as show_config + except ImportError as e: + msg = """Error importing numpy: you should not try to import numpy from + its source directory; please exit the numpy source tree, and relaunch + your python interpreter from there.""" + raise ImportError(msg) from e + + __all__ = [ + 'exceptions', 'ModuleDeprecationWarning', 'VisibleDeprecationWarning', + 'ComplexWarning', 'TooHardError', 'AxisError'] + + # mapping of {name: (value, deprecation_msg)} + __deprecated_attrs__ = {} + + from . import core + from .core import * + from . import compat + from . import exceptions + from . import dtypes + from . import lib + # NOTE: to be revisited following future namespace cleanup. + # See gh-14454 and gh-15672 for discussion. + from .lib import * + + from . import linalg + from . import fft + from . import polynomial + from . import random + from . import ctypeslib + from . import ma + from . import matrixlib as _mat + from .matrixlib import * + + # Deprecations introduced in NumPy 1.20.0, 2020-06-06 + import builtins as _builtins + + _msg = ( + "module 'numpy' has no attribute '{n}'.\n" + "`np.{n}` was a deprecated alias for the builtin `{n}`. " + "To avoid this error in existing code, use `{n}` by itself. " + "Doing this will not modify any behavior and is safe. {extended_msg}\n" + "The aliases was originally deprecated in NumPy 1.20; for more " + "details and guidance see the original release note at:\n" + " https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations") + + _specific_msg = ( + "If you specifically wanted the numpy scalar type, use `np.{}` here.") + + _int_extended_msg = ( + "When replacing `np.{}`, you may wish to use e.g. `np.int64` " + "or `np.int32` to specify the precision. If you wish to review " + "your current use, check the release note link for " + "additional information.") + + _type_info = [ + ("object", ""), # The NumPy scalar only exists by name. + ("bool", _specific_msg.format("bool_")), + ("float", _specific_msg.format("float64")), + ("complex", _specific_msg.format("complex128")), + ("str", _specific_msg.format("str_")), + ("int", _int_extended_msg.format("int"))] + + __former_attrs__ = { + n: _msg.format(n=n, extended_msg=extended_msg) + for n, extended_msg in _type_info + } + + # Future warning introduced in NumPy 1.24.0, 2022-11-17 + _msg = ( + "`np.{n}` is a deprecated alias for `{an}`. (Deprecated NumPy 1.24)") + + # Some of these are awkward (since `np.str` may be preferable in the long + # term), but overall the names ending in 0 seem undesirable + _type_info = [ + ("bool8", bool_, "np.bool_"), + ("int0", intp, "np.intp"), + ("uint0", uintp, "np.uintp"), + ("str0", str_, "np.str_"), + ("bytes0", bytes_, "np.bytes_"), + ("void0", void, "np.void"), + ("object0", object_, + "`np.object0` is a deprecated alias for `np.object_`. " + "`object` can be used instead. (Deprecated NumPy 1.24)")] + + # Some of these could be defined right away, but most were aliases to + # the Python objects and only removed in NumPy 1.24. Defining them should + # probably wait for NumPy 1.26 or 2.0. + # When defined, these should possibly not be added to `__all__` to avoid + # import with `from numpy import *`. + __future_scalars__ = {"bool", "long", "ulong", "str", "bytes", "object"} + + __deprecated_attrs__.update({ + n: (alias, _msg.format(n=n, an=an)) for n, alias, an in _type_info}) + + import math + + __deprecated_attrs__['math'] = (math, + "`np.math` is a deprecated alias for the standard library `math` " + "module (Deprecated Numpy 1.25). Replace usages of `np.math` with " + "`math`") + + del math, _msg, _type_info + + from .core import abs + # now that numpy modules are imported, can initialize limits + core.getlimits._register_known_types() + + __all__.extend(['__version__', 'show_config']) + __all__.extend(core.__all__) + __all__.extend(_mat.__all__) + __all__.extend(lib.__all__) + __all__.extend(['linalg', 'fft', 'random', 'ctypeslib', 'ma']) + + # Remove min and max from __all__ to avoid `from numpy import *` override + # the builtins min/max. Temporary fix for 1.25.x/1.26.x, see gh-24229. + __all__.remove('min') + __all__.remove('max') + __all__.remove('round') + + # Remove one of the two occurrences of `issubdtype`, which is exposed as + # both `numpy.core.issubdtype` and `numpy.lib.issubdtype`. + __all__.remove('issubdtype') + + # These are exported by np.core, but are replaced by the builtins below + # remove them to ensure that we don't end up with `np.long == np.int_`, + # which would be a breaking change. + del long, unicode + __all__.remove('long') + __all__.remove('unicode') + + # Remove things that are in the numpy.lib but not in the numpy namespace + # Note that there is a test (numpy/tests/test_public_api.py:test_numpy_namespace) + # that prevents adding more things to the main namespace by accident. + # The list below will grow until the `from .lib import *` fixme above is + # taken care of + __all__.remove('Arrayterator') + del Arrayterator + + # These names were removed in NumPy 1.20. For at least one release, + # attempts to access these names in the numpy namespace will trigger + # a warning, and calling the function will raise an exception. + _financial_names = ['fv', 'ipmt', 'irr', 'mirr', 'nper', 'npv', 'pmt', + 'ppmt', 'pv', 'rate'] + __expired_functions__ = { + name: (f'In accordance with NEP 32, the function {name} was removed ' + 'from NumPy version 1.20. A replacement for this function ' + 'is available in the numpy_financial library: ' + 'https://pypi.org/project/numpy-financial') + for name in _financial_names} + + # Filter out Cython harmless warnings + warnings.filterwarnings("ignore", message="numpy.dtype size changed") + warnings.filterwarnings("ignore", message="numpy.ufunc size changed") + warnings.filterwarnings("ignore", message="numpy.ndarray size changed") + + # oldnumeric and numarray were removed in 1.9. In case some packages import + # but do not use them, we define them here for backward compatibility. + oldnumeric = 'removed' + numarray = 'removed' + + def __getattr__(attr): + # Warn for expired attributes, and return a dummy function + # that always raises an exception. + import warnings + import math + try: + msg = __expired_functions__[attr] + except KeyError: + pass + else: + warnings.warn(msg, DeprecationWarning, stacklevel=2) + + def _expired(*args, **kwds): + raise RuntimeError(msg) + + return _expired + + # Emit warnings for deprecated attributes + try: + val, msg = __deprecated_attrs__[attr] + except KeyError: + pass + else: + warnings.warn(msg, DeprecationWarning, stacklevel=2) + return val + + if attr in __future_scalars__: + # And future warnings for those that will change, but also give + # the AttributeError + warnings.warn( + f"In the future `np.{attr}` will be defined as the " + "corresponding NumPy scalar.", FutureWarning, stacklevel=2) + + if attr in __former_attrs__: + raise AttributeError(__former_attrs__[attr]) + + if attr == 'testing': + import numpy.testing as testing + return testing + elif attr == 'Tester': + "Removed in NumPy 1.25.0" + raise RuntimeError("Tester was removed in NumPy 1.25.") + + raise AttributeError("module {!r} has no attribute " + "{!r}".format(__name__, attr)) + + def __dir__(): + public_symbols = globals().keys() | {'testing'} + public_symbols -= { + "core", "matrixlib", + # These were moved in 1.25 and may be deprecated eventually: + "ModuleDeprecationWarning", "VisibleDeprecationWarning", + "ComplexWarning", "TooHardError", "AxisError" + } + return list(public_symbols) + + # Pytest testing + from numpy._pytesttester import PytestTester + test = PytestTester(__name__) + del PytestTester + + def _sanity_check(): + """ + Quick sanity checks for common bugs caused by environment. + There are some cases e.g. with wrong BLAS ABI that cause wrong + results under specific runtime conditions that are not necessarily + achieved during test suite runs, and it is useful to catch those early. + + See https://github.com/numpy/numpy/issues/8577 and other + similar bug reports. + + """ + try: + x = ones(2, dtype=float32) + if not abs(x.dot(x) - float32(2.0)) < 1e-5: + raise AssertionError() + except AssertionError: + msg = ("The current Numpy installation ({!r}) fails to " + "pass simple sanity checks. This can be caused for example " + "by incorrect BLAS library being linked in, or by mixing " + "package managers (pip, conda, apt, ...). Search closed " + "numpy issues for similar problems.") + raise RuntimeError(msg.format(__file__)) from None + + _sanity_check() + del _sanity_check + + def _mac_os_check(): + """ + Quick Sanity check for Mac OS look for accelerate build bugs. + Testing numpy polyfit calls init_dgelsd(LAPACK) + """ + try: + c = array([3., 2., 1.]) + x = linspace(0, 2, 5) + y = polyval(c, x) + _ = polyfit(x, y, 2, cov=True) + except ValueError: + pass + + if sys.platform == "darwin": + from . import exceptions + with warnings.catch_warnings(record=True) as w: + _mac_os_check() + # Throw runtime error, if the test failed Check for warning and error_message + if len(w) > 0: + for _wn in w: + if _wn.category is exceptions.RankWarning: + # Ignore other warnings, they may not be relevant (see gh-25433). + error_message = f"{_wn.category.__name__}: {str(_wn.message)}" + msg = ( + "Polyfit sanity test emitted a warning, most likely due " + "to using a buggy Accelerate backend." + "\nIf you compiled yourself, more information is available at:" + "\nhttps://numpy.org/devdocs/building/index.html" + "\nOtherwise report this to the vendor " + "that provided NumPy.\n\n{}\n".format(error_message)) + raise RuntimeError(msg) + del _wn + del w + del _mac_os_check + + # We usually use madvise hugepages support, but on some old kernels it + # is slow and thus better avoided. + # Specifically kernel version 4.6 had a bug fix which probably fixed this: + # https://github.com/torvalds/linux/commit/7cf91a98e607c2f935dbcc177d70011e95b8faff + import os + use_hugepage = os.environ.get("NUMPY_MADVISE_HUGEPAGE", None) + if sys.platform == "linux" and use_hugepage is None: + # If there is an issue with parsing the kernel version, + # set use_hugepages to 0. Usage of LooseVersion will handle + # the kernel version parsing better, but avoided since it + # will increase the import time. See: #16679 for related discussion. + try: + use_hugepage = 1 + kernel_version = os.uname().release.split(".")[:2] + kernel_version = tuple(int(v) for v in kernel_version) + if kernel_version < (4, 6): + use_hugepage = 0 + except ValueError: + use_hugepages = 0 + elif use_hugepage is None: + # This is not Linux, so it should not matter, just enable anyway + use_hugepage = 1 + else: + use_hugepage = int(use_hugepage) + + # Note that this will currently only make a difference on Linux + core.multiarray._set_madvise_hugepage(use_hugepage) + del use_hugepage + + # Give a warning if NumPy is reloaded or imported on a sub-interpreter + # We do this from python, since the C-module may not be reloaded and + # it is tidier organized. + core.multiarray._multiarray_umath._reload_guard() + + # default to "weak" promotion for "NumPy 2". + core._set_promotion_state( + os.environ.get("NPY_PROMOTION_STATE", + "weak" if _using_numpy2_behavior() else "legacy")) + + # Tell PyInstaller where to find hook-numpy.py + def _pyinstaller_hooks_dir(): + from pathlib import Path + return [str(Path(__file__).with_name("_pyinstaller").resolve())] + + # Remove symbols imported for internal use + del os + + +# Remove symbols imported for internal use +del sys, warnings diff --git a/lib/python3.12/site-packages/numpy/__init__.pyi b/lib/python3.12/site-packages/numpy/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a185bfe754e3d6ba5004f8cca06177e60a7aa13c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/__init__.pyi @@ -0,0 +1,4422 @@ +import builtins +import sys +import os +import mmap +import ctypes as ct +import array as _array +import datetime as dt +import enum +from abc import abstractmethod +from types import TracebackType, MappingProxyType, GenericAlias +from contextlib import ContextDecorator +from contextlib import contextmanager + +from numpy._pytesttester import PytestTester +from numpy.core._internal import _ctypes + +from numpy._typing import ( + # Arrays + ArrayLike, + NDArray, + _SupportsArray, + _NestedSequence, + _FiniteNestedSequence, + _SupportsArray, + _ArrayLikeBool_co, + _ArrayLikeUInt_co, + _ArrayLikeInt_co, + _ArrayLikeFloat_co, + _ArrayLikeComplex_co, + _ArrayLikeNumber_co, + _ArrayLikeTD64_co, + _ArrayLikeDT64_co, + _ArrayLikeObject_co, + _ArrayLikeStr_co, + _ArrayLikeBytes_co, + _ArrayLikeUnknown, + _UnknownType, + + # DTypes + DTypeLike, + _DTypeLike, + _DTypeLikeVoid, + _SupportsDType, + _VoidDTypeLike, + + # Shapes + _Shape, + _ShapeLike, + + # Scalars + _CharLike_co, + _BoolLike_co, + _IntLike_co, + _FloatLike_co, + _ComplexLike_co, + _TD64Like_co, + _NumberLike_co, + _ScalarLike_co, + + # `number` precision + NBitBase, + _256Bit, + _128Bit, + _96Bit, + _80Bit, + _64Bit, + _32Bit, + _16Bit, + _8Bit, + _NBitByte, + _NBitShort, + _NBitIntC, + _NBitIntP, + _NBitInt, + _NBitLongLong, + _NBitHalf, + _NBitSingle, + _NBitDouble, + _NBitLongDouble, + + # Character codes + _BoolCodes, + _UInt8Codes, + _UInt16Codes, + _UInt32Codes, + _UInt64Codes, + _Int8Codes, + _Int16Codes, + _Int32Codes, + _Int64Codes, + _Float16Codes, + _Float32Codes, + _Float64Codes, + _Complex64Codes, + _Complex128Codes, + _ByteCodes, + _ShortCodes, + _IntCCodes, + _IntPCodes, + _IntCodes, + _LongLongCodes, + _UByteCodes, + _UShortCodes, + _UIntCCodes, + _UIntPCodes, + _UIntCodes, + _ULongLongCodes, + _HalfCodes, + _SingleCodes, + _DoubleCodes, + _LongDoubleCodes, + _CSingleCodes, + _CDoubleCodes, + _CLongDoubleCodes, + _DT64Codes, + _TD64Codes, + _StrCodes, + _BytesCodes, + _VoidCodes, + _ObjectCodes, + + # Ufuncs + _UFunc_Nin1_Nout1, + _UFunc_Nin2_Nout1, + _UFunc_Nin1_Nout2, + _UFunc_Nin2_Nout2, + _GUFunc_Nin2_Nout1, +) + +from numpy._typing._callable import ( + _BoolOp, + _BoolBitOp, + _BoolSub, + _BoolTrueDiv, + _BoolMod, + _BoolDivMod, + _TD64Div, + _IntTrueDiv, + _UnsignedIntOp, + _UnsignedIntBitOp, + _UnsignedIntMod, + _UnsignedIntDivMod, + _SignedIntOp, + _SignedIntBitOp, + _SignedIntMod, + _SignedIntDivMod, + _FloatOp, + _FloatMod, + _FloatDivMod, + _ComplexOp, + _NumberOp, + _ComparisonOp, +) + +# NOTE: Numpy's mypy plugin is used for removing the types unavailable +# to the specific platform +from numpy._typing._extended_precision import ( + uint128 as uint128, + uint256 as uint256, + int128 as int128, + int256 as int256, + float80 as float80, + float96 as float96, + float128 as float128, + float256 as float256, + complex160 as complex160, + complex192 as complex192, + complex256 as complex256, + complex512 as complex512, +) + +from collections.abc import ( + Callable, + Container, + Iterable, + Iterator, + Mapping, + Sequence, + Sized, +) +from typing import ( + Literal as L, + Any, + Generator, + Generic, + IO, + NoReturn, + overload, + SupportsComplex, + SupportsFloat, + SupportsInt, + TypeVar, + Union, + Protocol, + SupportsIndex, + Final, + final, + ClassVar, +) + +# Ensures that the stubs are picked up +from numpy import ( + ctypeslib as ctypeslib, + exceptions as exceptions, + fft as fft, + lib as lib, + linalg as linalg, + ma as ma, + polynomial as polynomial, + random as random, + testing as testing, + version as version, + exceptions as exceptions, + dtypes as dtypes, +) + +from numpy.core import defchararray, records +char = defchararray +rec = records + +from numpy.core.function_base import ( + linspace as linspace, + logspace as logspace, + geomspace as geomspace, +) + +from numpy.core.fromnumeric import ( + take as take, + reshape as reshape, + choose as choose, + repeat as repeat, + put as put, + swapaxes as swapaxes, + transpose as transpose, + partition as partition, + argpartition as argpartition, + sort as sort, + argsort as argsort, + argmax as argmax, + argmin as argmin, + searchsorted as searchsorted, + resize as resize, + squeeze as squeeze, + diagonal as diagonal, + trace as trace, + ravel as ravel, + nonzero as nonzero, + shape as shape, + compress as compress, + clip as clip, + sum as sum, + all as all, + any as any, + cumsum as cumsum, + ptp as ptp, + max as max, + min as min, + amax as amax, + amin as amin, + prod as prod, + cumprod as cumprod, + ndim as ndim, + size as size, + around as around, + round as round, + mean as mean, + std as std, + var as var, +) + +from numpy.core._asarray import ( + require as require, +) + +from numpy.core._type_aliases import ( + sctypes as sctypes, + sctypeDict as sctypeDict, +) + +from numpy.core._ufunc_config import ( + seterr as seterr, + geterr as geterr, + setbufsize as setbufsize, + getbufsize as getbufsize, + seterrcall as seterrcall, + geterrcall as geterrcall, + _ErrKind, + _ErrFunc, + _ErrDictOptional, +) + +from numpy.core.arrayprint import ( + set_printoptions as set_printoptions, + get_printoptions as get_printoptions, + array2string as array2string, + format_float_scientific as format_float_scientific, + format_float_positional as format_float_positional, + array_repr as array_repr, + array_str as array_str, + set_string_function as set_string_function, + printoptions as printoptions, +) + +from numpy.core.einsumfunc import ( + einsum as einsum, + einsum_path as einsum_path, +) + +from numpy.core.multiarray import ( + ALLOW_THREADS as ALLOW_THREADS, + BUFSIZE as BUFSIZE, + CLIP as CLIP, + MAXDIMS as MAXDIMS, + MAY_SHARE_BOUNDS as MAY_SHARE_BOUNDS, + MAY_SHARE_EXACT as MAY_SHARE_EXACT, + RAISE as RAISE, + WRAP as WRAP, + tracemalloc_domain as tracemalloc_domain, + array as array, + empty_like as empty_like, + empty as empty, + zeros as zeros, + concatenate as concatenate, + inner as inner, + where as where, + lexsort as lexsort, + can_cast as can_cast, + min_scalar_type as min_scalar_type, + result_type as result_type, + dot as dot, + vdot as vdot, + bincount as bincount, + copyto as copyto, + putmask as putmask, + packbits as packbits, + unpackbits as unpackbits, + shares_memory as shares_memory, + may_share_memory as may_share_memory, + asarray as asarray, + asanyarray as asanyarray, + ascontiguousarray as ascontiguousarray, + asfortranarray as asfortranarray, + arange as arange, + busday_count as busday_count, + busday_offset as busday_offset, + compare_chararrays as compare_chararrays, + datetime_as_string as datetime_as_string, + datetime_data as datetime_data, + frombuffer as frombuffer, + fromfile as fromfile, + fromiter as fromiter, + is_busday as is_busday, + promote_types as promote_types, + seterrobj as seterrobj, + geterrobj as geterrobj, + fromstring as fromstring, + frompyfunc as frompyfunc, + nested_iters as nested_iters, + flagsobj, +) + +from numpy.core.numeric import ( + zeros_like as zeros_like, + ones as ones, + ones_like as ones_like, + full as full, + full_like as full_like, + count_nonzero as count_nonzero, + isfortran as isfortran, + argwhere as argwhere, + flatnonzero as flatnonzero, + correlate as correlate, + convolve as convolve, + outer as outer, + tensordot as tensordot, + roll as roll, + rollaxis as rollaxis, + moveaxis as moveaxis, + cross as cross, + indices as indices, + fromfunction as fromfunction, + isscalar as isscalar, + binary_repr as binary_repr, + base_repr as base_repr, + identity as identity, + allclose as allclose, + isclose as isclose, + array_equal as array_equal, + array_equiv as array_equiv, +) + +from numpy.core.numerictypes import ( + maximum_sctype as maximum_sctype, + issctype as issctype, + obj2sctype as obj2sctype, + issubclass_ as issubclass_, + issubsctype as issubsctype, + issubdtype as issubdtype, + sctype2char as sctype2char, + nbytes as nbytes, + cast as cast, + ScalarType as ScalarType, + typecodes as typecodes, +) + +from numpy.core.shape_base import ( + atleast_1d as atleast_1d, + atleast_2d as atleast_2d, + atleast_3d as atleast_3d, + block as block, + hstack as hstack, + stack as stack, + vstack as vstack, +) + +from numpy.exceptions import ( + ComplexWarning as ComplexWarning, + ModuleDeprecationWarning as ModuleDeprecationWarning, + VisibleDeprecationWarning as VisibleDeprecationWarning, + TooHardError as TooHardError, + DTypePromotionError as DTypePromotionError, + AxisError as AxisError, +) + +from numpy.lib import ( + emath as emath, +) + +from numpy.lib.arraypad import ( + pad as pad, +) + +from numpy.lib.arraysetops import ( + ediff1d as ediff1d, + intersect1d as intersect1d, + setxor1d as setxor1d, + union1d as union1d, + setdiff1d as setdiff1d, + unique as unique, + in1d as in1d, + isin as isin, +) + +from numpy.lib.arrayterator import ( + Arrayterator as Arrayterator, +) + +from numpy.lib.function_base import ( + select as select, + piecewise as piecewise, + trim_zeros as trim_zeros, + copy as copy, + iterable as iterable, + percentile as percentile, + diff as diff, + gradient as gradient, + angle as angle, + unwrap as unwrap, + sort_complex as sort_complex, + disp as disp, + flip as flip, + rot90 as rot90, + extract as extract, + place as place, + asarray_chkfinite as asarray_chkfinite, + average as average, + bincount as bincount, + digitize as digitize, + cov as cov, + corrcoef as corrcoef, + median as median, + sinc as sinc, + hamming as hamming, + hanning as hanning, + bartlett as bartlett, + blackman as blackman, + kaiser as kaiser, + trapz as trapz, + i0 as i0, + add_newdoc as add_newdoc, + add_docstring as add_docstring, + meshgrid as meshgrid, + delete as delete, + insert as insert, + append as append, + interp as interp, + add_newdoc_ufunc as add_newdoc_ufunc, + quantile as quantile, +) + +from numpy.lib.histograms import ( + histogram_bin_edges as histogram_bin_edges, + histogram as histogram, + histogramdd as histogramdd, +) + +from numpy.lib.index_tricks import ( + ravel_multi_index as ravel_multi_index, + unravel_index as unravel_index, + mgrid as mgrid, + ogrid as ogrid, + r_ as r_, + c_ as c_, + s_ as s_, + index_exp as index_exp, + ix_ as ix_, + fill_diagonal as fill_diagonal, + diag_indices as diag_indices, + diag_indices_from as diag_indices_from, +) + +from numpy.lib.nanfunctions import ( + nansum as nansum, + nanmax as nanmax, + nanmin as nanmin, + nanargmax as nanargmax, + nanargmin as nanargmin, + nanmean as nanmean, + nanmedian as nanmedian, + nanpercentile as nanpercentile, + nanvar as nanvar, + nanstd as nanstd, + nanprod as nanprod, + nancumsum as nancumsum, + nancumprod as nancumprod, + nanquantile as nanquantile, +) + +from numpy.lib.npyio import ( + savetxt as savetxt, + loadtxt as loadtxt, + genfromtxt as genfromtxt, + recfromtxt as recfromtxt, + recfromcsv as recfromcsv, + load as load, + save as save, + savez as savez, + savez_compressed as savez_compressed, + packbits as packbits, + unpackbits as unpackbits, + fromregex as fromregex, +) + +from numpy.lib.polynomial import ( + poly as poly, + roots as roots, + polyint as polyint, + polyder as polyder, + polyadd as polyadd, + polysub as polysub, + polymul as polymul, + polydiv as polydiv, + polyval as polyval, + polyfit as polyfit, +) + +from numpy.lib.shape_base import ( + column_stack as column_stack, + row_stack as row_stack, + dstack as dstack, + array_split as array_split, + split as split, + hsplit as hsplit, + vsplit as vsplit, + dsplit as dsplit, + apply_over_axes as apply_over_axes, + expand_dims as expand_dims, + apply_along_axis as apply_along_axis, + kron as kron, + tile as tile, + get_array_wrap as get_array_wrap, + take_along_axis as take_along_axis, + put_along_axis as put_along_axis, +) + +from numpy.lib.stride_tricks import ( + broadcast_to as broadcast_to, + broadcast_arrays as broadcast_arrays, + broadcast_shapes as broadcast_shapes, +) + +from numpy.lib.twodim_base import ( + diag as diag, + diagflat as diagflat, + eye as eye, + fliplr as fliplr, + flipud as flipud, + tri as tri, + triu as triu, + tril as tril, + vander as vander, + histogram2d as histogram2d, + mask_indices as mask_indices, + tril_indices as tril_indices, + tril_indices_from as tril_indices_from, + triu_indices as triu_indices, + triu_indices_from as triu_indices_from, +) + +from numpy.lib.type_check import ( + mintypecode as mintypecode, + asfarray as asfarray, + real as real, + imag as imag, + iscomplex as iscomplex, + isreal as isreal, + iscomplexobj as iscomplexobj, + isrealobj as isrealobj, + nan_to_num as nan_to_num, + real_if_close as real_if_close, + typename as typename, + common_type as common_type, +) + +from numpy.lib.ufunclike import ( + fix as fix, + isposinf as isposinf, + isneginf as isneginf, +) + +from numpy.lib.utils import ( + issubclass_ as issubclass_, + issubsctype as issubsctype, + issubdtype as issubdtype, + deprecate as deprecate, + deprecate_with_doc as deprecate_with_doc, + get_include as get_include, + info as info, + source as source, + who as who, + lookfor as lookfor, + byte_bounds as byte_bounds, + safe_eval as safe_eval, + show_runtime as show_runtime, +) + +from numpy.matrixlib import ( + asmatrix as asmatrix, + mat as mat, + bmat as bmat, +) + +_AnyStr_contra = TypeVar("_AnyStr_contra", str, bytes, contravariant=True) + +# Protocol for representing file-like-objects accepted +# by `ndarray.tofile` and `fromfile` +class _IOProtocol(Protocol): + def flush(self) -> object: ... + def fileno(self) -> int: ... + def tell(self) -> SupportsIndex: ... + def seek(self, offset: int, whence: int, /) -> object: ... + +# NOTE: `seek`, `write` and `flush` are technically only required +# for `readwrite`/`write` modes +class _MemMapIOProtocol(Protocol): + def flush(self) -> object: ... + def fileno(self) -> SupportsIndex: ... + def tell(self) -> int: ... + def seek(self, offset: int, whence: int, /) -> object: ... + def write(self, s: bytes, /) -> object: ... + @property + def read(self) -> object: ... + +class _SupportsWrite(Protocol[_AnyStr_contra]): + def write(self, s: _AnyStr_contra, /) -> object: ... + +__all__: list[str] +__path__: list[str] +__version__: str +test: PytestTester + +# TODO: Move placeholders to their respective module once +# their annotations are properly implemented +# +# Placeholders for classes + +def show_config() -> None: ... + +_NdArraySubClass = TypeVar("_NdArraySubClass", bound=ndarray[Any, Any]) +_DTypeScalar_co = TypeVar("_DTypeScalar_co", covariant=True, bound=generic) +_ByteOrder = L["S", "<", ">", "=", "|", "L", "B", "N", "I"] + +@final +class dtype(Generic[_DTypeScalar_co]): + names: None | tuple[builtins.str, ...] + def __hash__(self) -> int: ... + # Overload for subclass of generic + @overload + def __new__( + cls, + dtype: type[_DTypeScalar_co], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + # Overloads for string aliases, Python types, and some assorted + # other special cases. Order is sometimes important because of the + # subtype relationships + # + # bool < int < float < complex < object + # + # so we have to make sure the overloads for the narrowest type is + # first. + # Builtin types + @overload + def __new__(cls, dtype: type[bool], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bool_]: ... + @overload + def __new__(cls, dtype: type[int], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int_]: ... + @overload + def __new__(cls, dtype: None | type[float], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float_]: ... + @overload + def __new__(cls, dtype: type[complex], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex_]: ... + @overload + def __new__(cls, dtype: type[builtins.str], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[str_]: ... + @overload + def __new__(cls, dtype: type[bytes], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bytes_]: ... + + # `unsignedinteger` string-based representations and ctypes + @overload + def __new__(cls, dtype: _UInt8Codes | type[ct.c_uint8], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint8]: ... + @overload + def __new__(cls, dtype: _UInt16Codes | type[ct.c_uint16], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint16]: ... + @overload + def __new__(cls, dtype: _UInt32Codes | type[ct.c_uint32], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint32]: ... + @overload + def __new__(cls, dtype: _UInt64Codes | type[ct.c_uint64], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint64]: ... + @overload + def __new__(cls, dtype: _UByteCodes | type[ct.c_ubyte], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ubyte]: ... + @overload + def __new__(cls, dtype: _UShortCodes | type[ct.c_ushort], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ushort]: ... + @overload + def __new__(cls, dtype: _UIntCCodes | type[ct.c_uint], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uintc]: ... + + # NOTE: We're assuming here that `uint_ptr_t == size_t`, + # an assumption that does not hold in rare cases (same for `ssize_t`) + @overload + def __new__(cls, dtype: _UIntPCodes | type[ct.c_void_p] | type[ct.c_size_t], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uintp]: ... + @overload + def __new__(cls, dtype: _UIntCodes | type[ct.c_ulong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[uint]: ... + @overload + def __new__(cls, dtype: _ULongLongCodes | type[ct.c_ulonglong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[ulonglong]: ... + + # `signedinteger` string-based representations and ctypes + @overload + def __new__(cls, dtype: _Int8Codes | type[ct.c_int8], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int8]: ... + @overload + def __new__(cls, dtype: _Int16Codes | type[ct.c_int16], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int16]: ... + @overload + def __new__(cls, dtype: _Int32Codes | type[ct.c_int32], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int32]: ... + @overload + def __new__(cls, dtype: _Int64Codes | type[ct.c_int64], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int64]: ... + @overload + def __new__(cls, dtype: _ByteCodes | type[ct.c_byte], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[byte]: ... + @overload + def __new__(cls, dtype: _ShortCodes | type[ct.c_short], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[short]: ... + @overload + def __new__(cls, dtype: _IntCCodes | type[ct.c_int], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[intc]: ... + @overload + def __new__(cls, dtype: _IntPCodes | type[ct.c_ssize_t], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[intp]: ... + @overload + def __new__(cls, dtype: _IntCodes | type[ct.c_long], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[int_]: ... + @overload + def __new__(cls, dtype: _LongLongCodes | type[ct.c_longlong], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[longlong]: ... + + # `floating` string-based representations and ctypes + @overload + def __new__(cls, dtype: _Float16Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float16]: ... + @overload + def __new__(cls, dtype: _Float32Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float32]: ... + @overload + def __new__(cls, dtype: _Float64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[float64]: ... + @overload + def __new__(cls, dtype: _HalfCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[half]: ... + @overload + def __new__(cls, dtype: _SingleCodes | type[ct.c_float], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[single]: ... + @overload + def __new__(cls, dtype: _DoubleCodes | type[ct.c_double], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[double]: ... + @overload + def __new__(cls, dtype: _LongDoubleCodes | type[ct.c_longdouble], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[longdouble]: ... + + # `complexfloating` string-based representations + @overload + def __new__(cls, dtype: _Complex64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex64]: ... + @overload + def __new__(cls, dtype: _Complex128Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[complex128]: ... + @overload + def __new__(cls, dtype: _CSingleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[csingle]: ... + @overload + def __new__(cls, dtype: _CDoubleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[cdouble]: ... + @overload + def __new__(cls, dtype: _CLongDoubleCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[clongdouble]: ... + + # Miscellaneous string-based representations and ctypes + @overload + def __new__(cls, dtype: _BoolCodes | type[ct.c_bool], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bool_]: ... + @overload + def __new__(cls, dtype: _TD64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[timedelta64]: ... + @overload + def __new__(cls, dtype: _DT64Codes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[datetime64]: ... + @overload + def __new__(cls, dtype: _StrCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[str_]: ... + @overload + def __new__(cls, dtype: _BytesCodes | type[ct.c_char], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[bytes_]: ... + @overload + def __new__(cls, dtype: _VoidCodes, align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[void]: ... + @overload + def __new__(cls, dtype: _ObjectCodes | type[ct.py_object[Any]], align: bool = ..., copy: bool = ..., metadata: dict[builtins.str, Any] = ...) -> dtype[object_]: ... + + # dtype of a dtype is the same dtype + @overload + def __new__( + cls, + dtype: dtype[_DTypeScalar_co], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + @overload + def __new__( + cls, + dtype: _SupportsDType[dtype[_DTypeScalar_co]], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[_DTypeScalar_co]: ... + # Handle strings that can't be expressed as literals; i.e. s1, s2, ... + @overload + def __new__( + cls, + dtype: builtins.str, + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[Any]: ... + # Catchall overload for void-likes + @overload + def __new__( + cls, + dtype: _VoidDTypeLike, + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[void]: ... + # Catchall overload for object-likes + @overload + def __new__( + cls, + dtype: type[object], + align: bool = ..., + copy: bool = ..., + metadata: dict[builtins.str, Any] = ..., + ) -> dtype[object_]: ... + + def __class_getitem__(self, item: Any) -> GenericAlias: ... + + @overload + def __getitem__(self: dtype[void], key: list[builtins.str]) -> dtype[void]: ... + @overload + def __getitem__(self: dtype[void], key: builtins.str | SupportsIndex) -> dtype[Any]: ... + + # NOTE: In the future 1-based multiplications will also yield `flexible` dtypes + @overload + def __mul__(self: _DType, value: L[1]) -> _DType: ... + @overload + def __mul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ... + @overload + def __mul__(self, value: SupportsIndex) -> dtype[void]: ... + + # NOTE: `__rmul__` seems to be broken when used in combination with + # literals as of mypy 0.902. Set the return-type to `dtype[Any]` for + # now for non-flexible dtypes. + @overload + def __rmul__(self: _FlexDType, value: SupportsIndex) -> _FlexDType: ... + @overload + def __rmul__(self, value: SupportsIndex) -> dtype[Any]: ... + + def __gt__(self, other: DTypeLike) -> bool: ... + def __ge__(self, other: DTypeLike) -> bool: ... + def __lt__(self, other: DTypeLike) -> bool: ... + def __le__(self, other: DTypeLike) -> bool: ... + + # Explicitly defined `__eq__` and `__ne__` to get around mypy's + # `strict_equality` option; even though their signatures are + # identical to their `object`-based counterpart + def __eq__(self, other: Any) -> bool: ... + def __ne__(self, other: Any) -> bool: ... + + @property + def alignment(self) -> int: ... + @property + def base(self) -> dtype[Any]: ... + @property + def byteorder(self) -> builtins.str: ... + @property + def char(self) -> builtins.str: ... + @property + def descr(self) -> list[tuple[builtins.str, builtins.str] | tuple[builtins.str, builtins.str, _Shape]]: ... + @property + def fields( + self, + ) -> None | MappingProxyType[builtins.str, tuple[dtype[Any], int] | tuple[dtype[Any], int, Any]]: ... + @property + def flags(self) -> int: ... + @property + def hasobject(self) -> bool: ... + @property + def isbuiltin(self) -> int: ... + @property + def isnative(self) -> bool: ... + @property + def isalignedstruct(self) -> bool: ... + @property + def itemsize(self) -> int: ... + @property + def kind(self) -> builtins.str: ... + @property + def metadata(self) -> None | MappingProxyType[builtins.str, Any]: ... + @property + def name(self) -> builtins.str: ... + @property + def num(self) -> int: ... + @property + def shape(self) -> _Shape: ... + @property + def ndim(self) -> int: ... + @property + def subdtype(self) -> None | tuple[dtype[Any], _Shape]: ... + def newbyteorder(self: _DType, __new_order: _ByteOrder = ...) -> _DType: ... + @property + def str(self) -> builtins.str: ... + @property + def type(self) -> type[_DTypeScalar_co]: ... + +_ArrayLikeInt = Union[ + int, + integer[Any], + Sequence[Union[int, integer[Any]]], + Sequence[Sequence[Any]], # TODO: wait for support for recursive types + ndarray[Any, Any] +] + +_FlatIterSelf = TypeVar("_FlatIterSelf", bound=flatiter[Any]) + +@final +class flatiter(Generic[_NdArraySubClass]): + __hash__: ClassVar[None] + @property + def base(self) -> _NdArraySubClass: ... + @property + def coords(self) -> _Shape: ... + @property + def index(self) -> int: ... + def copy(self) -> _NdArraySubClass: ... + def __iter__(self: _FlatIterSelf) -> _FlatIterSelf: ... + def __next__(self: flatiter[ndarray[Any, dtype[_ScalarType]]]) -> _ScalarType: ... + def __len__(self) -> int: ... + @overload + def __getitem__( + self: flatiter[ndarray[Any, dtype[_ScalarType]]], + key: int | integer[Any] | tuple[int | integer[Any]], + ) -> _ScalarType: ... + @overload + def __getitem__( + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], + ) -> _NdArraySubClass: ... + # TODO: `__setitem__` operates via `unsafe` casting rules, and can + # thus accept any type accepted by the relevant underlying `np.generic` + # constructor. + # This means that `value` must in reality be a supertype of `npt.ArrayLike`. + def __setitem__( + self, + key: _ArrayLikeInt | slice | ellipsis | tuple[_ArrayLikeInt | slice | ellipsis], + value: Any, + ) -> None: ... + @overload + def __array__(self: flatiter[ndarray[Any, _DType]], dtype: None = ..., /) -> ndarray[Any, _DType]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + +_OrderKACF = L[None, "K", "A", "C", "F"] +_OrderACF = L[None, "A", "C", "F"] +_OrderCF = L[None, "C", "F"] + +_ModeKind = L["raise", "wrap", "clip"] +_PartitionKind = L["introselect"] +_SortKind = L["quicksort", "mergesort", "heapsort", "stable"] +_SortSide = L["left", "right"] + +_ArraySelf = TypeVar("_ArraySelf", bound=_ArrayOrScalarCommon) + +class _ArrayOrScalarCommon: + @property + def T(self: _ArraySelf) -> _ArraySelf: ... + @property + def data(self) -> memoryview: ... + @property + def flags(self) -> flagsobj: ... + @property + def itemsize(self) -> int: ... + @property + def nbytes(self) -> int: ... + def __bool__(self) -> bool: ... + def __bytes__(self) -> bytes: ... + def __str__(self) -> str: ... + def __repr__(self) -> str: ... + def __copy__(self: _ArraySelf) -> _ArraySelf: ... + def __deepcopy__(self: _ArraySelf, memo: None | dict[int, Any], /) -> _ArraySelf: ... + + # TODO: How to deal with the non-commutative nature of `==` and `!=`? + # xref numpy/numpy#17368 + def __eq__(self, other: Any) -> Any: ... + def __ne__(self, other: Any) -> Any: ... + def copy(self: _ArraySelf, order: _OrderKACF = ...) -> _ArraySelf: ... + def dump(self, file: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _SupportsWrite[bytes]) -> None: ... + def dumps(self) -> bytes: ... + def tobytes(self, order: _OrderKACF = ...) -> bytes: ... + # NOTE: `tostring()` is deprecated and therefore excluded + # def tostring(self, order=...): ... + def tofile( + self, + fid: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _IOProtocol, + sep: str = ..., + format: str = ..., + ) -> None: ... + # generics and 0d arrays return builtin scalars + def tolist(self) -> Any: ... + + @property + def __array_interface__(self) -> dict[str, Any]: ... + @property + def __array_priority__(self) -> float: ... + @property + def __array_struct__(self) -> Any: ... # builtins.PyCapsule + def __setstate__(self, state: tuple[ + SupportsIndex, # version + _ShapeLike, # Shape + _DType_co, # DType + bool, # F-continuous + bytes | list[Any], # Data + ], /) -> None: ... + # a `bool_` is returned when `keepdims=True` and `self` is a 0d array + + @overload + def all( + self, + axis: None = ..., + out: None = ..., + keepdims: L[False] = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> bool_: ... + @overload + def all( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def all( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def any( + self, + axis: None = ..., + out: None = ..., + keepdims: L[False] = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> bool_: ... + @overload + def any( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def any( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def argmax( + self, + axis: None = ..., + out: None = ..., + *, + keepdims: L[False] = ..., + ) -> intp: ... + @overload + def argmax( + self, + axis: SupportsIndex = ..., + out: None = ..., + *, + keepdims: bool = ..., + ) -> Any: ... + @overload + def argmax( + self, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + *, + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + @overload + def argmin( + self, + axis: None = ..., + out: None = ..., + *, + keepdims: L[False] = ..., + ) -> intp: ... + @overload + def argmin( + self, + axis: SupportsIndex = ..., + out: None = ..., + *, + keepdims: bool = ..., + ) -> Any: ... + @overload + def argmin( + self, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + *, + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + def argsort( + self, + axis: None | SupportsIndex = ..., + kind: None | _SortKind = ..., + order: None | str | Sequence[str] = ..., + ) -> ndarray[Any, Any]: ... + + @overload + def choose( + self, + choices: ArrayLike, + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, Any]: ... + @overload + def choose( + self, + choices: ArrayLike, + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + @overload + def clip( + self, + min: ArrayLike = ..., + max: None | ArrayLike = ..., + out: None = ..., + **kwargs: Any, + ) -> ndarray[Any, Any]: ... + @overload + def clip( + self, + min: None = ..., + max: ArrayLike = ..., + out: None = ..., + **kwargs: Any, + ) -> ndarray[Any, Any]: ... + @overload + def clip( + self, + min: ArrayLike = ..., + max: None | ArrayLike = ..., + out: _NdArraySubClass = ..., + **kwargs: Any, + ) -> _NdArraySubClass: ... + @overload + def clip( + self, + min: None = ..., + max: ArrayLike = ..., + out: _NdArraySubClass = ..., + **kwargs: Any, + ) -> _NdArraySubClass: ... + + @overload + def compress( + self, + a: ArrayLike, + axis: None | SupportsIndex = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def compress( + self, + a: ArrayLike, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + def conj(self: _ArraySelf) -> _ArraySelf: ... + + def conjugate(self: _ArraySelf) -> _ArraySelf: ... + + @overload + def cumprod( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def cumprod( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def cumsum( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> ndarray[Any, Any]: ... + @overload + def cumsum( + self, + axis: None | SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def max( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def max( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def mean( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def mean( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def min( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def min( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + def newbyteorder( + self: _ArraySelf, + __new_order: _ByteOrder = ..., + ) -> _ArraySelf: ... + + @overload + def prod( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def prod( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def ptp( + self, + axis: None | _ShapeLike = ..., + out: None = ..., + keepdims: bool = ..., + ) -> Any: ... + @overload + def ptp( + self, + axis: None | _ShapeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + ) -> _NdArraySubClass: ... + + @overload + def round( + self: _ArraySelf, + decimals: SupportsIndex = ..., + out: None = ..., + ) -> _ArraySelf: ... + @overload + def round( + self, + decimals: SupportsIndex = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def std( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def std( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def sum( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def sum( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + keepdims: bool = ..., + initial: _NumberLike_co = ..., + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + + @overload + def var( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: None = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> Any: ... + @overload + def var( + self, + axis: None | _ShapeLike = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ddof: float = ..., + keepdims: bool = ..., + *, + where: _ArrayLikeBool_co = ..., + ) -> _NdArraySubClass: ... + +_DType = TypeVar("_DType", bound=dtype[Any]) +_DType_co = TypeVar("_DType_co", covariant=True, bound=dtype[Any]) +_FlexDType = TypeVar("_FlexDType", bound=dtype[flexible]) + +# TODO: Set the `bound` to something more suitable once we +# have proper shape support +_ShapeType = TypeVar("_ShapeType", bound=Any) +_ShapeType2 = TypeVar("_ShapeType2", bound=Any) +_NumberType = TypeVar("_NumberType", bound=number[Any]) + +if sys.version_info >= (3, 12): + from collections.abc import Buffer as _SupportsBuffer +else: + _SupportsBuffer = ( + bytes + | bytearray + | memoryview + | _array.array[Any] + | mmap.mmap + | NDArray[Any] + | generic + ) + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +_T_contra = TypeVar("_T_contra", contravariant=True) +_2Tuple = tuple[_T, _T] +_CastingKind = L["no", "equiv", "safe", "same_kind", "unsafe"] + +_ArrayUInt_co = NDArray[Union[bool_, unsignedinteger[Any]]] +_ArrayInt_co = NDArray[Union[bool_, integer[Any]]] +_ArrayFloat_co = NDArray[Union[bool_, integer[Any], floating[Any]]] +_ArrayComplex_co = NDArray[Union[bool_, integer[Any], floating[Any], complexfloating[Any, Any]]] +_ArrayNumber_co = NDArray[Union[bool_, number[Any]]] +_ArrayTD64_co = NDArray[Union[bool_, integer[Any], timedelta64]] + +# Introduce an alias for `dtype` to avoid naming conflicts. +_dtype = dtype + +# `builtins.PyCapsule` unfortunately lacks annotations as of the moment; +# use `Any` as a stopgap measure +_PyCapsule = Any + +class _SupportsItem(Protocol[_T_co]): + def item(self, args: Any, /) -> _T_co: ... + +class _SupportsReal(Protocol[_T_co]): + @property + def real(self) -> _T_co: ... + +class _SupportsImag(Protocol[_T_co]): + @property + def imag(self) -> _T_co: ... + +class ndarray(_ArrayOrScalarCommon, Generic[_ShapeType, _DType_co]): + __hash__: ClassVar[None] + @property + def base(self) -> None | ndarray[Any, Any]: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def real( + self: ndarray[_ShapeType, dtype[_SupportsReal[_ScalarType]]], # type: ignore[type-var] + ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ... + @real.setter + def real(self, value: ArrayLike) -> None: ... + @property + def imag( + self: ndarray[_ShapeType, dtype[_SupportsImag[_ScalarType]]], # type: ignore[type-var] + ) -> ndarray[_ShapeType, _dtype[_ScalarType]]: ... + @imag.setter + def imag(self, value: ArrayLike) -> None: ... + def __new__( + cls: type[_ArraySelf], + shape: _ShapeLike, + dtype: DTypeLike = ..., + buffer: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> _ArraySelf: ... + + if sys.version_info >= (3, 12): + def __buffer__(self, flags: int, /) -> memoryview: ... + + def __class_getitem__(self, item: Any) -> GenericAlias: ... + + @overload + def __array__(self, dtype: None = ..., /) -> ndarray[Any, _DType_co]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + + def __array_ufunc__( + self, + ufunc: ufunc, + method: L["__call__", "reduce", "reduceat", "accumulate", "outer", "inner"], + *inputs: Any, + **kwargs: Any, + ) -> Any: ... + + def __array_function__( + self, + func: Callable[..., Any], + types: Iterable[type], + args: Iterable[Any], + kwargs: Mapping[str, Any], + ) -> Any: ... + + # NOTE: In practice any object is accepted by `obj`, but as `__array_finalize__` + # is a pseudo-abstract method the type has been narrowed down in order to + # grant subclasses a bit more flexibility + def __array_finalize__(self, obj: None | NDArray[Any], /) -> None: ... + + def __array_wrap__( + self, + array: ndarray[_ShapeType2, _DType], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + /, + ) -> ndarray[_ShapeType2, _DType]: ... + + def __array_prepare__( + self, + array: ndarray[_ShapeType2, _DType], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + /, + ) -> ndarray[_ShapeType2, _DType]: ... + + @overload + def __getitem__(self, key: ( + NDArray[integer[Any]] + | NDArray[bool_] + | tuple[NDArray[integer[Any]] | NDArray[bool_], ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self, key: SupportsIndex | tuple[SupportsIndex, ...]) -> Any: ... + @overload + def __getitem__(self, key: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self: NDArray[void], key: str) -> NDArray[Any]: ... + @overload + def __getitem__(self: NDArray[void], key: list[str]) -> ndarray[_ShapeType, _dtype[void]]: ... + + @property + def ctypes(self) -> _ctypes[int]: ... + @property + def shape(self) -> _Shape: ... + @shape.setter + def shape(self, value: _ShapeLike) -> None: ... + @property + def strides(self) -> _Shape: ... + @strides.setter + def strides(self, value: _ShapeLike) -> None: ... + def byteswap(self: _ArraySelf, inplace: bool = ...) -> _ArraySelf: ... + def fill(self, value: Any) -> None: ... + @property + def flat(self: _NdArraySubClass) -> flatiter[_NdArraySubClass]: ... + + # Use the same output type as that of the underlying `generic` + @overload + def item( + self: ndarray[Any, _dtype[_SupportsItem[_T]]], # type: ignore[type-var] + *args: SupportsIndex, + ) -> _T: ... + @overload + def item( + self: ndarray[Any, _dtype[_SupportsItem[_T]]], # type: ignore[type-var] + args: tuple[SupportsIndex, ...], + /, + ) -> _T: ... + + @overload + def itemset(self, value: Any, /) -> None: ... + @overload + def itemset(self, item: _ShapeLike, value: Any, /) -> None: ... + + @overload + def resize(self, new_shape: _ShapeLike, /, *, refcheck: bool = ...) -> None: ... + @overload + def resize(self, *new_shape: SupportsIndex, refcheck: bool = ...) -> None: ... + + def setflags( + self, write: bool = ..., align: bool = ..., uic: bool = ... + ) -> None: ... + + def squeeze( + self, + axis: None | SupportsIndex | tuple[SupportsIndex, ...] = ..., + ) -> ndarray[Any, _DType_co]: ... + + def swapaxes( + self, + axis1: SupportsIndex, + axis2: SupportsIndex, + ) -> ndarray[Any, _DType_co]: ... + + @overload + def transpose(self: _ArraySelf, axes: None | _ShapeLike, /) -> _ArraySelf: ... + @overload + def transpose(self: _ArraySelf, *axes: SupportsIndex) -> _ArraySelf: ... + + def argpartition( + self, + kth: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + kind: _PartitionKind = ..., + order: None | str | Sequence[str] = ..., + ) -> ndarray[Any, _dtype[intp]]: ... + + def diagonal( + self, + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + ) -> ndarray[Any, _DType_co]: ... + + # 1D + 1D returns a scalar; + # all other with at least 1 non-0D array return an ndarray. + @overload + def dot(self, b: _ScalarLike_co, out: None = ...) -> ndarray[Any, Any]: ... + @overload + def dot(self, b: ArrayLike, out: None = ...) -> Any: ... # type: ignore[misc] + @overload + def dot(self, b: ArrayLike, out: _NdArraySubClass) -> _NdArraySubClass: ... + + # `nonzero()` is deprecated for 0d arrays/generics + def nonzero(self) -> tuple[ndarray[Any, _dtype[intp]], ...]: ... + + def partition( + self, + kth: _ArrayLikeInt_co, + axis: SupportsIndex = ..., + kind: _PartitionKind = ..., + order: None | str | Sequence[str] = ..., + ) -> None: ... + + # `put` is technically available to `generic`, + # but is pointless as `generic`s are immutable + def put( + self, + ind: _ArrayLikeInt_co, + v: ArrayLike, + mode: _ModeKind = ..., + ) -> None: ... + + @overload + def searchsorted( # type: ignore[misc] + self, # >= 1D array + v: _ScalarLike_co, # 0D array-like + side: _SortSide = ..., + sorter: None | _ArrayLikeInt_co = ..., + ) -> intp: ... + @overload + def searchsorted( + self, # >= 1D array + v: ArrayLike, + side: _SortSide = ..., + sorter: None | _ArrayLikeInt_co = ..., + ) -> ndarray[Any, _dtype[intp]]: ... + + def setfield( + self, + val: ArrayLike, + dtype: DTypeLike, + offset: SupportsIndex = ..., + ) -> None: ... + + def sort( + self, + axis: SupportsIndex = ..., + kind: None | _SortKind = ..., + order: None | str | Sequence[str] = ..., + ) -> None: ... + + @overload + def trace( + self, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + out: None = ..., + ) -> Any: ... + @overload + def trace( + self, # >= 2D array + offset: SupportsIndex = ..., + axis1: SupportsIndex = ..., + axis2: SupportsIndex = ..., + dtype: DTypeLike = ..., + out: _NdArraySubClass = ..., + ) -> _NdArraySubClass: ... + + @overload + def take( # type: ignore[misc] + self: ndarray[Any, _dtype[_ScalarType]], + indices: _IntLike_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> _ScalarType: ... + @overload + def take( # type: ignore[misc] + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, _DType_co]: ... + @overload + def take( + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + def repeat( + self, + repeats: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + ) -> ndarray[Any, _DType_co]: ... + + def flatten( + self, + order: _OrderKACF = ..., + ) -> ndarray[Any, _DType_co]: ... + + def ravel( + self, + order: _OrderKACF = ..., + ) -> ndarray[Any, _DType_co]: ... + + @overload + def reshape( + self, shape: _ShapeLike, /, *, order: _OrderACF = ... + ) -> ndarray[Any, _DType_co]: ... + @overload + def reshape( + self, *shape: SupportsIndex, order: _OrderACF = ... + ) -> ndarray[Any, _DType_co]: ... + + @overload + def astype( + self, + dtype: _DTypeLike[_ScalarType], + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> NDArray[_ScalarType]: ... + @overload + def astype( + self, + dtype: DTypeLike, + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> NDArray[Any]: ... + + @overload + def view(self: _ArraySelf) -> _ArraySelf: ... + @overload + def view(self, type: type[_NdArraySubClass]) -> _NdArraySubClass: ... + @overload + def view(self, dtype: _DTypeLike[_ScalarType]) -> NDArray[_ScalarType]: ... + @overload + def view(self, dtype: DTypeLike) -> NDArray[Any]: ... + @overload + def view( + self, + dtype: DTypeLike, + type: type[_NdArraySubClass], + ) -> _NdArraySubClass: ... + + @overload + def getfield( + self, + dtype: _DTypeLike[_ScalarType], + offset: SupportsIndex = ... + ) -> NDArray[_ScalarType]: ... + @overload + def getfield( + self, + dtype: DTypeLike, + offset: SupportsIndex = ... + ) -> NDArray[Any]: ... + + # Dispatch to the underlying `generic` via protocols + def __int__( + self: ndarray[Any, _dtype[SupportsInt]], # type: ignore[type-var] + ) -> int: ... + + def __float__( + self: ndarray[Any, _dtype[SupportsFloat]], # type: ignore[type-var] + ) -> float: ... + + def __complex__( + self: ndarray[Any, _dtype[SupportsComplex]], # type: ignore[type-var] + ) -> complex: ... + + def __index__( + self: ndarray[Any, _dtype[SupportsIndex]], # type: ignore[type-var] + ) -> int: ... + + def __len__(self) -> int: ... + def __setitem__(self, key, value): ... + def __iter__(self) -> Any: ... + def __contains__(self, key) -> bool: ... + + # The last overload is for catching recursive objects whose + # nesting is too deep. + # The first overload is for catching `bytes` (as they are a subtype of + # `Sequence[int]`) and `str`. As `str` is a recursive sequence of + # strings, it will pass through the final overload otherwise + + @overload + def __lt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __lt__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __le__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __le__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __le__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __gt__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __gt__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + @overload + def __ge__(self: _ArrayNumber_co, other: _ArrayLikeNumber_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[object_], other: Any) -> NDArray[bool_]: ... + @overload + def __ge__(self: NDArray[Any], other: _ArrayLikeObject_co) -> NDArray[bool_]: ... + + # Unary ops + @overload + def __abs__(self: NDArray[bool_]) -> NDArray[bool_]: ... + @overload + def __abs__(self: NDArray[complexfloating[_NBit1, _NBit1]]) -> NDArray[floating[_NBit1]]: ... + @overload + def __abs__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __abs__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __abs__(self: NDArray[object_]) -> Any: ... + + @overload + def __invert__(self: NDArray[bool_]) -> NDArray[bool_]: ... + @overload + def __invert__(self: NDArray[_IntType]) -> NDArray[_IntType]: ... + @overload + def __invert__(self: NDArray[object_]) -> Any: ... + + @overload + def __pos__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __pos__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __pos__(self: NDArray[object_]) -> Any: ... + + @overload + def __neg__(self: NDArray[_NumberType]) -> NDArray[_NumberType]: ... + @overload + def __neg__(self: NDArray[timedelta64]) -> NDArray[timedelta64]: ... + @overload + def __neg__(self: NDArray[object_]) -> Any: ... + + # Binary ops + @overload + def __matmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __matmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __matmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __matmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __matmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmatmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmatmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __rmatmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rmatmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmatmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __mod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __mod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __mod__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __mod__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __rmod__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmod__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __divmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> _2Tuple[NDArray[int8]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> _2Tuple[NDArray[unsignedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> _2Tuple[NDArray[signedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ... # type: ignore[misc] + @overload + def __divmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> tuple[NDArray[int64], NDArray[timedelta64]]: ... + + @overload + def __rdivmod__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> _2Tuple[NDArray[int8]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> _2Tuple[NDArray[unsignedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> _2Tuple[NDArray[signedinteger[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> _2Tuple[NDArray[floating[Any]]]: ... # type: ignore[misc] + @overload + def __rdivmod__(self: _ArrayTD64_co, other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> tuple[NDArray[int64], NDArray[timedelta64]]: ... + + @overload + def __add__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __add__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __add__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __add__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... + @overload + def __add__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __add__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __add__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __radd__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __radd__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __radd__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __radd__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... + @overload + def __radd__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __radd__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __radd__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __sub__(self: NDArray[_UnknownType], other: _ArrayLikeUnknown) -> NDArray[Any]: ... + @overload + def __sub__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __sub__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __sub__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __sub__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __sub__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __sub__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[timedelta64]: ... + @overload + def __sub__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __sub__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rsub__(self: NDArray[_UnknownType], other: _ArrayLikeUnknown) -> NDArray[Any]: ... + @overload + def __rsub__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __rsub__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rsub__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rsub__(self: _ArrayTD64_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... # type: ignore[misc] + @overload + def __rsub__(self: _ArrayTD64_co, other: _ArrayLikeDT64_co) -> NDArray[datetime64]: ... # type: ignore[misc] + @overload + def __rsub__(self: NDArray[datetime64], other: _ArrayLikeDT64_co) -> NDArray[timedelta64]: ... + @overload + def __rsub__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rsub__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __mul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __mul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __mul__(self: _ArrayTD64_co, other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __mul__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __mul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __mul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rmul__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rmul__(self: _ArrayTD64_co, other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __rmul__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rmul__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rmul__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __floordiv__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __floordiv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[int64]: ... + @overload + def __floordiv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __floordiv__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __floordiv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __floordiv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rfloordiv__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rfloordiv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[int64]: ... + @overload + def __rfloordiv__(self: NDArray[bool_], other: _ArrayLikeTD64_co) -> NoReturn: ... + @overload + def __rfloordiv__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rfloordiv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rfloordiv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __pow__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __pow__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __pow__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __pow__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __pow__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rpow__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rpow__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... + @overload + def __rpow__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rpow__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rpow__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __truediv__(self: _ArrayInt_co, other: _ArrayInt_co) -> NDArray[float64]: ... # type: ignore[misc] + @overload + def __truediv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __truediv__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __truediv__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[float64]: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __truediv__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __truediv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __truediv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rtruediv__(self: _ArrayInt_co, other: _ArrayInt_co) -> NDArray[float64]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: _ArrayFloat_co, other: _ArrayLikeFloat_co) -> NDArray[floating[Any]]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: _ArrayComplex_co, other: _ArrayLikeComplex_co) -> NDArray[complexfloating[Any, Any]]: ... # type: ignore[misc] + @overload + def __rtruediv__(self: NDArray[number[Any]], other: _ArrayLikeNumber_co) -> NDArray[number[Any]]: ... + @overload + def __rtruediv__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[float64]: ... + @overload + def __rtruediv__(self: NDArray[bool_], other: _ArrayLikeTD64_co) -> NoReturn: ... + @overload + def __rtruediv__(self: _ArrayFloat_co, other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __rtruediv__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rtruediv__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __lshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __lshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __lshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __lshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __lshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rlshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rlshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rlshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rlshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rlshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rrshift__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[int8]: ... # type: ignore[misc] + @overload + def __rrshift__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rrshift__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rrshift__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rrshift__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __and__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __and__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __and__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __and__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __and__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rand__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rand__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rand__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rand__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rand__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __xor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __xor__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __xor__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __xor__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __xor__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __rxor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __rxor__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __rxor__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __rxor__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __rxor__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __or__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __or__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __or__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __or__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __or__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + @overload + def __ror__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... # type: ignore[misc] + @overload + def __ror__(self: _ArrayUInt_co, other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[Any]]: ... # type: ignore[misc] + @overload + def __ror__(self: _ArrayInt_co, other: _ArrayLikeInt_co) -> NDArray[signedinteger[Any]]: ... + @overload + def __ror__(self: NDArray[object_], other: Any) -> Any: ... + @overload + def __ror__(self: NDArray[Any], other: _ArrayLikeObject_co) -> Any: ... + + # `np.generic` does not support inplace operations + + # NOTE: Inplace ops generally use "same_kind" casting w.r.t. to the left + # operand. An exception to this rule are unsigned integers though, which + # also accepts a signed integer for the right operand as long it is a 0D + # object and its value is >= 0 + @overload + def __iadd__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __iadd__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __iadd__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __iadd__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __iadd__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __iadd__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __isub__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __isub__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __isub__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __isub__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __isub__(self: NDArray[timedelta64], other: _ArrayLikeTD64_co) -> NDArray[timedelta64]: ... + @overload + def __isub__(self: NDArray[datetime64], other: _ArrayLikeTD64_co) -> NDArray[datetime64]: ... + @overload + def __isub__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __imul__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imul__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __imul__(self: NDArray[timedelta64], other: _ArrayLikeFloat_co) -> NDArray[timedelta64]: ... + @overload + def __imul__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __itruediv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __itruediv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __itruediv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __itruediv__(self: NDArray[timedelta64], other: _ArrayLikeInt_co) -> NDArray[timedelta64]: ... + @overload + def __itruediv__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ifloordiv__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __ifloordiv__(self: NDArray[timedelta64], other: _ArrayLikeBool_co) -> NoReturn: ... + @overload + def __ifloordiv__(self: NDArray[timedelta64], other: _ArrayLikeInt_co) -> NDArray[timedelta64]: ... + @overload + def __ifloordiv__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ipow__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __ipow__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __ipow__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imod__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imod__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imod__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imod__(self: NDArray[timedelta64], other: _SupportsArray[_dtype[timedelta64]] | _NestedSequence[_SupportsArray[_dtype[timedelta64]]]) -> NDArray[timedelta64]: ... + @overload + def __imod__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ilshift__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ilshift__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ilshift__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __irshift__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __irshift__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __irshift__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __iand__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __iand__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __iand__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __iand__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ixor__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __ixor__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ixor__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ixor__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __ior__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __ior__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co | _IntLike_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __ior__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __ior__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + @overload + def __imatmul__(self: NDArray[bool_], other: _ArrayLikeBool_co) -> NDArray[bool_]: ... + @overload + def __imatmul__(self: NDArray[unsignedinteger[_NBit1]], other: _ArrayLikeUInt_co) -> NDArray[unsignedinteger[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[signedinteger[_NBit1]], other: _ArrayLikeInt_co) -> NDArray[signedinteger[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[floating[_NBit1]], other: _ArrayLikeFloat_co) -> NDArray[floating[_NBit1]]: ... + @overload + def __imatmul__(self: NDArray[complexfloating[_NBit1, _NBit1]], other: _ArrayLikeComplex_co) -> NDArray[complexfloating[_NBit1, _NBit1]]: ... + @overload + def __imatmul__(self: NDArray[object_], other: Any) -> NDArray[object_]: ... + + def __dlpack__(self: NDArray[number[Any]], *, stream: None = ...) -> _PyCapsule: ... + def __dlpack_device__(self) -> tuple[int, L[0]]: ... + + # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` + @property + def dtype(self) -> _DType_co: ... + +# NOTE: while `np.generic` is not technically an instance of `ABCMeta`, +# the `@abstractmethod` decorator is herein used to (forcefully) deny +# the creation of `np.generic` instances. +# The `# type: ignore` comments are necessary to silence mypy errors regarding +# the missing `ABCMeta` metaclass. + +# See https://github.com/numpy/numpy-stubs/pull/80 for more details. + +_ScalarType = TypeVar("_ScalarType", bound=generic) +_NBit1 = TypeVar("_NBit1", bound=NBitBase) +_NBit2 = TypeVar("_NBit2", bound=NBitBase) + +class generic(_ArrayOrScalarCommon): + @abstractmethod + def __init__(self, *args: Any, **kwargs: Any) -> None: ... + @overload + def __array__(self: _ScalarType, dtype: None = ..., /) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def __array__(self, dtype: _DType, /) -> ndarray[Any, _DType]: ... + def __hash__(self) -> int: ... + @property + def base(self) -> None: ... + @property + def ndim(self) -> L[0]: ... + @property + def size(self) -> L[1]: ... + @property + def shape(self) -> tuple[()]: ... + @property + def strides(self) -> tuple[()]: ... + def byteswap(self: _ScalarType, inplace: L[False] = ...) -> _ScalarType: ... + @property + def flat(self: _ScalarType) -> flatiter[ndarray[Any, _dtype[_ScalarType]]]: ... + + if sys.version_info >= (3, 12): + def __buffer__(self, flags: int, /) -> memoryview: ... + + @overload + def astype( + self, + dtype: _DTypeLike[_ScalarType], + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> _ScalarType: ... + @overload + def astype( + self, + dtype: DTypeLike, + order: _OrderKACF = ..., + casting: _CastingKind = ..., + subok: bool = ..., + copy: bool | _CopyMode = ..., + ) -> Any: ... + + # NOTE: `view` will perform a 0D->scalar cast, + # thus the array `type` is irrelevant to the output type + @overload + def view( + self: _ScalarType, + type: type[ndarray[Any, Any]] = ..., + ) -> _ScalarType: ... + @overload + def view( + self, + dtype: _DTypeLike[_ScalarType], + type: type[ndarray[Any, Any]] = ..., + ) -> _ScalarType: ... + @overload + def view( + self, + dtype: DTypeLike, + type: type[ndarray[Any, Any]] = ..., + ) -> Any: ... + + @overload + def getfield( + self, + dtype: _DTypeLike[_ScalarType], + offset: SupportsIndex = ... + ) -> _ScalarType: ... + @overload + def getfield( + self, + dtype: DTypeLike, + offset: SupportsIndex = ... + ) -> Any: ... + + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> Any: ... + + @overload + def take( # type: ignore[misc] + self: _ScalarType, + indices: _IntLike_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> _ScalarType: ... + @overload + def take( # type: ignore[misc] + self: _ScalarType, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: None = ..., + mode: _ModeKind = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def take( + self, + indices: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + out: _NdArraySubClass = ..., + mode: _ModeKind = ..., + ) -> _NdArraySubClass: ... + + def repeat( + self: _ScalarType, + repeats: _ArrayLikeInt_co, + axis: None | SupportsIndex = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def flatten( + self: _ScalarType, + order: _OrderKACF = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def ravel( + self: _ScalarType, + order: _OrderKACF = ..., + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + @overload + def reshape( + self: _ScalarType, shape: _ShapeLike, /, *, order: _OrderACF = ... + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + @overload + def reshape( + self: _ScalarType, *shape: SupportsIndex, order: _OrderACF = ... + ) -> ndarray[Any, _dtype[_ScalarType]]: ... + + def squeeze( + self: _ScalarType, axis: None | L[0] | tuple[()] = ... + ) -> _ScalarType: ... + def transpose(self: _ScalarType, axes: None | tuple[()] = ..., /) -> _ScalarType: ... + # Keep `dtype` at the bottom to avoid name conflicts with `np.dtype` + @property + def dtype(self: _ScalarType) -> _dtype[_ScalarType]: ... + +class number(generic, Generic[_NBit1]): # type: ignore + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def __class_getitem__(self, item: Any) -> GenericAlias: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __neg__(self: _ArraySelf) -> _ArraySelf: ... + def __pos__(self: _ArraySelf) -> _ArraySelf: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + # Ensure that objects annotated as `number` support arithmetic operations + __add__: _NumberOp + __radd__: _NumberOp + __sub__: _NumberOp + __rsub__: _NumberOp + __mul__: _NumberOp + __rmul__: _NumberOp + __floordiv__: _NumberOp + __rfloordiv__: _NumberOp + __pow__: _NumberOp + __rpow__: _NumberOp + __truediv__: _NumberOp + __rtruediv__: _NumberOp + __lt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __le__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __gt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __ge__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + +class bool_(generic): + def __init__(self, value: object = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> bool: ... + def tolist(self) -> bool: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + __add__: _BoolOp[bool_] + __radd__: _BoolOp[bool_] + __sub__: _BoolSub + __rsub__: _BoolSub + __mul__: _BoolOp[bool_] + __rmul__: _BoolOp[bool_] + __floordiv__: _BoolOp[int8] + __rfloordiv__: _BoolOp[int8] + __pow__: _BoolOp[int8] + __rpow__: _BoolOp[int8] + __truediv__: _BoolTrueDiv + __rtruediv__: _BoolTrueDiv + def __invert__(self) -> bool_: ... + __lshift__: _BoolBitOp[int8] + __rlshift__: _BoolBitOp[int8] + __rshift__: _BoolBitOp[int8] + __rrshift__: _BoolBitOp[int8] + __and__: _BoolBitOp[bool_] + __rand__: _BoolBitOp[bool_] + __xor__: _BoolBitOp[bool_] + __rxor__: _BoolBitOp[bool_] + __or__: _BoolBitOp[bool_] + __ror__: _BoolBitOp[bool_] + __mod__: _BoolMod + __rmod__: _BoolMod + __divmod__: _BoolDivMod + __rdivmod__: _BoolDivMod + __lt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __le__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __gt__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + __ge__: _ComparisonOp[_NumberLike_co, _ArrayLikeNumber_co] + +class object_(generic): + def __init__(self, value: object = ..., /) -> None: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + # The 3 protocols below may or may not raise, + # depending on the underlying object + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + + if sys.version_info >= (3, 12): + def __release_buffer__(self, buffer: memoryview, /) -> None: ... + +# The `datetime64` constructors requires an object with the three attributes below, +# and thus supports datetime duck typing +class _DatetimeScalar(Protocol): + @property + def day(self) -> int: ... + @property + def month(self) -> int: ... + @property + def year(self) -> int: ... + +# TODO: `item`/`tolist` returns either `dt.date`, `dt.datetime` or `int` +# depending on the unit +class datetime64(generic): + @overload + def __init__( + self, + value: None | datetime64 | _CharLike_co | _DatetimeScalar = ..., + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., + /, + ) -> None: ... + @overload + def __init__( + self, + value: int, + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co], + /, + ) -> None: ... + def __add__(self, other: _TD64Like_co) -> datetime64: ... + def __radd__(self, other: _TD64Like_co) -> datetime64: ... + @overload + def __sub__(self, other: datetime64) -> timedelta64: ... + @overload + def __sub__(self, other: _TD64Like_co) -> datetime64: ... + def __rsub__(self, other: datetime64) -> timedelta64: ... + __lt__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __le__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __gt__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + __ge__: _ComparisonOp[datetime64, _ArrayLikeDT64_co] + +_IntValue = Union[SupportsInt, _CharLike_co, SupportsIndex] +_FloatValue = Union[None, _CharLike_co, SupportsFloat, SupportsIndex] +_ComplexValue = Union[ + None, + _CharLike_co, + SupportsFloat, + SupportsComplex, + SupportsIndex, + complex, # `complex` is not a subtype of `SupportsComplex` +] + +class integer(number[_NBit1]): # type: ignore + @property + def numerator(self: _ScalarType) -> _ScalarType: ... + @property + def denominator(self) -> L[1]: ... + @overload + def __round__(self, ndigits: None = ...) -> int: ... + @overload + def __round__(self: _ScalarType, ndigits: SupportsIndex) -> _ScalarType: ... + + # NOTE: `__index__` is technically defined in the bottom-most + # sub-classes (`int64`, `uint32`, etc) + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> int: ... + def tolist(self) -> int: ... + def is_integer(self) -> L[True]: ... + def bit_count(self: _ScalarType) -> int: ... + def __index__(self) -> int: ... + __truediv__: _IntTrueDiv[_NBit1] + __rtruediv__: _IntTrueDiv[_NBit1] + def __mod__(self, value: _IntLike_co) -> integer[Any]: ... + def __rmod__(self, value: _IntLike_co) -> integer[Any]: ... + def __invert__(self: _IntType) -> _IntType: ... + # Ensure that objects annotated as `integer` support bit-wise operations + def __lshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rlshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __rrshift__(self, other: _IntLike_co) -> integer[Any]: ... + def __and__(self, other: _IntLike_co) -> integer[Any]: ... + def __rand__(self, other: _IntLike_co) -> integer[Any]: ... + def __or__(self, other: _IntLike_co) -> integer[Any]: ... + def __ror__(self, other: _IntLike_co) -> integer[Any]: ... + def __xor__(self, other: _IntLike_co) -> integer[Any]: ... + def __rxor__(self, other: _IntLike_co) -> integer[Any]: ... + +class signedinteger(integer[_NBit1]): + def __init__(self, value: _IntValue = ..., /) -> None: ... + __add__: _SignedIntOp[_NBit1] + __radd__: _SignedIntOp[_NBit1] + __sub__: _SignedIntOp[_NBit1] + __rsub__: _SignedIntOp[_NBit1] + __mul__: _SignedIntOp[_NBit1] + __rmul__: _SignedIntOp[_NBit1] + __floordiv__: _SignedIntOp[_NBit1] + __rfloordiv__: _SignedIntOp[_NBit1] + __pow__: _SignedIntOp[_NBit1] + __rpow__: _SignedIntOp[_NBit1] + __lshift__: _SignedIntBitOp[_NBit1] + __rlshift__: _SignedIntBitOp[_NBit1] + __rshift__: _SignedIntBitOp[_NBit1] + __rrshift__: _SignedIntBitOp[_NBit1] + __and__: _SignedIntBitOp[_NBit1] + __rand__: _SignedIntBitOp[_NBit1] + __xor__: _SignedIntBitOp[_NBit1] + __rxor__: _SignedIntBitOp[_NBit1] + __or__: _SignedIntBitOp[_NBit1] + __ror__: _SignedIntBitOp[_NBit1] + __mod__: _SignedIntMod[_NBit1] + __rmod__: _SignedIntMod[_NBit1] + __divmod__: _SignedIntDivMod[_NBit1] + __rdivmod__: _SignedIntDivMod[_NBit1] + +int8 = signedinteger[_8Bit] +int16 = signedinteger[_16Bit] +int32 = signedinteger[_32Bit] +int64 = signedinteger[_64Bit] + +byte = signedinteger[_NBitByte] +short = signedinteger[_NBitShort] +intc = signedinteger[_NBitIntC] +intp = signedinteger[_NBitIntP] +int_ = signedinteger[_NBitInt] +longlong = signedinteger[_NBitLongLong] + +# TODO: `item`/`tolist` returns either `dt.timedelta` or `int` +# depending on the unit +class timedelta64(generic): + def __init__( + self, + value: None | int | _CharLike_co | dt.timedelta | timedelta64 = ..., + format: _CharLike_co | tuple[_CharLike_co, _IntLike_co] = ..., + /, + ) -> None: ... + @property + def numerator(self: _ScalarType) -> _ScalarType: ... + @property + def denominator(self) -> L[1]: ... + + # NOTE: Only a limited number of units support conversion + # to builtin scalar types: `Y`, `M`, `ns`, `ps`, `fs`, `as` + def __int__(self) -> int: ... + def __float__(self) -> float: ... + def __complex__(self) -> complex: ... + def __neg__(self: _ArraySelf) -> _ArraySelf: ... + def __pos__(self: _ArraySelf) -> _ArraySelf: ... + def __abs__(self: _ArraySelf) -> _ArraySelf: ... + def __add__(self, other: _TD64Like_co) -> timedelta64: ... + def __radd__(self, other: _TD64Like_co) -> timedelta64: ... + def __sub__(self, other: _TD64Like_co) -> timedelta64: ... + def __rsub__(self, other: _TD64Like_co) -> timedelta64: ... + def __mul__(self, other: _FloatLike_co) -> timedelta64: ... + def __rmul__(self, other: _FloatLike_co) -> timedelta64: ... + __truediv__: _TD64Div[float64] + __floordiv__: _TD64Div[int64] + def __rtruediv__(self, other: timedelta64) -> float64: ... + def __rfloordiv__(self, other: timedelta64) -> int64: ... + def __mod__(self, other: timedelta64) -> timedelta64: ... + def __rmod__(self, other: timedelta64) -> timedelta64: ... + def __divmod__(self, other: timedelta64) -> tuple[int64, timedelta64]: ... + def __rdivmod__(self, other: timedelta64) -> tuple[int64, timedelta64]: ... + __lt__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __le__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __gt__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + __ge__: _ComparisonOp[_TD64Like_co, _ArrayLikeTD64_co] + +class unsignedinteger(integer[_NBit1]): + # NOTE: `uint64 + signedinteger -> float64` + def __init__(self, value: _IntValue = ..., /) -> None: ... + __add__: _UnsignedIntOp[_NBit1] + __radd__: _UnsignedIntOp[_NBit1] + __sub__: _UnsignedIntOp[_NBit1] + __rsub__: _UnsignedIntOp[_NBit1] + __mul__: _UnsignedIntOp[_NBit1] + __rmul__: _UnsignedIntOp[_NBit1] + __floordiv__: _UnsignedIntOp[_NBit1] + __rfloordiv__: _UnsignedIntOp[_NBit1] + __pow__: _UnsignedIntOp[_NBit1] + __rpow__: _UnsignedIntOp[_NBit1] + __lshift__: _UnsignedIntBitOp[_NBit1] + __rlshift__: _UnsignedIntBitOp[_NBit1] + __rshift__: _UnsignedIntBitOp[_NBit1] + __rrshift__: _UnsignedIntBitOp[_NBit1] + __and__: _UnsignedIntBitOp[_NBit1] + __rand__: _UnsignedIntBitOp[_NBit1] + __xor__: _UnsignedIntBitOp[_NBit1] + __rxor__: _UnsignedIntBitOp[_NBit1] + __or__: _UnsignedIntBitOp[_NBit1] + __ror__: _UnsignedIntBitOp[_NBit1] + __mod__: _UnsignedIntMod[_NBit1] + __rmod__: _UnsignedIntMod[_NBit1] + __divmod__: _UnsignedIntDivMod[_NBit1] + __rdivmod__: _UnsignedIntDivMod[_NBit1] + +uint8 = unsignedinteger[_8Bit] +uint16 = unsignedinteger[_16Bit] +uint32 = unsignedinteger[_32Bit] +uint64 = unsignedinteger[_64Bit] + +ubyte = unsignedinteger[_NBitByte] +ushort = unsignedinteger[_NBitShort] +uintc = unsignedinteger[_NBitIntC] +uintp = unsignedinteger[_NBitIntP] +uint = unsignedinteger[_NBitInt] +ulonglong = unsignedinteger[_NBitLongLong] + +class inexact(number[_NBit1]): # type: ignore + def __getnewargs__(self: inexact[_64Bit]) -> tuple[float, ...]: ... + +_IntType = TypeVar("_IntType", bound=integer[Any]) +_FloatType = TypeVar('_FloatType', bound=floating[Any]) + +class floating(inexact[_NBit1]): + def __init__(self, value: _FloatValue = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., + /, + ) -> float: ... + def tolist(self) -> float: ... + def is_integer(self) -> bool: ... + def hex(self: float64) -> str: ... + @classmethod + def fromhex(cls: type[float64], string: str, /) -> float64: ... + def as_integer_ratio(self) -> tuple[int, int]: ... + def __ceil__(self: float64) -> int: ... + def __floor__(self: float64) -> int: ... + def __trunc__(self: float64) -> int: ... + def __getnewargs__(self: float64) -> tuple[float]: ... + def __getformat__(self: float64, typestr: L["double", "float"], /) -> str: ... + @overload + def __round__(self, ndigits: None = ...) -> int: ... + @overload + def __round__(self: _ScalarType, ndigits: SupportsIndex) -> _ScalarType: ... + __add__: _FloatOp[_NBit1] + __radd__: _FloatOp[_NBit1] + __sub__: _FloatOp[_NBit1] + __rsub__: _FloatOp[_NBit1] + __mul__: _FloatOp[_NBit1] + __rmul__: _FloatOp[_NBit1] + __truediv__: _FloatOp[_NBit1] + __rtruediv__: _FloatOp[_NBit1] + __floordiv__: _FloatOp[_NBit1] + __rfloordiv__: _FloatOp[_NBit1] + __pow__: _FloatOp[_NBit1] + __rpow__: _FloatOp[_NBit1] + __mod__: _FloatMod[_NBit1] + __rmod__: _FloatMod[_NBit1] + __divmod__: _FloatDivMod[_NBit1] + __rdivmod__: _FloatDivMod[_NBit1] + +float16 = floating[_16Bit] +float32 = floating[_32Bit] +float64 = floating[_64Bit] + +half = floating[_NBitHalf] +single = floating[_NBitSingle] +double = floating[_NBitDouble] +float_ = floating[_NBitDouble] +longdouble = floating[_NBitLongDouble] +longfloat = floating[_NBitLongDouble] + +# The main reason for `complexfloating` having two typevars is cosmetic. +# It is used to clarify why `complex128`s precision is `_64Bit`, the latter +# describing the two 64 bit floats representing its real and imaginary component + +class complexfloating(inexact[_NBit1], Generic[_NBit1, _NBit2]): + def __init__(self, value: _ComplexValue = ..., /) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> complex: ... + def tolist(self) -> complex: ... + @property + def real(self) -> floating[_NBit1]: ... # type: ignore[override] + @property + def imag(self) -> floating[_NBit2]: ... # type: ignore[override] + def __abs__(self) -> floating[_NBit1]: ... # type: ignore[override] + def __getnewargs__(self: complex128) -> tuple[float, float]: ... + # NOTE: Deprecated + # def __round__(self, ndigits=...): ... + __add__: _ComplexOp[_NBit1] + __radd__: _ComplexOp[_NBit1] + __sub__: _ComplexOp[_NBit1] + __rsub__: _ComplexOp[_NBit1] + __mul__: _ComplexOp[_NBit1] + __rmul__: _ComplexOp[_NBit1] + __truediv__: _ComplexOp[_NBit1] + __rtruediv__: _ComplexOp[_NBit1] + __pow__: _ComplexOp[_NBit1] + __rpow__: _ComplexOp[_NBit1] + +complex64 = complexfloating[_32Bit, _32Bit] +complex128 = complexfloating[_64Bit, _64Bit] + +csingle = complexfloating[_NBitSingle, _NBitSingle] +singlecomplex = complexfloating[_NBitSingle, _NBitSingle] +cdouble = complexfloating[_NBitDouble, _NBitDouble] +complex_ = complexfloating[_NBitDouble, _NBitDouble] +cfloat = complexfloating[_NBitDouble, _NBitDouble] +clongdouble = complexfloating[_NBitLongDouble, _NBitLongDouble] +clongfloat = complexfloating[_NBitLongDouble, _NBitLongDouble] +longcomplex = complexfloating[_NBitLongDouble, _NBitLongDouble] + +class flexible(generic): ... # type: ignore + +# TODO: `item`/`tolist` returns either `bytes` or `tuple` +# depending on whether or not it's used as an opaque bytes sequence +# or a structure +class void(flexible): + @overload + def __init__(self, value: _IntLike_co | bytes, /, dtype : None = ...) -> None: ... + @overload + def __init__(self, value: Any, /, dtype: _DTypeLikeVoid) -> None: ... + @property + def real(self: _ArraySelf) -> _ArraySelf: ... + @property + def imag(self: _ArraySelf) -> _ArraySelf: ... + def setfield( + self, val: ArrayLike, dtype: DTypeLike, offset: int = ... + ) -> None: ... + @overload + def __getitem__(self, key: str | SupportsIndex) -> Any: ... + @overload + def __getitem__(self, key: list[str]) -> void: ... + def __setitem__( + self, + key: str | list[str] | SupportsIndex, + value: ArrayLike, + ) -> None: ... + +class character(flexible): # type: ignore + def __int__(self) -> int: ... + def __float__(self) -> float: ... + +# NOTE: Most `np.bytes_` / `np.str_` methods return their +# builtin `bytes` / `str` counterpart + +class bytes_(character, bytes): + @overload + def __init__(self, value: object = ..., /) -> None: ... + @overload + def __init__( + self, value: str, /, encoding: str = ..., errors: str = ... + ) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> bytes: ... + def tolist(self) -> bytes: ... + +string_ = bytes_ + +class str_(character, str): + @overload + def __init__(self, value: object = ..., /) -> None: ... + @overload + def __init__( + self, value: bytes, /, encoding: str = ..., errors: str = ... + ) -> None: ... + def item( + self, args: L[0] | tuple[()] | tuple[L[0]] = ..., /, + ) -> str: ... + def tolist(self) -> str: ... + +unicode_ = str_ + +# +# Constants +# + +Inf: Final[float] +Infinity: Final[float] +NAN: Final[float] +NINF: Final[float] +NZERO: Final[float] +NaN: Final[float] +PINF: Final[float] +PZERO: Final[float] +e: Final[float] +euler_gamma: Final[float] +inf: Final[float] +infty: Final[float] +nan: Final[float] +pi: Final[float] + +ERR_IGNORE: L[0] +ERR_WARN: L[1] +ERR_RAISE: L[2] +ERR_CALL: L[3] +ERR_PRINT: L[4] +ERR_LOG: L[5] +ERR_DEFAULT: L[521] + +SHIFT_DIVIDEBYZERO: L[0] +SHIFT_OVERFLOW: L[3] +SHIFT_UNDERFLOW: L[6] +SHIFT_INVALID: L[9] + +FPE_DIVIDEBYZERO: L[1] +FPE_OVERFLOW: L[2] +FPE_UNDERFLOW: L[4] +FPE_INVALID: L[8] + +FLOATING_POINT_SUPPORT: L[1] +UFUNC_BUFSIZE_DEFAULT = BUFSIZE + +little_endian: Final[bool] +True_: Final[bool_] +False_: Final[bool_] + +UFUNC_PYVALS_NAME: L["UFUNC_PYVALS"] + +newaxis: None + +# See `numpy._typing._ufunc` for more concrete nin-/nout-specific stubs +@final +class ufunc: + @property + def __name__(self) -> str: ... + @property + def __doc__(self) -> str: ... + __call__: Callable[..., Any] + @property + def nin(self) -> int: ... + @property + def nout(self) -> int: ... + @property + def nargs(self) -> int: ... + @property + def ntypes(self) -> int: ... + @property + def types(self) -> list[str]: ... + # Broad return type because it has to encompass things like + # + # >>> np.logical_and.identity is True + # True + # >>> np.add.identity is 0 + # True + # >>> np.sin.identity is None + # True + # + # and any user-defined ufuncs. + @property + def identity(self) -> Any: ... + # This is None for ufuncs and a string for gufuncs. + @property + def signature(self) -> None | str: ... + # The next four methods will always exist, but they will just + # raise a ValueError ufuncs with that don't accept two input + # arguments and return one output argument. Because of that we + # can't type them very precisely. + reduce: Any + accumulate: Any + reduceat: Any + outer: Any + # Similarly at won't be defined for ufuncs that return multiple + # outputs, so we can't type it very precisely. + at: Any + +# Parameters: `__name__`, `ntypes` and `identity` +absolute: _UFunc_Nin1_Nout1[L['absolute'], L[20], None] +add: _UFunc_Nin2_Nout1[L['add'], L[22], L[0]] +arccos: _UFunc_Nin1_Nout1[L['arccos'], L[8], None] +arccosh: _UFunc_Nin1_Nout1[L['arccosh'], L[8], None] +arcsin: _UFunc_Nin1_Nout1[L['arcsin'], L[8], None] +arcsinh: _UFunc_Nin1_Nout1[L['arcsinh'], L[8], None] +arctan2: _UFunc_Nin2_Nout1[L['arctan2'], L[5], None] +arctan: _UFunc_Nin1_Nout1[L['arctan'], L[8], None] +arctanh: _UFunc_Nin1_Nout1[L['arctanh'], L[8], None] +bitwise_and: _UFunc_Nin2_Nout1[L['bitwise_and'], L[12], L[-1]] +bitwise_not: _UFunc_Nin1_Nout1[L['invert'], L[12], None] +bitwise_or: _UFunc_Nin2_Nout1[L['bitwise_or'], L[12], L[0]] +bitwise_xor: _UFunc_Nin2_Nout1[L['bitwise_xor'], L[12], L[0]] +cbrt: _UFunc_Nin1_Nout1[L['cbrt'], L[5], None] +ceil: _UFunc_Nin1_Nout1[L['ceil'], L[7], None] +conj: _UFunc_Nin1_Nout1[L['conjugate'], L[18], None] +conjugate: _UFunc_Nin1_Nout1[L['conjugate'], L[18], None] +copysign: _UFunc_Nin2_Nout1[L['copysign'], L[4], None] +cos: _UFunc_Nin1_Nout1[L['cos'], L[9], None] +cosh: _UFunc_Nin1_Nout1[L['cosh'], L[8], None] +deg2rad: _UFunc_Nin1_Nout1[L['deg2rad'], L[5], None] +degrees: _UFunc_Nin1_Nout1[L['degrees'], L[5], None] +divide: _UFunc_Nin2_Nout1[L['true_divide'], L[11], None] +divmod: _UFunc_Nin2_Nout2[L['divmod'], L[15], None] +equal: _UFunc_Nin2_Nout1[L['equal'], L[23], None] +exp2: _UFunc_Nin1_Nout1[L['exp2'], L[8], None] +exp: _UFunc_Nin1_Nout1[L['exp'], L[10], None] +expm1: _UFunc_Nin1_Nout1[L['expm1'], L[8], None] +fabs: _UFunc_Nin1_Nout1[L['fabs'], L[5], None] +float_power: _UFunc_Nin2_Nout1[L['float_power'], L[4], None] +floor: _UFunc_Nin1_Nout1[L['floor'], L[7], None] +floor_divide: _UFunc_Nin2_Nout1[L['floor_divide'], L[21], None] +fmax: _UFunc_Nin2_Nout1[L['fmax'], L[21], None] +fmin: _UFunc_Nin2_Nout1[L['fmin'], L[21], None] +fmod: _UFunc_Nin2_Nout1[L['fmod'], L[15], None] +frexp: _UFunc_Nin1_Nout2[L['frexp'], L[4], None] +gcd: _UFunc_Nin2_Nout1[L['gcd'], L[11], L[0]] +greater: _UFunc_Nin2_Nout1[L['greater'], L[23], None] +greater_equal: _UFunc_Nin2_Nout1[L['greater_equal'], L[23], None] +heaviside: _UFunc_Nin2_Nout1[L['heaviside'], L[4], None] +hypot: _UFunc_Nin2_Nout1[L['hypot'], L[5], L[0]] +invert: _UFunc_Nin1_Nout1[L['invert'], L[12], None] +isfinite: _UFunc_Nin1_Nout1[L['isfinite'], L[20], None] +isinf: _UFunc_Nin1_Nout1[L['isinf'], L[20], None] +isnan: _UFunc_Nin1_Nout1[L['isnan'], L[20], None] +isnat: _UFunc_Nin1_Nout1[L['isnat'], L[2], None] +lcm: _UFunc_Nin2_Nout1[L['lcm'], L[11], None] +ldexp: _UFunc_Nin2_Nout1[L['ldexp'], L[8], None] +left_shift: _UFunc_Nin2_Nout1[L['left_shift'], L[11], None] +less: _UFunc_Nin2_Nout1[L['less'], L[23], None] +less_equal: _UFunc_Nin2_Nout1[L['less_equal'], L[23], None] +log10: _UFunc_Nin1_Nout1[L['log10'], L[8], None] +log1p: _UFunc_Nin1_Nout1[L['log1p'], L[8], None] +log2: _UFunc_Nin1_Nout1[L['log2'], L[8], None] +log: _UFunc_Nin1_Nout1[L['log'], L[10], None] +logaddexp2: _UFunc_Nin2_Nout1[L['logaddexp2'], L[4], float] +logaddexp: _UFunc_Nin2_Nout1[L['logaddexp'], L[4], float] +logical_and: _UFunc_Nin2_Nout1[L['logical_and'], L[20], L[True]] +logical_not: _UFunc_Nin1_Nout1[L['logical_not'], L[20], None] +logical_or: _UFunc_Nin2_Nout1[L['logical_or'], L[20], L[False]] +logical_xor: _UFunc_Nin2_Nout1[L['logical_xor'], L[19], L[False]] +matmul: _GUFunc_Nin2_Nout1[L['matmul'], L[19], None] +maximum: _UFunc_Nin2_Nout1[L['maximum'], L[21], None] +minimum: _UFunc_Nin2_Nout1[L['minimum'], L[21], None] +mod: _UFunc_Nin2_Nout1[L['remainder'], L[16], None] +modf: _UFunc_Nin1_Nout2[L['modf'], L[4], None] +multiply: _UFunc_Nin2_Nout1[L['multiply'], L[23], L[1]] +negative: _UFunc_Nin1_Nout1[L['negative'], L[19], None] +nextafter: _UFunc_Nin2_Nout1[L['nextafter'], L[4], None] +not_equal: _UFunc_Nin2_Nout1[L['not_equal'], L[23], None] +positive: _UFunc_Nin1_Nout1[L['positive'], L[19], None] +power: _UFunc_Nin2_Nout1[L['power'], L[18], None] +rad2deg: _UFunc_Nin1_Nout1[L['rad2deg'], L[5], None] +radians: _UFunc_Nin1_Nout1[L['radians'], L[5], None] +reciprocal: _UFunc_Nin1_Nout1[L['reciprocal'], L[18], None] +remainder: _UFunc_Nin2_Nout1[L['remainder'], L[16], None] +right_shift: _UFunc_Nin2_Nout1[L['right_shift'], L[11], None] +rint: _UFunc_Nin1_Nout1[L['rint'], L[10], None] +sign: _UFunc_Nin1_Nout1[L['sign'], L[19], None] +signbit: _UFunc_Nin1_Nout1[L['signbit'], L[4], None] +sin: _UFunc_Nin1_Nout1[L['sin'], L[9], None] +sinh: _UFunc_Nin1_Nout1[L['sinh'], L[8], None] +spacing: _UFunc_Nin1_Nout1[L['spacing'], L[4], None] +sqrt: _UFunc_Nin1_Nout1[L['sqrt'], L[10], None] +square: _UFunc_Nin1_Nout1[L['square'], L[18], None] +subtract: _UFunc_Nin2_Nout1[L['subtract'], L[21], None] +tan: _UFunc_Nin1_Nout1[L['tan'], L[8], None] +tanh: _UFunc_Nin1_Nout1[L['tanh'], L[8], None] +true_divide: _UFunc_Nin2_Nout1[L['true_divide'], L[11], None] +trunc: _UFunc_Nin1_Nout1[L['trunc'], L[7], None] + +abs = absolute + +class _CopyMode(enum.Enum): + ALWAYS: L[True] + IF_NEEDED: L[False] + NEVER: L[2] + +# Warnings +class RankWarning(UserWarning): ... + +_CallType = TypeVar("_CallType", bound=_ErrFunc | _SupportsWrite[str]) + +class errstate(Generic[_CallType], ContextDecorator): + call: _CallType + kwargs: _ErrDictOptional + + # Expand `**kwargs` into explicit keyword-only arguments + def __init__( + self, + *, + call: _CallType = ..., + all: None | _ErrKind = ..., + divide: None | _ErrKind = ..., + over: None | _ErrKind = ..., + under: None | _ErrKind = ..., + invalid: None | _ErrKind = ..., + ) -> None: ... + def __enter__(self) -> None: ... + def __exit__( + self, + exc_type: None | type[BaseException], + exc_value: None | BaseException, + traceback: None | TracebackType, + /, + ) -> None: ... + +@contextmanager +def _no_nep50_warning() -> Generator[None, None, None]: ... +def _get_promotion_state() -> str: ... +def _set_promotion_state(state: str, /) -> None: ... + +class ndenumerate(Generic[_ScalarType]): + iter: flatiter[NDArray[_ScalarType]] + @overload + def __new__( + cls, arr: _FiniteNestedSequence[_SupportsArray[dtype[_ScalarType]]], + ) -> ndenumerate[_ScalarType]: ... + @overload + def __new__(cls, arr: str | _NestedSequence[str]) -> ndenumerate[str_]: ... + @overload + def __new__(cls, arr: bytes | _NestedSequence[bytes]) -> ndenumerate[bytes_]: ... + @overload + def __new__(cls, arr: bool | _NestedSequence[bool]) -> ndenumerate[bool_]: ... + @overload + def __new__(cls, arr: int | _NestedSequence[int]) -> ndenumerate[int_]: ... + @overload + def __new__(cls, arr: float | _NestedSequence[float]) -> ndenumerate[float_]: ... + @overload + def __new__(cls, arr: complex | _NestedSequence[complex]) -> ndenumerate[complex_]: ... + def __next__(self: ndenumerate[_ScalarType]) -> tuple[_Shape, _ScalarType]: ... + def __iter__(self: _T) -> _T: ... + +class ndindex: + @overload + def __init__(self, shape: tuple[SupportsIndex, ...], /) -> None: ... + @overload + def __init__(self, *shape: SupportsIndex) -> None: ... + def __iter__(self: _T) -> _T: ... + def __next__(self) -> _Shape: ... + +class DataSource: + def __init__( + self, + destpath: None | str | os.PathLike[str] = ..., + ) -> None: ... + def __del__(self) -> None: ... + def abspath(self, path: str) -> str: ... + def exists(self, path: str) -> bool: ... + + # Whether the file-object is opened in string or bytes mode (by default) + # depends on the file-extension of `path` + def open( + self, + path: str, + mode: str = ..., + encoding: None | str = ..., + newline: None | str = ..., + ) -> IO[Any]: ... + +# TODO: The type of each `__next__` and `iters` return-type depends +# on the length and dtype of `args`; we can't describe this behavior yet +# as we lack variadics (PEP 646). +@final +class broadcast: + def __new__(cls, *args: ArrayLike) -> broadcast: ... + @property + def index(self) -> int: ... + @property + def iters(self) -> tuple[flatiter[Any], ...]: ... + @property + def nd(self) -> int: ... + @property + def ndim(self) -> int: ... + @property + def numiter(self) -> int: ... + @property + def shape(self) -> _Shape: ... + @property + def size(self) -> int: ... + def __next__(self) -> tuple[Any, ...]: ... + def __iter__(self: _T) -> _T: ... + def reset(self) -> None: ... + +@final +class busdaycalendar: + def __new__( + cls, + weekmask: ArrayLike = ..., + holidays: ArrayLike | dt.date | _NestedSequence[dt.date] = ..., + ) -> busdaycalendar: ... + @property + def weekmask(self) -> NDArray[bool_]: ... + @property + def holidays(self) -> NDArray[datetime64]: ... + +class finfo(Generic[_FloatType]): + dtype: dtype[_FloatType] + bits: int + eps: _FloatType + epsneg: _FloatType + iexp: int + machep: int + max: _FloatType + maxexp: int + min: _FloatType + minexp: int + negep: int + nexp: int + nmant: int + precision: int + resolution: _FloatType + smallest_subnormal: _FloatType + @property + def smallest_normal(self) -> _FloatType: ... + @property + def tiny(self) -> _FloatType: ... + @overload + def __new__( + cls, dtype: inexact[_NBit1] | _DTypeLike[inexact[_NBit1]] + ) -> finfo[floating[_NBit1]]: ... + @overload + def __new__( + cls, dtype: complex | float | type[complex] | type[float] + ) -> finfo[float_]: ... + @overload + def __new__( + cls, dtype: str + ) -> finfo[floating[Any]]: ... + +class iinfo(Generic[_IntType]): + dtype: dtype[_IntType] + kind: str + bits: int + key: str + @property + def min(self) -> int: ... + @property + def max(self) -> int: ... + + @overload + def __new__(cls, dtype: _IntType | _DTypeLike[_IntType]) -> iinfo[_IntType]: ... + @overload + def __new__(cls, dtype: int | type[int]) -> iinfo[int_]: ... + @overload + def __new__(cls, dtype: str) -> iinfo[Any]: ... + +class format_parser: + dtype: dtype[void] + def __init__( + self, + formats: DTypeLike, + names: None | str | Sequence[str], + titles: None | str | Sequence[str], + aligned: bool = ..., + byteorder: None | _ByteOrder = ..., + ) -> None: ... + +class recarray(ndarray[_ShapeType, _DType_co]): + # NOTE: While not strictly mandatory, we're demanding here that arguments + # for the `format_parser`- and `dtype`-based dtype constructors are + # mutually exclusive + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: None = ..., + buf: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + *, + formats: DTypeLike, + names: None | str | Sequence[str] = ..., + titles: None | str | Sequence[str] = ..., + byteorder: None | _ByteOrder = ..., + aligned: bool = ..., + order: _OrderKACF = ..., + ) -> recarray[Any, dtype[record]]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + dtype: DTypeLike, + buf: None | _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: None | _ShapeLike = ..., + formats: None = ..., + names: None = ..., + titles: None = ..., + byteorder: None = ..., + aligned: L[False] = ..., + order: _OrderKACF = ..., + ) -> recarray[Any, dtype[Any]]: ... + def __array_finalize__(self, obj: object) -> None: ... + def __getattribute__(self, attr: str) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike) -> None: ... + @overload + def __getitem__(self, indx: ( + SupportsIndex + | _ArrayLikeInt_co + | tuple[SupportsIndex | _ArrayLikeInt_co, ...] + )) -> Any: ... + @overload + def __getitem__(self: recarray[Any, dtype[void]], indx: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> recarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> ndarray[Any, _DType_co]: ... + @overload + def __getitem__(self, indx: str) -> NDArray[Any]: ... + @overload + def __getitem__(self, indx: list[str]) -> recarray[_ShapeType, dtype[record]]: ... + @overload + def field(self, attr: int | str, val: None = ...) -> Any: ... + @overload + def field(self, attr: int | str, val: ArrayLike) -> None: ... + +class record(void): + def __getattribute__(self, attr: str) -> Any: ... + def __setattr__(self, attr: str, val: ArrayLike) -> None: ... + def pprint(self) -> str: ... + @overload + def __getitem__(self, key: str | SupportsIndex) -> Any: ... + @overload + def __getitem__(self, key: list[str]) -> record: ... + +_NDIterFlagsKind = L[ + "buffered", + "c_index", + "copy_if_overlap", + "common_dtype", + "delay_bufalloc", + "external_loop", + "f_index", + "grow_inner", "growinner", + "multi_index", + "ranged", + "refs_ok", + "reduce_ok", + "zerosize_ok", +] + +_NDIterOpFlagsKind = L[ + "aligned", + "allocate", + "arraymask", + "copy", + "config", + "nbo", + "no_subtype", + "no_broadcast", + "overlap_assume_elementwise", + "readonly", + "readwrite", + "updateifcopy", + "virtual", + "writeonly", + "writemasked" +] + +@final +class nditer: + def __new__( + cls, + op: ArrayLike | Sequence[ArrayLike], + flags: None | Sequence[_NDIterFlagsKind] = ..., + op_flags: None | Sequence[Sequence[_NDIterOpFlagsKind]] = ..., + op_dtypes: DTypeLike | Sequence[DTypeLike] = ..., + order: _OrderKACF = ..., + casting: _CastingKind = ..., + op_axes: None | Sequence[Sequence[SupportsIndex]] = ..., + itershape: None | _ShapeLike = ..., + buffersize: SupportsIndex = ..., + ) -> nditer: ... + def __enter__(self) -> nditer: ... + def __exit__( + self, + exc_type: None | type[BaseException], + exc_value: None | BaseException, + traceback: None | TracebackType, + ) -> None: ... + def __iter__(self) -> nditer: ... + def __next__(self) -> tuple[NDArray[Any], ...]: ... + def __len__(self) -> int: ... + def __copy__(self) -> nditer: ... + @overload + def __getitem__(self, index: SupportsIndex) -> NDArray[Any]: ... + @overload + def __getitem__(self, index: slice) -> tuple[NDArray[Any], ...]: ... + def __setitem__(self, index: slice | SupportsIndex, value: ArrayLike) -> None: ... + def close(self) -> None: ... + def copy(self) -> nditer: ... + def debug_print(self) -> None: ... + def enable_external_loop(self) -> None: ... + def iternext(self) -> bool: ... + def remove_axis(self, i: SupportsIndex, /) -> None: ... + def remove_multi_index(self) -> None: ... + def reset(self) -> None: ... + @property + def dtypes(self) -> tuple[dtype[Any], ...]: ... + @property + def finished(self) -> bool: ... + @property + def has_delayed_bufalloc(self) -> bool: ... + @property + def has_index(self) -> bool: ... + @property + def has_multi_index(self) -> bool: ... + @property + def index(self) -> int: ... + @property + def iterationneedsapi(self) -> bool: ... + @property + def iterindex(self) -> int: ... + @property + def iterrange(self) -> tuple[int, ...]: ... + @property + def itersize(self) -> int: ... + @property + def itviews(self) -> tuple[NDArray[Any], ...]: ... + @property + def multi_index(self) -> tuple[int, ...]: ... + @property + def ndim(self) -> int: ... + @property + def nop(self) -> int: ... + @property + def operands(self) -> tuple[NDArray[Any], ...]: ... + @property + def shape(self) -> tuple[int, ...]: ... + @property + def value(self) -> tuple[NDArray[Any], ...]: ... + +_MemMapModeKind = L[ + "readonly", "r", + "copyonwrite", "c", + "readwrite", "r+", + "write", "w+", +] + +class memmap(ndarray[_ShapeType, _DType_co]): + __array_priority__: ClassVar[float] + filename: str | None + offset: int + mode: str + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: type[uint8] = ..., + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[uint8]]: ... + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: _DTypeLike[_ScalarType], + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[_ScalarType]]: ... + @overload + def __new__( + subtype, + filename: str | bytes | os.PathLike[str] | os.PathLike[bytes] | _MemMapIOProtocol, + dtype: DTypeLike, + mode: _MemMapModeKind = ..., + offset: int = ..., + shape: None | int | tuple[int, ...] = ..., + order: _OrderKACF = ..., + ) -> memmap[Any, dtype[Any]]: ... + def __array_finalize__(self, obj: object) -> None: ... + def __array_wrap__( + self, + array: memmap[_ShapeType, _DType_co], + context: None | tuple[ufunc, tuple[Any, ...], int] = ..., + ) -> Any: ... + def flush(self) -> None: ... + +# TODO: Add a mypy plugin for managing functions whose output type is dependent +# on the literal value of some sort of signature (e.g. `einsum` and `vectorize`) +class vectorize: + pyfunc: Callable[..., Any] + cache: bool + signature: None | str + otypes: None | str + excluded: set[int | str] + __doc__: None | str + def __init__( + self, + pyfunc: Callable[..., Any], + otypes: None | str | Iterable[DTypeLike] = ..., + doc: None | str = ..., + excluded: None | Iterable[int | str] = ..., + cache: bool = ..., + signature: None | str = ..., + ) -> None: ... + def __call__(self, *args: Any, **kwargs: Any) -> Any: ... + +class poly1d: + @property + def variable(self) -> str: ... + @property + def order(self) -> int: ... + @property + def o(self) -> int: ... + @property + def roots(self) -> NDArray[Any]: ... + @property + def r(self) -> NDArray[Any]: ... + + @property + def coeffs(self) -> NDArray[Any]: ... + @coeffs.setter + def coeffs(self, value: NDArray[Any]) -> None: ... + + @property + def c(self) -> NDArray[Any]: ... + @c.setter + def c(self, value: NDArray[Any]) -> None: ... + + @property + def coef(self) -> NDArray[Any]: ... + @coef.setter + def coef(self, value: NDArray[Any]) -> None: ... + + @property + def coefficients(self) -> NDArray[Any]: ... + @coefficients.setter + def coefficients(self, value: NDArray[Any]) -> None: ... + + __hash__: ClassVar[None] # type: ignore + + @overload + def __array__(self, t: None = ...) -> NDArray[Any]: ... + @overload + def __array__(self, t: _DType) -> ndarray[Any, _DType]: ... + + @overload + def __call__(self, val: _ScalarLike_co) -> Any: ... + @overload + def __call__(self, val: poly1d) -> poly1d: ... + @overload + def __call__(self, val: ArrayLike) -> NDArray[Any]: ... + + def __init__( + self, + c_or_r: ArrayLike, + r: bool = ..., + variable: None | str = ..., + ) -> None: ... + def __len__(self) -> int: ... + def __neg__(self) -> poly1d: ... + def __pos__(self) -> poly1d: ... + def __mul__(self, other: ArrayLike) -> poly1d: ... + def __rmul__(self, other: ArrayLike) -> poly1d: ... + def __add__(self, other: ArrayLike) -> poly1d: ... + def __radd__(self, other: ArrayLike) -> poly1d: ... + def __pow__(self, val: _FloatLike_co) -> poly1d: ... # Integral floats are accepted + def __sub__(self, other: ArrayLike) -> poly1d: ... + def __rsub__(self, other: ArrayLike) -> poly1d: ... + def __div__(self, other: ArrayLike) -> poly1d: ... + def __truediv__(self, other: ArrayLike) -> poly1d: ... + def __rdiv__(self, other: ArrayLike) -> poly1d: ... + def __rtruediv__(self, other: ArrayLike) -> poly1d: ... + def __getitem__(self, val: int) -> Any: ... + def __setitem__(self, key: int, val: Any) -> None: ... + def __iter__(self) -> Iterator[Any]: ... + def deriv(self, m: SupportsInt | SupportsIndex = ...) -> poly1d: ... + def integ( + self, + m: SupportsInt | SupportsIndex = ..., + k: None | _ArrayLikeComplex_co | _ArrayLikeObject_co = ..., + ) -> poly1d: ... + +class matrix(ndarray[_ShapeType, _DType_co]): + __array_priority__: ClassVar[float] + def __new__( + subtype, + data: ArrayLike, + dtype: DTypeLike = ..., + copy: bool = ..., + ) -> matrix[Any, Any]: ... + def __array_finalize__(self, obj: object) -> None: ... + + @overload + def __getitem__(self, key: ( + SupportsIndex + | _ArrayLikeInt_co + | tuple[SupportsIndex | _ArrayLikeInt_co, ...] + )) -> Any: ... + @overload + def __getitem__(self, key: ( + None + | slice + | ellipsis + | SupportsIndex + | _ArrayLikeInt_co + | tuple[None | slice | ellipsis | _ArrayLikeInt_co | SupportsIndex, ...] + )) -> matrix[Any, _DType_co]: ... + @overload + def __getitem__(self: NDArray[void], key: str) -> matrix[Any, dtype[Any]]: ... + @overload + def __getitem__(self: NDArray[void], key: list[str]) -> matrix[_ShapeType, dtype[void]]: ... + + def __mul__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __rmul__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __imul__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ... + def __pow__(self, other: ArrayLike) -> matrix[Any, Any]: ... + def __ipow__(self, other: ArrayLike) -> matrix[_ShapeType, _DType_co]: ... + + @overload + def sum(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def sum(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def sum(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def mean(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def mean(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def mean(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def std(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ... + @overload + def std(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ... + @overload + def std(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ... + + @overload + def var(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> Any: ... + @overload + def var(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ..., ddof: float = ...) -> matrix[Any, Any]: ... + @overload + def var(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ..., ddof: float = ...) -> _NdArraySubClass: ... + + @overload + def prod(self, axis: None = ..., dtype: DTypeLike = ..., out: None = ...) -> Any: ... + @overload + def prod(self, axis: _ShapeLike, dtype: DTypeLike = ..., out: None = ...) -> matrix[Any, Any]: ... + @overload + def prod(self, axis: None | _ShapeLike = ..., dtype: DTypeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def any(self, axis: None = ..., out: None = ...) -> bool_: ... + @overload + def any(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ... + @overload + def any(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def all(self, axis: None = ..., out: None = ...) -> bool_: ... + @overload + def all(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[bool_]]: ... + @overload + def all(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def max(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def max(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def max(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def min(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def min(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def min(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def argmax(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ... + @overload + def argmax(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ... + @overload + def argmax(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def argmin(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> intp: ... + @overload + def argmin(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, dtype[intp]]: ... + @overload + def argmin(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + @overload + def ptp(self: NDArray[_ScalarType], axis: None = ..., out: None = ...) -> _ScalarType: ... + @overload + def ptp(self, axis: _ShapeLike, out: None = ...) -> matrix[Any, _DType_co]: ... + @overload + def ptp(self, axis: None | _ShapeLike = ..., out: _NdArraySubClass = ...) -> _NdArraySubClass: ... + + def squeeze(self, axis: None | _ShapeLike = ...) -> matrix[Any, _DType_co]: ... + def tolist(self: matrix[Any, dtype[_SupportsItem[_T]]]) -> list[list[_T]]: ... # type: ignore[typevar] + def ravel(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ... + def flatten(self, order: _OrderKACF = ...) -> matrix[Any, _DType_co]: ... + + @property + def T(self) -> matrix[Any, _DType_co]: ... + @property + def I(self) -> matrix[Any, Any]: ... + @property + def A(self) -> ndarray[_ShapeType, _DType_co]: ... + @property + def A1(self) -> ndarray[Any, _DType_co]: ... + @property + def H(self) -> matrix[Any, _DType_co]: ... + def getT(self) -> matrix[Any, _DType_co]: ... + def getI(self) -> matrix[Any, Any]: ... + def getA(self) -> ndarray[_ShapeType, _DType_co]: ... + def getA1(self) -> ndarray[Any, _DType_co]: ... + def getH(self) -> matrix[Any, _DType_co]: ... + +_CharType = TypeVar("_CharType", str_, bytes_) +_CharDType = TypeVar("_CharDType", dtype[str_], dtype[bytes_]) +_CharArray = chararray[Any, dtype[_CharType]] + +class chararray(ndarray[_ShapeType, _CharDType]): + @overload + def __new__( + subtype, + shape: _ShapeLike, + itemsize: SupportsIndex | SupportsInt = ..., + unicode: L[False] = ..., + buffer: _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> chararray[Any, dtype[bytes_]]: ... + @overload + def __new__( + subtype, + shape: _ShapeLike, + itemsize: SupportsIndex | SupportsInt = ..., + unicode: L[True] = ..., + buffer: _SupportsBuffer = ..., + offset: SupportsIndex = ..., + strides: _ShapeLike = ..., + order: _OrderKACF = ..., + ) -> chararray[Any, dtype[str_]]: ... + + def __array_finalize__(self, obj: object) -> None: ... + def __mul__(self, other: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def __rmul__(self, other: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def __mod__(self, i: Any) -> chararray[Any, _CharDType]: ... + + @overload + def __eq__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __eq__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __ne__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __ne__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __ge__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __ge__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __le__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __le__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __gt__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __gt__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __lt__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> NDArray[bool_]: ... + @overload + def __lt__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> NDArray[bool_]: ... + + @overload + def __add__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def __add__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def __radd__( + self: _CharArray[str_], + other: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def __radd__( + self: _CharArray[bytes_], + other: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def center( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def center( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def count( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def count( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + def decode( + self: _CharArray[bytes_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[str_]: ... + + def encode( + self: _CharArray[str_], + encoding: None | str = ..., + errors: None | str = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def endswith( + self: _CharArray[str_], + suffix: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + @overload + def endswith( + self: _CharArray[bytes_], + suffix: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + + def expandtabs( + self, + tabsize: _ArrayLikeInt_co = ..., + ) -> chararray[Any, _CharDType]: ... + + @overload + def find( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def find( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def index( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def index( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def join( + self: _CharArray[str_], + seq: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def join( + self: _CharArray[bytes_], + seq: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def ljust( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def ljust( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def lstrip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def lstrip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def partition( + self: _CharArray[str_], + sep: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def partition( + self: _CharArray[bytes_], + sep: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def replace( + self: _CharArray[str_], + old: _ArrayLikeStr_co, + new: _ArrayLikeStr_co, + count: None | _ArrayLikeInt_co = ..., + ) -> _CharArray[str_]: ... + @overload + def replace( + self: _CharArray[bytes_], + old: _ArrayLikeBytes_co, + new: _ArrayLikeBytes_co, + count: None | _ArrayLikeInt_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rfind( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def rfind( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rindex( + self: _CharArray[str_], + sub: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + @overload + def rindex( + self: _CharArray[bytes_], + sub: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[int_]: ... + + @overload + def rjust( + self: _CharArray[str_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rjust( + self: _CharArray[bytes_], + width: _ArrayLikeInt_co, + fillchar: _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def rpartition( + self: _CharArray[str_], + sep: _ArrayLikeStr_co, + ) -> _CharArray[str_]: ... + @overload + def rpartition( + self: _CharArray[bytes_], + sep: _ArrayLikeBytes_co, + ) -> _CharArray[bytes_]: ... + + @overload + def rsplit( + self: _CharArray[str_], + sep: None | _ArrayLikeStr_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + @overload + def rsplit( + self: _CharArray[bytes_], + sep: None | _ArrayLikeBytes_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + + @overload + def rstrip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def rstrip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def split( + self: _CharArray[str_], + sep: None | _ArrayLikeStr_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + @overload + def split( + self: _CharArray[bytes_], + sep: None | _ArrayLikeBytes_co = ..., + maxsplit: None | _ArrayLikeInt_co = ..., + ) -> NDArray[object_]: ... + + def splitlines(self, keepends: None | _ArrayLikeBool_co = ...) -> NDArray[object_]: ... + + @overload + def startswith( + self: _CharArray[str_], + prefix: _ArrayLikeStr_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + @overload + def startswith( + self: _CharArray[bytes_], + prefix: _ArrayLikeBytes_co, + start: _ArrayLikeInt_co = ..., + end: None | _ArrayLikeInt_co = ..., + ) -> NDArray[bool_]: ... + + @overload + def strip( + self: _CharArray[str_], + chars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def strip( + self: _CharArray[bytes_], + chars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + @overload + def translate( + self: _CharArray[str_], + table: _ArrayLikeStr_co, + deletechars: None | _ArrayLikeStr_co = ..., + ) -> _CharArray[str_]: ... + @overload + def translate( + self: _CharArray[bytes_], + table: _ArrayLikeBytes_co, + deletechars: None | _ArrayLikeBytes_co = ..., + ) -> _CharArray[bytes_]: ... + + def zfill(self, width: _ArrayLikeInt_co) -> chararray[Any, _CharDType]: ... + def capitalize(self) -> chararray[_ShapeType, _CharDType]: ... + def title(self) -> chararray[_ShapeType, _CharDType]: ... + def swapcase(self) -> chararray[_ShapeType, _CharDType]: ... + def lower(self) -> chararray[_ShapeType, _CharDType]: ... + def upper(self) -> chararray[_ShapeType, _CharDType]: ... + def isalnum(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isalpha(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isdigit(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def islower(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isspace(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def istitle(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isupper(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isnumeric(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + def isdecimal(self) -> ndarray[_ShapeType, dtype[bool_]]: ... + +# NOTE: Deprecated +# class MachAr: ... + +class _SupportsDLPack(Protocol[_T_contra]): + def __dlpack__(self, *, stream: None | _T_contra = ...) -> _PyCapsule: ... + +def from_dlpack(obj: _SupportsDLPack[None], /) -> NDArray[Any]: ... diff --git a/lib/python3.12/site-packages/numpy/_distributor_init.py b/lib/python3.12/site-packages/numpy/_distributor_init.py new file mode 100644 index 0000000000000000000000000000000000000000..25b0eed79fcabe6d6ad5a7b2bf45e5371f37d4a0 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/_distributor_init.py @@ -0,0 +1,15 @@ +""" Distributor init file + +Distributors: you can add custom code here to support particular distributions +of numpy. + +For example, this is a good place to put any BLAS/LAPACK initialization code. + +The numpy standard source distribution will not put code in this file, so you +can safely replace this file with your own version. +""" + +try: + from . import _distributor_init_local +except ImportError: + pass diff --git a/lib/python3.12/site-packages/numpy/_globals.py b/lib/python3.12/site-packages/numpy/_globals.py new file mode 100644 index 0000000000000000000000000000000000000000..416a20f5e11b14b1da34e2bfb45c7961edc9097c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/_globals.py @@ -0,0 +1,95 @@ +""" +Module defining global singleton classes. + +This module raises a RuntimeError if an attempt to reload it is made. In that +way the identities of the classes defined here are fixed and will remain so +even if numpy itself is reloaded. In particular, a function like the following +will still work correctly after numpy is reloaded:: + + def foo(arg=np._NoValue): + if arg is np._NoValue: + ... + +That was not the case when the singleton classes were defined in the numpy +``__init__.py`` file. See gh-7844 for a discussion of the reload problem that +motivated this module. + +""" +import enum + +from ._utils import set_module as _set_module + +__all__ = ['_NoValue', '_CopyMode'] + + +# Disallow reloading this module so as to preserve the identities of the +# classes defined here. +if '_is_loaded' in globals(): + raise RuntimeError('Reloading numpy._globals is not allowed') +_is_loaded = True + + +class _NoValueType: + """Special keyword value. + + The instance of this class may be used as the default value assigned to a + keyword if no other obvious default (e.g., `None`) is suitable, + + Common reasons for using this keyword are: + + - A new keyword is added to a function, and that function forwards its + inputs to another function or method which can be defined outside of + NumPy. For example, ``np.std(x)`` calls ``x.std``, so when a ``keepdims`` + keyword was added that could only be forwarded if the user explicitly + specified ``keepdims``; downstream array libraries may not have added + the same keyword, so adding ``x.std(..., keepdims=keepdims)`` + unconditionally could have broken previously working code. + - A keyword is being deprecated, and a deprecation warning must only be + emitted when the keyword is used. + + """ + __instance = None + def __new__(cls): + # ensure that only one instance exists + if not cls.__instance: + cls.__instance = super().__new__(cls) + return cls.__instance + + def __repr__(self): + return "" + + +_NoValue = _NoValueType() + + +@_set_module("numpy") +class _CopyMode(enum.Enum): + """ + An enumeration for the copy modes supported + by numpy.copy() and numpy.array(). The following three modes are supported, + + - ALWAYS: This means that a deep copy of the input + array will always be taken. + - IF_NEEDED: This means that a deep copy of the input + array will be taken only if necessary. + - NEVER: This means that the deep copy will never be taken. + If a copy cannot be avoided then a `ValueError` will be + raised. + + Note that the buffer-protocol could in theory do copies. NumPy currently + assumes an object exporting the buffer protocol will never do this. + """ + + ALWAYS = True + IF_NEEDED = False + NEVER = 2 + + def __bool__(self): + # For backwards compatibility + if self == _CopyMode.ALWAYS: + return True + + if self == _CopyMode.IF_NEEDED: + return False + + raise ValueError(f"{self} is neither True nor False.") diff --git a/lib/python3.12/site-packages/numpy/_pytesttester.py b/lib/python3.12/site-packages/numpy/_pytesttester.py new file mode 100644 index 0000000000000000000000000000000000000000..1c38291ae3319a08bb665fe5c86dfa13e1655a4c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/_pytesttester.py @@ -0,0 +1,207 @@ +""" +Pytest test running. + +This module implements the ``test()`` function for NumPy modules. The usual +boiler plate for doing that is to put the following in the module +``__init__.py`` file:: + + from numpy._pytesttester import PytestTester + test = PytestTester(__name__) + del PytestTester + + +Warnings filtering and other runtime settings should be dealt with in the +``pytest.ini`` file in the numpy repo root. The behavior of the test depends on +whether or not that file is found as follows: + +* ``pytest.ini`` is present (develop mode) + All warnings except those explicitly filtered out are raised as error. +* ``pytest.ini`` is absent (release mode) + DeprecationWarnings and PendingDeprecationWarnings are ignored, other + warnings are passed through. + +In practice, tests run from the numpy repo are run in develop mode. That +includes the standard ``python runtests.py`` invocation. + +This module is imported by every numpy subpackage, so lies at the top level to +simplify circular import issues. For the same reason, it contains no numpy +imports at module scope, instead importing numpy within function calls. +""" +import sys +import os + +__all__ = ['PytestTester'] + + +def _show_numpy_info(): + import numpy as np + + print("NumPy version %s" % np.__version__) + relaxed_strides = np.ones((10, 1), order="C").flags.f_contiguous + print("NumPy relaxed strides checking option:", relaxed_strides) + info = np.lib.utils._opt_info() + print("NumPy CPU features: ", (info if info else 'nothing enabled')) + + +class PytestTester: + """ + Pytest test runner. + + A test function is typically added to a package's __init__.py like so:: + + from numpy._pytesttester import PytestTester + test = PytestTester(__name__).test + del PytestTester + + Calling this test function finds and runs all tests associated with the + module and all its sub-modules. + + Attributes + ---------- + module_name : str + Full path to the package to test. + + Parameters + ---------- + module_name : module name + The name of the module to test. + + Notes + ----- + Unlike the previous ``nose``-based implementation, this class is not + publicly exposed as it performs some ``numpy``-specific warning + suppression. + + """ + def __init__(self, module_name): + self.module_name = module_name + + def __call__(self, label='fast', verbose=1, extra_argv=None, + doctests=False, coverage=False, durations=-1, tests=None): + """ + Run tests for module using pytest. + + Parameters + ---------- + label : {'fast', 'full'}, optional + Identifies the tests to run. When set to 'fast', tests decorated + with `pytest.mark.slow` are skipped, when 'full', the slow marker + is ignored. + verbose : int, optional + Verbosity value for test outputs, in the range 1-3. Default is 1. + extra_argv : list, optional + List with any extra arguments to pass to pytests. + doctests : bool, optional + .. note:: Not supported + coverage : bool, optional + If True, report coverage of NumPy code. Default is False. + Requires installation of (pip) pytest-cov. + durations : int, optional + If < 0, do nothing, If 0, report time of all tests, if > 0, + report the time of the slowest `timer` tests. Default is -1. + tests : test or list of tests + Tests to be executed with pytest '--pyargs' + + Returns + ------- + result : bool + Return True on success, false otherwise. + + Notes + ----- + Each NumPy module exposes `test` in its namespace to run all tests for + it. For example, to run all tests for numpy.lib: + + >>> np.lib.test() #doctest: +SKIP + + Examples + -------- + >>> result = np.lib.test() #doctest: +SKIP + ... + 1023 passed, 2 skipped, 6 deselected, 1 xfailed in 10.39 seconds + >>> result + True + + """ + import pytest + import warnings + + module = sys.modules[self.module_name] + module_path = os.path.abspath(module.__path__[0]) + + # setup the pytest arguments + pytest_args = ["-l"] + + # offset verbosity. The "-q" cancels a "-v". + pytest_args += ["-q"] + + if sys.version_info < (3, 12): + with warnings.catch_warnings(): + warnings.simplefilter("always") + # Filter out distutils cpu warnings (could be localized to + # distutils tests). ASV has problems with top level import, + # so fetch module for suppression here. + from numpy.distutils import cpuinfo + + with warnings.catch_warnings(record=True): + # Ignore the warning from importing the array_api submodule. This + # warning is done on import, so it would break pytest collection, + # but importing it early here prevents the warning from being + # issued when it imported again. + import numpy.array_api + + # Filter out annoying import messages. Want these in both develop and + # release mode. + pytest_args += [ + "-W ignore:Not importing directory", + "-W ignore:numpy.dtype size changed", + "-W ignore:numpy.ufunc size changed", + "-W ignore::UserWarning:cpuinfo", + ] + + # When testing matrices, ignore their PendingDeprecationWarnings + pytest_args += [ + "-W ignore:the matrix subclass is not", + "-W ignore:Importing from numpy.matlib is", + ] + + if doctests: + pytest_args += ["--doctest-modules"] + + if extra_argv: + pytest_args += list(extra_argv) + + if verbose > 1: + pytest_args += ["-" + "v"*(verbose - 1)] + + if coverage: + pytest_args += ["--cov=" + module_path] + + if label == "fast": + # not importing at the top level to avoid circular import of module + from numpy.testing import IS_PYPY + if IS_PYPY: + pytest_args += ["-m", "not slow and not slow_pypy"] + else: + pytest_args += ["-m", "not slow"] + + elif label != "full": + pytest_args += ["-m", label] + + if durations >= 0: + pytest_args += ["--durations=%s" % durations] + + if tests is None: + tests = [self.module_name] + + pytest_args += ["--pyargs"] + list(tests) + + # run tests. + _show_numpy_info() + + try: + code = pytest.main(pytest_args) + except SystemExit as exc: + code = exc.code + + return code == 0 diff --git a/lib/python3.12/site-packages/numpy/_pytesttester.pyi b/lib/python3.12/site-packages/numpy/_pytesttester.pyi new file mode 100644 index 0000000000000000000000000000000000000000..67ac87b33de164c710a25110d45545e24a06d42e --- /dev/null +++ b/lib/python3.12/site-packages/numpy/_pytesttester.pyi @@ -0,0 +1,18 @@ +from collections.abc import Iterable +from typing import Literal as L + +__all__: list[str] + +class PytestTester: + module_name: str + def __init__(self, module_name: str) -> None: ... + def __call__( + self, + label: L["fast", "full"] = ..., + verbose: int = ..., + extra_argv: None | Iterable[str] = ..., + doctests: L[False] = ..., + coverage: bool = ..., + durations: int = ..., + tests: None | Iterable[str] = ..., + ) -> bool: ... diff --git a/lib/python3.12/site-packages/numpy/conftest.py b/lib/python3.12/site-packages/numpy/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..f1a3eda989057713f3576b60580f2d06b664873c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/conftest.py @@ -0,0 +1,138 @@ +""" +Pytest configuration and fixtures for the Numpy test suite. +""" +import os +import tempfile + +import hypothesis +import pytest +import numpy + +from numpy.core._multiarray_tests import get_fpu_mode + + +_old_fpu_mode = None +_collect_results = {} + +# Use a known and persistent tmpdir for hypothesis' caches, which +# can be automatically cleared by the OS or user. +hypothesis.configuration.set_hypothesis_home_dir( + os.path.join(tempfile.gettempdir(), ".hypothesis") +) + +# We register two custom profiles for Numpy - for details see +# https://hypothesis.readthedocs.io/en/latest/settings.html +# The first is designed for our own CI runs; the latter also +# forces determinism and is designed for use via np.test() +hypothesis.settings.register_profile( + name="numpy-profile", deadline=None, print_blob=True, +) +hypothesis.settings.register_profile( + name="np.test() profile", + deadline=None, print_blob=True, database=None, derandomize=True, + suppress_health_check=list(hypothesis.HealthCheck), +) +# Note that the default profile is chosen based on the presence +# of pytest.ini, but can be overridden by passing the +# --hypothesis-profile=NAME argument to pytest. +_pytest_ini = os.path.join(os.path.dirname(__file__), "..", "pytest.ini") +hypothesis.settings.load_profile( + "numpy-profile" if os.path.isfile(_pytest_ini) else "np.test() profile" +) + +# The experimentalAPI is used in _umath_tests +os.environ["NUMPY_EXPERIMENTAL_DTYPE_API"] = "1" + +def pytest_configure(config): + config.addinivalue_line("markers", + "valgrind_error: Tests that are known to error under valgrind.") + config.addinivalue_line("markers", + "leaks_references: Tests that are known to leak references.") + config.addinivalue_line("markers", + "slow: Tests that are very slow.") + config.addinivalue_line("markers", + "slow_pypy: Tests that are very slow on pypy.") + + +def pytest_addoption(parser): + parser.addoption("--available-memory", action="store", default=None, + help=("Set amount of memory available for running the " + "test suite. This can result to tests requiring " + "especially large amounts of memory to be skipped. " + "Equivalent to setting environment variable " + "NPY_AVAILABLE_MEM. Default: determined" + "automatically.")) + + +def pytest_sessionstart(session): + available_mem = session.config.getoption('available_memory') + if available_mem is not None: + os.environ['NPY_AVAILABLE_MEM'] = available_mem + + +#FIXME when yield tests are gone. +@pytest.hookimpl() +def pytest_itemcollected(item): + """ + Check FPU precision mode was not changed during test collection. + + The clumsy way we do it here is mainly necessary because numpy + still uses yield tests, which can execute code at test collection + time. + """ + global _old_fpu_mode + + mode = get_fpu_mode() + + if _old_fpu_mode is None: + _old_fpu_mode = mode + elif mode != _old_fpu_mode: + _collect_results[item] = (_old_fpu_mode, mode) + _old_fpu_mode = mode + + +@pytest.fixture(scope="function", autouse=True) +def check_fpu_mode(request): + """ + Check FPU precision mode was not changed during the test. + """ + old_mode = get_fpu_mode() + yield + new_mode = get_fpu_mode() + + if old_mode != new_mode: + raise AssertionError("FPU precision mode changed from {0:#x} to {1:#x}" + " during the test".format(old_mode, new_mode)) + + collect_result = _collect_results.get(request.node) + if collect_result is not None: + old_mode, new_mode = collect_result + raise AssertionError("FPU precision mode changed from {0:#x} to {1:#x}" + " when collecting the test".format(old_mode, + new_mode)) + + +@pytest.fixture(autouse=True) +def add_np(doctest_namespace): + doctest_namespace['np'] = numpy + +@pytest.fixture(autouse=True) +def env_setup(monkeypatch): + monkeypatch.setenv('PYTHONHASHSEED', '0') + + +@pytest.fixture(params=[True, False]) +def weak_promotion(request): + """ + Fixture to ensure "legacy" promotion state or change it to use the new + weak promotion (plus warning). `old_promotion` should be used as a + parameter in the function. + """ + state = numpy._get_promotion_state() + if request.param: + numpy._set_promotion_state("weak_and_warn") + else: + numpy._set_promotion_state("legacy") + + yield request.param + numpy._set_promotion_state(state) diff --git a/lib/python3.12/site-packages/numpy/ctypeslib.py b/lib/python3.12/site-packages/numpy/ctypeslib.py new file mode 100644 index 0000000000000000000000000000000000000000..d9f64fd9e716830ff33d4d787a0492c65d517603 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/ctypeslib.py @@ -0,0 +1,545 @@ +""" +============================ +``ctypes`` Utility Functions +============================ + +See Also +-------- +load_library : Load a C library. +ndpointer : Array restype/argtype with verification. +as_ctypes : Create a ctypes array from an ndarray. +as_array : Create an ndarray from a ctypes array. + +References +---------- +.. [1] "SciPy Cookbook: ctypes", https://scipy-cookbook.readthedocs.io/items/Ctypes.html + +Examples +-------- +Load the C library: + +>>> _lib = np.ctypeslib.load_library('libmystuff', '.') #doctest: +SKIP + +Our result type, an ndarray that must be of type double, be 1-dimensional +and is C-contiguous in memory: + +>>> array_1d_double = np.ctypeslib.ndpointer( +... dtype=np.double, +... ndim=1, flags='CONTIGUOUS') #doctest: +SKIP + +Our C-function typically takes an array and updates its values +in-place. For example:: + + void foo_func(double* x, int length) + { + int i; + for (i = 0; i < length; i++) { + x[i] = i*i; + } + } + +We wrap it using: + +>>> _lib.foo_func.restype = None #doctest: +SKIP +>>> _lib.foo_func.argtypes = [array_1d_double, c_int] #doctest: +SKIP + +Then, we're ready to call ``foo_func``: + +>>> out = np.empty(15, dtype=np.double) +>>> _lib.foo_func(out, len(out)) #doctest: +SKIP + +""" +__all__ = ['load_library', 'ndpointer', 'c_intp', 'as_ctypes', 'as_array', + 'as_ctypes_type'] + +import os +from numpy import ( + integer, ndarray, dtype as _dtype, asarray, frombuffer +) +from numpy.core.multiarray import _flagdict, flagsobj + +try: + import ctypes +except ImportError: + ctypes = None + +if ctypes is None: + def _dummy(*args, **kwds): + """ + Dummy object that raises an ImportError if ctypes is not available. + + Raises + ------ + ImportError + If ctypes is not available. + + """ + raise ImportError("ctypes is not available.") + load_library = _dummy + as_ctypes = _dummy + as_array = _dummy + from numpy import intp as c_intp + _ndptr_base = object +else: + import numpy.core._internal as nic + c_intp = nic._getintp_ctype() + del nic + _ndptr_base = ctypes.c_void_p + + # Adapted from Albert Strasheim + def load_library(libname, loader_path): + """ + It is possible to load a library using + + >>> lib = ctypes.cdll[] # doctest: +SKIP + + But there are cross-platform considerations, such as library file extensions, + plus the fact Windows will just load the first library it finds with that name. + NumPy supplies the load_library function as a convenience. + + .. versionchanged:: 1.20.0 + Allow libname and loader_path to take any + :term:`python:path-like object`. + + Parameters + ---------- + libname : path-like + Name of the library, which can have 'lib' as a prefix, + but without an extension. + loader_path : path-like + Where the library can be found. + + Returns + ------- + ctypes.cdll[libpath] : library object + A ctypes library object + + Raises + ------ + OSError + If there is no library with the expected extension, or the + library is defective and cannot be loaded. + """ + # Convert path-like objects into strings + libname = os.fsdecode(libname) + loader_path = os.fsdecode(loader_path) + + ext = os.path.splitext(libname)[1] + if not ext: + import sys + import sysconfig + # Try to load library with platform-specific name, otherwise + # default to libname.[so|dll|dylib]. Sometimes, these files are + # built erroneously on non-linux platforms. + base_ext = ".so" + if sys.platform.startswith("darwin"): + base_ext = ".dylib" + elif sys.platform.startswith("win"): + base_ext = ".dll" + libname_ext = [libname + base_ext] + so_ext = sysconfig.get_config_var("EXT_SUFFIX") + if not so_ext == base_ext: + libname_ext.insert(0, libname + so_ext) + else: + libname_ext = [libname] + + loader_path = os.path.abspath(loader_path) + if not os.path.isdir(loader_path): + libdir = os.path.dirname(loader_path) + else: + libdir = loader_path + + for ln in libname_ext: + libpath = os.path.join(libdir, ln) + if os.path.exists(libpath): + try: + return ctypes.cdll[libpath] + except OSError: + ## defective lib file + raise + ## if no successful return in the libname_ext loop: + raise OSError("no file with expected extension") + + +def _num_fromflags(flaglist): + num = 0 + for val in flaglist: + num += _flagdict[val] + return num + +_flagnames = ['C_CONTIGUOUS', 'F_CONTIGUOUS', 'ALIGNED', 'WRITEABLE', + 'OWNDATA', 'WRITEBACKIFCOPY'] +def _flags_fromnum(num): + res = [] + for key in _flagnames: + value = _flagdict[key] + if (num & value): + res.append(key) + return res + + +class _ndptr(_ndptr_base): + @classmethod + def from_param(cls, obj): + if not isinstance(obj, ndarray): + raise TypeError("argument must be an ndarray") + if cls._dtype_ is not None \ + and obj.dtype != cls._dtype_: + raise TypeError("array must have data type %s" % cls._dtype_) + if cls._ndim_ is not None \ + and obj.ndim != cls._ndim_: + raise TypeError("array must have %d dimension(s)" % cls._ndim_) + if cls._shape_ is not None \ + and obj.shape != cls._shape_: + raise TypeError("array must have shape %s" % str(cls._shape_)) + if cls._flags_ is not None \ + and ((obj.flags.num & cls._flags_) != cls._flags_): + raise TypeError("array must have flags %s" % + _flags_fromnum(cls._flags_)) + return obj.ctypes + + +class _concrete_ndptr(_ndptr): + """ + Like _ndptr, but with `_shape_` and `_dtype_` specified. + + Notably, this means the pointer has enough information to reconstruct + the array, which is not generally true. + """ + def _check_retval_(self): + """ + This method is called when this class is used as the .restype + attribute for a shared-library function, to automatically wrap the + pointer into an array. + """ + return self.contents + + @property + def contents(self): + """ + Get an ndarray viewing the data pointed to by this pointer. + + This mirrors the `contents` attribute of a normal ctypes pointer + """ + full_dtype = _dtype((self._dtype_, self._shape_)) + full_ctype = ctypes.c_char * full_dtype.itemsize + buffer = ctypes.cast(self, ctypes.POINTER(full_ctype)).contents + return frombuffer(buffer, dtype=full_dtype).squeeze(axis=0) + + +# Factory for an array-checking class with from_param defined for +# use with ctypes argtypes mechanism +_pointer_type_cache = {} +def ndpointer(dtype=None, ndim=None, shape=None, flags=None): + """ + Array-checking restype/argtypes. + + An ndpointer instance is used to describe an ndarray in restypes + and argtypes specifications. This approach is more flexible than + using, for example, ``POINTER(c_double)``, since several restrictions + can be specified, which are verified upon calling the ctypes function. + These include data type, number of dimensions, shape and flags. If a + given array does not satisfy the specified restrictions, + a ``TypeError`` is raised. + + Parameters + ---------- + dtype : data-type, optional + Array data-type. + ndim : int, optional + Number of array dimensions. + shape : tuple of ints, optional + Array shape. + flags : str or tuple of str + Array flags; may be one or more of: + + - C_CONTIGUOUS / C / CONTIGUOUS + - F_CONTIGUOUS / F / FORTRAN + - OWNDATA / O + - WRITEABLE / W + - ALIGNED / A + - WRITEBACKIFCOPY / X + + Returns + ------- + klass : ndpointer type object + A type object, which is an ``_ndtpr`` instance containing + dtype, ndim, shape and flags information. + + Raises + ------ + TypeError + If a given array does not satisfy the specified restrictions. + + Examples + -------- + >>> clib.somefunc.argtypes = [np.ctypeslib.ndpointer(dtype=np.float64, + ... ndim=1, + ... flags='C_CONTIGUOUS')] + ... #doctest: +SKIP + >>> clib.somefunc(np.array([1, 2, 3], dtype=np.float64)) + ... #doctest: +SKIP + + """ + + # normalize dtype to an Optional[dtype] + if dtype is not None: + dtype = _dtype(dtype) + + # normalize flags to an Optional[int] + num = None + if flags is not None: + if isinstance(flags, str): + flags = flags.split(',') + elif isinstance(flags, (int, integer)): + num = flags + flags = _flags_fromnum(num) + elif isinstance(flags, flagsobj): + num = flags.num + flags = _flags_fromnum(num) + if num is None: + try: + flags = [x.strip().upper() for x in flags] + except Exception as e: + raise TypeError("invalid flags specification") from e + num = _num_fromflags(flags) + + # normalize shape to an Optional[tuple] + if shape is not None: + try: + shape = tuple(shape) + except TypeError: + # single integer -> 1-tuple + shape = (shape,) + + cache_key = (dtype, ndim, shape, num) + + try: + return _pointer_type_cache[cache_key] + except KeyError: + pass + + # produce a name for the new type + if dtype is None: + name = 'any' + elif dtype.names is not None: + name = str(id(dtype)) + else: + name = dtype.str + if ndim is not None: + name += "_%dd" % ndim + if shape is not None: + name += "_"+"x".join(str(x) for x in shape) + if flags is not None: + name += "_"+"_".join(flags) + + if dtype is not None and shape is not None: + base = _concrete_ndptr + else: + base = _ndptr + + klass = type("ndpointer_%s"%name, (base,), + {"_dtype_": dtype, + "_shape_" : shape, + "_ndim_" : ndim, + "_flags_" : num}) + _pointer_type_cache[cache_key] = klass + return klass + + +if ctypes is not None: + def _ctype_ndarray(element_type, shape): + """ Create an ndarray of the given element type and shape """ + for dim in shape[::-1]: + element_type = dim * element_type + # prevent the type name include np.ctypeslib + element_type.__module__ = None + return element_type + + + def _get_scalar_type_map(): + """ + Return a dictionary mapping native endian scalar dtype to ctypes types + """ + ct = ctypes + simple_types = [ + ct.c_byte, ct.c_short, ct.c_int, ct.c_long, ct.c_longlong, + ct.c_ubyte, ct.c_ushort, ct.c_uint, ct.c_ulong, ct.c_ulonglong, + ct.c_float, ct.c_double, + ct.c_bool, + ] + return {_dtype(ctype): ctype for ctype in simple_types} + + + _scalar_type_map = _get_scalar_type_map() + + + def _ctype_from_dtype_scalar(dtype): + # swapping twice ensure that `=` is promoted to <, >, or | + dtype_with_endian = dtype.newbyteorder('S').newbyteorder('S') + dtype_native = dtype.newbyteorder('=') + try: + ctype = _scalar_type_map[dtype_native] + except KeyError as e: + raise NotImplementedError( + "Converting {!r} to a ctypes type".format(dtype) + ) from None + + if dtype_with_endian.byteorder == '>': + ctype = ctype.__ctype_be__ + elif dtype_with_endian.byteorder == '<': + ctype = ctype.__ctype_le__ + + return ctype + + + def _ctype_from_dtype_subarray(dtype): + element_dtype, shape = dtype.subdtype + ctype = _ctype_from_dtype(element_dtype) + return _ctype_ndarray(ctype, shape) + + + def _ctype_from_dtype_structured(dtype): + # extract offsets of each field + field_data = [] + for name in dtype.names: + field_dtype, offset = dtype.fields[name][:2] + field_data.append((offset, name, _ctype_from_dtype(field_dtype))) + + # ctypes doesn't care about field order + field_data = sorted(field_data, key=lambda f: f[0]) + + if len(field_data) > 1 and all(offset == 0 for offset, name, ctype in field_data): + # union, if multiple fields all at address 0 + size = 0 + _fields_ = [] + for offset, name, ctype in field_data: + _fields_.append((name, ctype)) + size = max(size, ctypes.sizeof(ctype)) + + # pad to the right size + if dtype.itemsize != size: + _fields_.append(('', ctypes.c_char * dtype.itemsize)) + + # we inserted manual padding, so always `_pack_` + return type('union', (ctypes.Union,), dict( + _fields_=_fields_, + _pack_=1, + __module__=None, + )) + else: + last_offset = 0 + _fields_ = [] + for offset, name, ctype in field_data: + padding = offset - last_offset + if padding < 0: + raise NotImplementedError("Overlapping fields") + if padding > 0: + _fields_.append(('', ctypes.c_char * padding)) + + _fields_.append((name, ctype)) + last_offset = offset + ctypes.sizeof(ctype) + + + padding = dtype.itemsize - last_offset + if padding > 0: + _fields_.append(('', ctypes.c_char * padding)) + + # we inserted manual padding, so always `_pack_` + return type('struct', (ctypes.Structure,), dict( + _fields_=_fields_, + _pack_=1, + __module__=None, + )) + + + def _ctype_from_dtype(dtype): + if dtype.fields is not None: + return _ctype_from_dtype_structured(dtype) + elif dtype.subdtype is not None: + return _ctype_from_dtype_subarray(dtype) + else: + return _ctype_from_dtype_scalar(dtype) + + + def as_ctypes_type(dtype): + r""" + Convert a dtype into a ctypes type. + + Parameters + ---------- + dtype : dtype + The dtype to convert + + Returns + ------- + ctype + A ctype scalar, union, array, or struct + + Raises + ------ + NotImplementedError + If the conversion is not possible + + Notes + ----- + This function does not losslessly round-trip in either direction. + + ``np.dtype(as_ctypes_type(dt))`` will: + + - insert padding fields + - reorder fields to be sorted by offset + - discard field titles + + ``as_ctypes_type(np.dtype(ctype))`` will: + + - discard the class names of `ctypes.Structure`\ s and + `ctypes.Union`\ s + - convert single-element `ctypes.Union`\ s into single-element + `ctypes.Structure`\ s + - insert padding fields + + """ + return _ctype_from_dtype(_dtype(dtype)) + + + def as_array(obj, shape=None): + """ + Create a numpy array from a ctypes array or POINTER. + + The numpy array shares the memory with the ctypes object. + + The shape parameter must be given if converting from a ctypes POINTER. + The shape parameter is ignored if converting from a ctypes array + """ + if isinstance(obj, ctypes._Pointer): + # convert pointers to an array of the desired shape + if shape is None: + raise TypeError( + 'as_array() requires a shape argument when called on a ' + 'pointer') + p_arr_type = ctypes.POINTER(_ctype_ndarray(obj._type_, shape)) + obj = ctypes.cast(obj, p_arr_type).contents + + return asarray(obj) + + + def as_ctypes(obj): + """Create and return a ctypes object from a numpy array. Actually + anything that exposes the __array_interface__ is accepted.""" + ai = obj.__array_interface__ + if ai["strides"]: + raise TypeError("strided arrays not supported") + if ai["version"] != 3: + raise TypeError("only __array_interface__ version 3 supported") + addr, readonly = ai["data"] + if readonly: + raise TypeError("readonly arrays unsupported") + + # can't use `_dtype((ai["typestr"], ai["shape"]))` here, as it overflows + # dtype.itemsize (gh-14214) + ctype_scalar = as_ctypes_type(ai["typestr"]) + result_type = _ctype_ndarray(ctype_scalar, ai["shape"]) + result = result_type.from_address(addr) + result.__keep = obj + return result diff --git a/lib/python3.12/site-packages/numpy/ctypeslib.pyi b/lib/python3.12/site-packages/numpy/ctypeslib.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3edf98e143cf17e8a14fe585be75f4735b1fa6e7 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/ctypeslib.pyi @@ -0,0 +1,251 @@ +# NOTE: Numpy's mypy plugin is used for importing the correct +# platform-specific `ctypes._SimpleCData[int]` sub-type +from ctypes import c_int64 as _c_intp + +import os +import sys +import ctypes +from collections.abc import Iterable, Sequence +from typing import ( + Literal as L, + Any, + Union, + TypeVar, + Generic, + overload, + ClassVar, +) + +from numpy import ( + ndarray, + dtype, + generic, + bool_, + byte, + short, + intc, + int_, + longlong, + ubyte, + ushort, + uintc, + uint, + ulonglong, + single, + double, + longdouble, + void, +) +from numpy.core._internal import _ctypes +from numpy.core.multiarray import flagsobj +from numpy._typing import ( + # Arrays + NDArray, + _ArrayLike, + + # Shapes + _ShapeLike, + + # DTypes + DTypeLike, + _DTypeLike, + _VoidDTypeLike, + _BoolCodes, + _UByteCodes, + _UShortCodes, + _UIntCCodes, + _UIntCodes, + _ULongLongCodes, + _ByteCodes, + _ShortCodes, + _IntCCodes, + _IntCodes, + _LongLongCodes, + _SingleCodes, + _DoubleCodes, + _LongDoubleCodes, +) + +# TODO: Add a proper `_Shape` bound once we've got variadic typevars +_DType = TypeVar("_DType", bound=dtype[Any]) +_DTypeOptional = TypeVar("_DTypeOptional", bound=None | dtype[Any]) +_SCT = TypeVar("_SCT", bound=generic) + +_FlagsKind = L[ + 'C_CONTIGUOUS', 'CONTIGUOUS', 'C', + 'F_CONTIGUOUS', 'FORTRAN', 'F', + 'ALIGNED', 'A', + 'WRITEABLE', 'W', + 'OWNDATA', 'O', + 'WRITEBACKIFCOPY', 'X', +] + +# TODO: Add a shape typevar once we have variadic typevars (PEP 646) +class _ndptr(ctypes.c_void_p, Generic[_DTypeOptional]): + # In practice these 4 classvars are defined in the dynamic class + # returned by `ndpointer` + _dtype_: ClassVar[_DTypeOptional] + _shape_: ClassVar[None] + _ndim_: ClassVar[None | int] + _flags_: ClassVar[None | list[_FlagsKind]] + + @overload + @classmethod + def from_param(cls: type[_ndptr[None]], obj: ndarray[Any, Any]) -> _ctypes[Any]: ... + @overload + @classmethod + def from_param(cls: type[_ndptr[_DType]], obj: ndarray[Any, _DType]) -> _ctypes[Any]: ... + +class _concrete_ndptr(_ndptr[_DType]): + _dtype_: ClassVar[_DType] + _shape_: ClassVar[tuple[int, ...]] + @property + def contents(self) -> ndarray[Any, _DType]: ... + +def load_library( + libname: str | bytes | os.PathLike[str] | os.PathLike[bytes], + loader_path: str | bytes | os.PathLike[str] | os.PathLike[bytes], +) -> ctypes.CDLL: ... + +__all__: list[str] + +c_intp = _c_intp + +@overload +def ndpointer( + dtype: None = ..., + ndim: int = ..., + shape: None | _ShapeLike = ..., + flags: None | _FlagsKind | Iterable[_FlagsKind] | int | flagsobj = ..., +) -> type[_ndptr[None]]: ... +@overload +def ndpointer( + dtype: _DTypeLike[_SCT], + ndim: int = ..., + *, + shape: _ShapeLike, + flags: None | _FlagsKind | Iterable[_FlagsKind] | int | flagsobj = ..., +) -> type[_concrete_ndptr[dtype[_SCT]]]: ... +@overload +def ndpointer( + dtype: DTypeLike, + ndim: int = ..., + *, + shape: _ShapeLike, + flags: None | _FlagsKind | Iterable[_FlagsKind] | int | flagsobj = ..., +) -> type[_concrete_ndptr[dtype[Any]]]: ... +@overload +def ndpointer( + dtype: _DTypeLike[_SCT], + ndim: int = ..., + shape: None = ..., + flags: None | _FlagsKind | Iterable[_FlagsKind] | int | flagsobj = ..., +) -> type[_ndptr[dtype[_SCT]]]: ... +@overload +def ndpointer( + dtype: DTypeLike, + ndim: int = ..., + shape: None = ..., + flags: None | _FlagsKind | Iterable[_FlagsKind] | int | flagsobj = ..., +) -> type[_ndptr[dtype[Any]]]: ... + +@overload +def as_ctypes_type(dtype: _BoolCodes | _DTypeLike[bool_] | type[ctypes.c_bool]) -> type[ctypes.c_bool]: ... +@overload +def as_ctypes_type(dtype: _ByteCodes | _DTypeLike[byte] | type[ctypes.c_byte]) -> type[ctypes.c_byte]: ... +@overload +def as_ctypes_type(dtype: _ShortCodes | _DTypeLike[short] | type[ctypes.c_short]) -> type[ctypes.c_short]: ... +@overload +def as_ctypes_type(dtype: _IntCCodes | _DTypeLike[intc] | type[ctypes.c_int]) -> type[ctypes.c_int]: ... +@overload +def as_ctypes_type(dtype: _IntCodes | _DTypeLike[int_] | type[int | ctypes.c_long]) -> type[ctypes.c_long]: ... +@overload +def as_ctypes_type(dtype: _LongLongCodes | _DTypeLike[longlong] | type[ctypes.c_longlong]) -> type[ctypes.c_longlong]: ... +@overload +def as_ctypes_type(dtype: _UByteCodes | _DTypeLike[ubyte] | type[ctypes.c_ubyte]) -> type[ctypes.c_ubyte]: ... +@overload +def as_ctypes_type(dtype: _UShortCodes | _DTypeLike[ushort] | type[ctypes.c_ushort]) -> type[ctypes.c_ushort]: ... +@overload +def as_ctypes_type(dtype: _UIntCCodes | _DTypeLike[uintc] | type[ctypes.c_uint]) -> type[ctypes.c_uint]: ... +@overload +def as_ctypes_type(dtype: _UIntCodes | _DTypeLike[uint] | type[ctypes.c_ulong]) -> type[ctypes.c_ulong]: ... +@overload +def as_ctypes_type(dtype: _ULongLongCodes | _DTypeLike[ulonglong] | type[ctypes.c_ulonglong]) -> type[ctypes.c_ulonglong]: ... +@overload +def as_ctypes_type(dtype: _SingleCodes | _DTypeLike[single] | type[ctypes.c_float]) -> type[ctypes.c_float]: ... +@overload +def as_ctypes_type(dtype: _DoubleCodes | _DTypeLike[double] | type[float | ctypes.c_double]) -> type[ctypes.c_double]: ... +@overload +def as_ctypes_type(dtype: _LongDoubleCodes | _DTypeLike[longdouble] | type[ctypes.c_longdouble]) -> type[ctypes.c_longdouble]: ... +@overload +def as_ctypes_type(dtype: _VoidDTypeLike) -> type[Any]: ... # `ctypes.Union` or `ctypes.Structure` +@overload +def as_ctypes_type(dtype: str) -> type[Any]: ... + +@overload +def as_array(obj: ctypes._PointerLike, shape: Sequence[int]) -> NDArray[Any]: ... +@overload +def as_array(obj: _ArrayLike[_SCT], shape: None | _ShapeLike = ...) -> NDArray[_SCT]: ... +@overload +def as_array(obj: object, shape: None | _ShapeLike = ...) -> NDArray[Any]: ... + +@overload +def as_ctypes(obj: bool_) -> ctypes.c_bool: ... +@overload +def as_ctypes(obj: byte) -> ctypes.c_byte: ... +@overload +def as_ctypes(obj: short) -> ctypes.c_short: ... +@overload +def as_ctypes(obj: intc) -> ctypes.c_int: ... +@overload +def as_ctypes(obj: int_) -> ctypes.c_long: ... +@overload +def as_ctypes(obj: longlong) -> ctypes.c_longlong: ... +@overload +def as_ctypes(obj: ubyte) -> ctypes.c_ubyte: ... +@overload +def as_ctypes(obj: ushort) -> ctypes.c_ushort: ... +@overload +def as_ctypes(obj: uintc) -> ctypes.c_uint: ... +@overload +def as_ctypes(obj: uint) -> ctypes.c_ulong: ... +@overload +def as_ctypes(obj: ulonglong) -> ctypes.c_ulonglong: ... +@overload +def as_ctypes(obj: single) -> ctypes.c_float: ... +@overload +def as_ctypes(obj: double) -> ctypes.c_double: ... +@overload +def as_ctypes(obj: longdouble) -> ctypes.c_longdouble: ... +@overload +def as_ctypes(obj: void) -> Any: ... # `ctypes.Union` or `ctypes.Structure` +@overload +def as_ctypes(obj: NDArray[bool_]) -> ctypes.Array[ctypes.c_bool]: ... +@overload +def as_ctypes(obj: NDArray[byte]) -> ctypes.Array[ctypes.c_byte]: ... +@overload +def as_ctypes(obj: NDArray[short]) -> ctypes.Array[ctypes.c_short]: ... +@overload +def as_ctypes(obj: NDArray[intc]) -> ctypes.Array[ctypes.c_int]: ... +@overload +def as_ctypes(obj: NDArray[int_]) -> ctypes.Array[ctypes.c_long]: ... +@overload +def as_ctypes(obj: NDArray[longlong]) -> ctypes.Array[ctypes.c_longlong]: ... +@overload +def as_ctypes(obj: NDArray[ubyte]) -> ctypes.Array[ctypes.c_ubyte]: ... +@overload +def as_ctypes(obj: NDArray[ushort]) -> ctypes.Array[ctypes.c_ushort]: ... +@overload +def as_ctypes(obj: NDArray[uintc]) -> ctypes.Array[ctypes.c_uint]: ... +@overload +def as_ctypes(obj: NDArray[uint]) -> ctypes.Array[ctypes.c_ulong]: ... +@overload +def as_ctypes(obj: NDArray[ulonglong]) -> ctypes.Array[ctypes.c_ulonglong]: ... +@overload +def as_ctypes(obj: NDArray[single]) -> ctypes.Array[ctypes.c_float]: ... +@overload +def as_ctypes(obj: NDArray[double]) -> ctypes.Array[ctypes.c_double]: ... +@overload +def as_ctypes(obj: NDArray[longdouble]) -> ctypes.Array[ctypes.c_longdouble]: ... +@overload +def as_ctypes(obj: NDArray[void]) -> ctypes.Array[Any]: ... # `ctypes.Union` or `ctypes.Structure` diff --git a/lib/python3.12/site-packages/numpy/dtypes.py b/lib/python3.12/site-packages/numpy/dtypes.py new file mode 100644 index 0000000000000000000000000000000000000000..068a6a1a0f5b5382a7d0c4fcc2b6cd33f989fdfa --- /dev/null +++ b/lib/python3.12/site-packages/numpy/dtypes.py @@ -0,0 +1,77 @@ +""" +DType classes and utility (:mod:`numpy.dtypes`) +=============================================== + +This module is home to specific dtypes related functionality and their classes. +For more general information about dtypes, also see `numpy.dtype` and +:ref:`arrays.dtypes`. + +Similar to the builtin ``types`` module, this submodule defines types (classes) +that are not widely used directly. + +.. versionadded:: NumPy 1.25 + + The dtypes module is new in NumPy 1.25. Previously DType classes were + only accessible indirectly. + + +DType classes +------------- + +The following are the classes of the corresponding NumPy dtype instances and +NumPy scalar types. The classes can be used in ``isinstance`` checks and can +also be instantiated or used directly. Direct use of these classes is not +typical, since their scalar counterparts (e.g. ``np.float64``) or strings +like ``"float64"`` can be used. + +.. list-table:: + :header-rows: 1 + + * - Group + - DType class + + * - Boolean + - ``BoolDType`` + + * - Bit-sized integers + - ``Int8DType``, ``UInt8DType``, ``Int16DType``, ``UInt16DType``, + ``Int32DType``, ``UInt32DType``, ``Int64DType``, ``UInt64DType`` + + * - C-named integers (may be aliases) + - ``ByteDType``, ``UByteDType``, ``ShortDType``, ``UShortDType``, + ``IntDType``, ``UIntDType``, ``LongDType``, ``ULongDType``, + ``LongLongDType``, ``ULongLongDType`` + + * - Floating point + - ``Float16DType``, ``Float32DType``, ``Float64DType``, + ``LongDoubleDType`` + + * - Complex + - ``Complex64DType``, ``Complex128DType``, ``CLongDoubleDType`` + + * - Strings + - ``BytesDType``, ``BytesDType`` + + * - Times + - ``DateTime64DType``, ``TimeDelta64DType`` + + * - Others + - ``ObjectDType``, ``VoidDType`` + +""" + +__all__ = [] + + +def _add_dtype_helper(DType, alias): + # Function to add DTypes a bit more conveniently without channeling them + # through `numpy.core._multiarray_umath` namespace or similar. + from numpy import dtypes + + setattr(dtypes, DType.__name__, DType) + __all__.append(DType.__name__) + + if alias: + alias = alias.removeprefix("numpy.dtypes.") + setattr(dtypes, alias, DType) + __all__.append(alias) diff --git a/lib/python3.12/site-packages/numpy/dtypes.pyi b/lib/python3.12/site-packages/numpy/dtypes.pyi new file mode 100644 index 0000000000000000000000000000000000000000..2f7e846f23d4de0dd7caa3198e3eb4fd339ebdbe --- /dev/null +++ b/lib/python3.12/site-packages/numpy/dtypes.pyi @@ -0,0 +1,43 @@ +import numpy as np + + +__all__: list[str] + +# Boolean: +BoolDType = np.dtype[np.bool_] +# Sized integers: +Int8DType = np.dtype[np.int8] +UInt8DType = np.dtype[np.uint8] +Int16DType = np.dtype[np.int16] +UInt16DType = np.dtype[np.uint16] +Int32DType = np.dtype[np.int32] +UInt32DType = np.dtype[np.uint32] +Int64DType = np.dtype[np.int64] +UInt64DType = np.dtype[np.uint64] +# Standard C-named version/alias: +ByteDType = np.dtype[np.byte] +UByteDType = np.dtype[np.ubyte] +ShortDType = np.dtype[np.short] +UShortDType = np.dtype[np.ushort] +IntDType = np.dtype[np.intc] +UIntDType = np.dtype[np.uintc] +LongDType = np.dtype[np.int_] # Unfortunately, the correct scalar +ULongDType = np.dtype[np.uint] # Unfortunately, the correct scalar +LongLongDType = np.dtype[np.longlong] +ULongLongDType = np.dtype[np.ulonglong] +# Floats +Float16DType = np.dtype[np.float16] +Float32DType = np.dtype[np.float32] +Float64DType = np.dtype[np.float64] +LongDoubleDType = np.dtype[np.longdouble] +# Complex: +Complex64DType = np.dtype[np.complex64] +Complex128DType = np.dtype[np.complex128] +CLongDoubleDType = np.dtype[np.clongdouble] +# Others: +ObjectDType = np.dtype[np.object_] +BytesDType = np.dtype[np.bytes_] +StrDType = np.dtype[np.str_] +VoidDType = np.dtype[np.void] +DateTime64DType = np.dtype[np.datetime64] +TimeDelta64DType = np.dtype[np.timedelta64] diff --git a/lib/python3.12/site-packages/numpy/exceptions.py b/lib/python3.12/site-packages/numpy/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..2f843810141a7c2d78de9ff75f5a0db9e592c981 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/exceptions.py @@ -0,0 +1,231 @@ +""" +Exceptions and Warnings (:mod:`numpy.exceptions`) +================================================= + +General exceptions used by NumPy. Note that some exceptions may be module +specific, such as linear algebra errors. + +.. versionadded:: NumPy 1.25 + + The exceptions module is new in NumPy 1.25. Older exceptions remain + available through the main NumPy namespace for compatibility. + +.. currentmodule:: numpy.exceptions + +Warnings +-------- +.. autosummary:: + :toctree: generated/ + + ComplexWarning Given when converting complex to real. + VisibleDeprecationWarning Same as a DeprecationWarning, but more visible. + +Exceptions +---------- +.. autosummary:: + :toctree: generated/ + + AxisError Given when an axis was invalid. + DTypePromotionError Given when no common dtype could be found. + TooHardError Error specific to `numpy.shares_memory`. + +""" + + +__all__ = [ + "ComplexWarning", "VisibleDeprecationWarning", "ModuleDeprecationWarning", + "TooHardError", "AxisError", "DTypePromotionError"] + + +# Disallow reloading this module so as to preserve the identities of the +# classes defined here. +if '_is_loaded' in globals(): + raise RuntimeError('Reloading numpy._globals is not allowed') +_is_loaded = True + + +class ComplexWarning(RuntimeWarning): + """ + The warning raised when casting a complex dtype to a real dtype. + + As implemented, casting a complex number to a real discards its imaginary + part, but this behavior may not be what the user actually wants. + + """ + pass + + +class ModuleDeprecationWarning(DeprecationWarning): + """Module deprecation warning. + + .. warning:: + + This warning should not be used, since nose testing is not relevant + anymore. + + The nose tester turns ordinary Deprecation warnings into test failures. + That makes it hard to deprecate whole modules, because they get + imported by default. So this is a special Deprecation warning that the + nose tester will let pass without making tests fail. + + """ + + +class VisibleDeprecationWarning(UserWarning): + """Visible deprecation warning. + + By default, python will not show deprecation warnings, so this class + can be used when a very visible warning is helpful, for example because + the usage is most likely a user bug. + + """ + + +# Exception used in shares_memory() +class TooHardError(RuntimeError): + """max_work was exceeded. + + This is raised whenever the maximum number of candidate solutions + to consider specified by the ``max_work`` parameter is exceeded. + Assigning a finite number to max_work may have caused the operation + to fail. + + """ + + pass + + +class AxisError(ValueError, IndexError): + """Axis supplied was invalid. + + This is raised whenever an ``axis`` parameter is specified that is larger + than the number of array dimensions. + For compatibility with code written against older numpy versions, which + raised a mixture of `ValueError` and `IndexError` for this situation, this + exception subclasses both to ensure that ``except ValueError`` and + ``except IndexError`` statements continue to catch `AxisError`. + + .. versionadded:: 1.13 + + Parameters + ---------- + axis : int or str + The out of bounds axis or a custom exception message. + If an axis is provided, then `ndim` should be specified as well. + ndim : int, optional + The number of array dimensions. + msg_prefix : str, optional + A prefix for the exception message. + + Attributes + ---------- + axis : int, optional + The out of bounds axis or ``None`` if a custom exception + message was provided. This should be the axis as passed by + the user, before any normalization to resolve negative indices. + + .. versionadded:: 1.22 + ndim : int, optional + The number of array dimensions or ``None`` if a custom exception + message was provided. + + .. versionadded:: 1.22 + + + Examples + -------- + >>> array_1d = np.arange(10) + >>> np.cumsum(array_1d, axis=1) + Traceback (most recent call last): + ... + numpy.exceptions.AxisError: axis 1 is out of bounds for array of dimension 1 + + Negative axes are preserved: + + >>> np.cumsum(array_1d, axis=-2) + Traceback (most recent call last): + ... + numpy.exceptions.AxisError: axis -2 is out of bounds for array of dimension 1 + + The class constructor generally takes the axis and arrays' + dimensionality as arguments: + + >>> print(np.AxisError(2, 1, msg_prefix='error')) + error: axis 2 is out of bounds for array of dimension 1 + + Alternatively, a custom exception message can be passed: + + >>> print(np.AxisError('Custom error message')) + Custom error message + + """ + + __slots__ = ("axis", "ndim", "_msg") + + def __init__(self, axis, ndim=None, msg_prefix=None): + if ndim is msg_prefix is None: + # single-argument form: directly set the error message + self._msg = axis + self.axis = None + self.ndim = None + else: + self._msg = msg_prefix + self.axis = axis + self.ndim = ndim + + def __str__(self): + axis = self.axis + ndim = self.ndim + + if axis is ndim is None: + return self._msg + else: + msg = f"axis {axis} is out of bounds for array of dimension {ndim}" + if self._msg is not None: + msg = f"{self._msg}: {msg}" + return msg + + +class DTypePromotionError(TypeError): + """Multiple DTypes could not be converted to a common one. + + This exception derives from ``TypeError`` and is raised whenever dtypes + cannot be converted to a single common one. This can be because they + are of a different category/class or incompatible instances of the same + one (see Examples). + + Notes + ----- + Many functions will use promotion to find the correct result and + implementation. For these functions the error will typically be chained + with a more specific error indicating that no implementation was found + for the input dtypes. + + Typically promotion should be considered "invalid" between the dtypes of + two arrays when `arr1 == arr2` can safely return all ``False`` because the + dtypes are fundamentally different. + + Examples + -------- + Datetimes and complex numbers are incompatible classes and cannot be + promoted: + + >>> np.result_type(np.dtype("M8[s]"), np.complex128) + DTypePromotionError: The DType could not + be promoted by . This means that no common + DType exists for the given inputs. For example they cannot be stored in a + single array unless the dtype is `object`. The full list of DTypes is: + (, ) + + For example for structured dtypes, the structure can mismatch and the + same ``DTypePromotionError`` is given when two structured dtypes with + a mismatch in their number of fields is given: + + >>> dtype1 = np.dtype([("field1", np.float64), ("field2", np.int64)]) + >>> dtype2 = np.dtype([("field1", np.float64)]) + >>> np.promote_types(dtype1, dtype2) + DTypePromotionError: field names `('field1', 'field2')` and `('field1',)` + mismatch. + + """ + pass diff --git a/lib/python3.12/site-packages/numpy/exceptions.pyi b/lib/python3.12/site-packages/numpy/exceptions.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c76a0946b97b088c9f0c431eb559b5a3c86a4f6b --- /dev/null +++ b/lib/python3.12/site-packages/numpy/exceptions.pyi @@ -0,0 +1,18 @@ +from typing import overload + +__all__: list[str] + +class ComplexWarning(RuntimeWarning): ... +class ModuleDeprecationWarning(DeprecationWarning): ... +class VisibleDeprecationWarning(UserWarning): ... +class TooHardError(RuntimeError): ... +class DTypePromotionError(TypeError): ... + +class AxisError(ValueError, IndexError): + axis: None | int + ndim: None | int + @overload + def __init__(self, axis: str, ndim: None = ..., msg_prefix: None = ...) -> None: ... + @overload + def __init__(self, axis: int, ndim: int, msg_prefix: None | str = ...) -> None: ... + def __str__(self) -> str: ... diff --git a/lib/python3.12/site-packages/numpy/matlib.py b/lib/python3.12/site-packages/numpy/matlib.py new file mode 100644 index 0000000000000000000000000000000000000000..e929fd9b1885f208afb6301f19cc21511adc098b --- /dev/null +++ b/lib/python3.12/site-packages/numpy/matlib.py @@ -0,0 +1,378 @@ +import warnings + +# 2018-05-29, PendingDeprecationWarning added to matrix.__new__ +# 2020-01-23, numpy 1.19.0 PendingDeprecatonWarning +warnings.warn("Importing from numpy.matlib is deprecated since 1.19.0. " + "The matrix subclass is not the recommended way to represent " + "matrices or deal with linear algebra (see " + "https://docs.scipy.org/doc/numpy/user/numpy-for-matlab-users.html). " + "Please adjust your code to use regular ndarray. ", + PendingDeprecationWarning, stacklevel=2) + +import numpy as np +from numpy.matrixlib.defmatrix import matrix, asmatrix +# Matlib.py contains all functions in the numpy namespace with a few +# replacements. See doc/source/reference/routines.matlib.rst for details. +# Need * as we're copying the numpy namespace. +from numpy import * # noqa: F403 + +__version__ = np.__version__ + +__all__ = np.__all__[:] # copy numpy namespace +__all__ += ['rand', 'randn', 'repmat'] + +def empty(shape, dtype=None, order='C'): + """Return a new matrix of given shape and type, without initializing entries. + + Parameters + ---------- + shape : int or tuple of int + Shape of the empty matrix. + dtype : data-type, optional + Desired output data-type. + order : {'C', 'F'}, optional + Whether to store multi-dimensional data in row-major + (C-style) or column-major (Fortran-style) order in + memory. + + See Also + -------- + empty_like, zeros + + Notes + ----- + `empty`, unlike `zeros`, does not set the matrix values to zero, + and may therefore be marginally faster. On the other hand, it requires + the user to manually set all the values in the array, and should be + used with caution. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.empty((2, 2)) # filled with random data + matrix([[ 6.76425276e-320, 9.79033856e-307], # random + [ 7.39337286e-309, 3.22135945e-309]]) + >>> np.matlib.empty((2, 2), dtype=int) + matrix([[ 6600475, 0], # random + [ 6586976, 22740995]]) + + """ + return ndarray.__new__(matrix, shape, dtype, order=order) + +def ones(shape, dtype=None, order='C'): + """ + Matrix of ones. + + Return a matrix of given shape and type, filled with ones. + + Parameters + ---------- + shape : {sequence of ints, int} + Shape of the matrix + dtype : data-type, optional + The desired data-type for the matrix, default is np.float64. + order : {'C', 'F'}, optional + Whether to store matrix in C- or Fortran-contiguous order, + default is 'C'. + + Returns + ------- + out : matrix + Matrix of ones of given shape, dtype, and order. + + See Also + -------- + ones : Array of ones. + matlib.zeros : Zero matrix. + + Notes + ----- + If `shape` has length one i.e. ``(N,)``, or is a scalar ``N``, + `out` becomes a single row matrix of shape ``(1,N)``. + + Examples + -------- + >>> np.matlib.ones((2,3)) + matrix([[1., 1., 1.], + [1., 1., 1.]]) + + >>> np.matlib.ones(2) + matrix([[1., 1.]]) + + """ + a = ndarray.__new__(matrix, shape, dtype, order=order) + a.fill(1) + return a + +def zeros(shape, dtype=None, order='C'): + """ + Return a matrix of given shape and type, filled with zeros. + + Parameters + ---------- + shape : int or sequence of ints + Shape of the matrix + dtype : data-type, optional + The desired data-type for the matrix, default is float. + order : {'C', 'F'}, optional + Whether to store the result in C- or Fortran-contiguous order, + default is 'C'. + + Returns + ------- + out : matrix + Zero matrix of given shape, dtype, and order. + + See Also + -------- + numpy.zeros : Equivalent array function. + matlib.ones : Return a matrix of ones. + + Notes + ----- + If `shape` has length one i.e. ``(N,)``, or is a scalar ``N``, + `out` becomes a single row matrix of shape ``(1,N)``. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.zeros((2, 3)) + matrix([[0., 0., 0.], + [0., 0., 0.]]) + + >>> np.matlib.zeros(2) + matrix([[0., 0.]]) + + """ + a = ndarray.__new__(matrix, shape, dtype, order=order) + a.fill(0) + return a + +def identity(n,dtype=None): + """ + Returns the square identity matrix of given size. + + Parameters + ---------- + n : int + Size of the returned identity matrix. + dtype : data-type, optional + Data-type of the output. Defaults to ``float``. + + Returns + ------- + out : matrix + `n` x `n` matrix with its main diagonal set to one, + and all other elements zero. + + See Also + -------- + numpy.identity : Equivalent array function. + matlib.eye : More general matrix identity function. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.identity(3, dtype=int) + matrix([[1, 0, 0], + [0, 1, 0], + [0, 0, 1]]) + + """ + a = array([1]+n*[0], dtype=dtype) + b = empty((n, n), dtype=dtype) + b.flat = a + return b + +def eye(n,M=None, k=0, dtype=float, order='C'): + """ + Return a matrix with ones on the diagonal and zeros elsewhere. + + Parameters + ---------- + n : int + Number of rows in the output. + M : int, optional + Number of columns in the output, defaults to `n`. + k : int, optional + Index of the diagonal: 0 refers to the main diagonal, + a positive value refers to an upper diagonal, + and a negative value to a lower diagonal. + dtype : dtype, optional + Data-type of the returned matrix. + order : {'C', 'F'}, optional + Whether the output should be stored in row-major (C-style) or + column-major (Fortran-style) order in memory. + + .. versionadded:: 1.14.0 + + Returns + ------- + I : matrix + A `n` x `M` matrix where all elements are equal to zero, + except for the `k`-th diagonal, whose values are equal to one. + + See Also + -------- + numpy.eye : Equivalent array function. + identity : Square identity matrix. + + Examples + -------- + >>> import numpy.matlib + >>> np.matlib.eye(3, k=1, dtype=float) + matrix([[0., 1., 0.], + [0., 0., 1.], + [0., 0., 0.]]) + + """ + return asmatrix(np.eye(n, M=M, k=k, dtype=dtype, order=order)) + +def rand(*args): + """ + Return a matrix of random values with given shape. + + Create a matrix of the given shape and propagate it with + random samples from a uniform distribution over ``[0, 1)``. + + Parameters + ---------- + \\*args : Arguments + Shape of the output. + If given as N integers, each integer specifies the size of one + dimension. + If given as a tuple, this tuple gives the complete shape. + + Returns + ------- + out : ndarray + The matrix of random values with shape given by `\\*args`. + + See Also + -------- + randn, numpy.random.RandomState.rand + + Examples + -------- + >>> np.random.seed(123) + >>> import numpy.matlib + >>> np.matlib.rand(2, 3) + matrix([[0.69646919, 0.28613933, 0.22685145], + [0.55131477, 0.71946897, 0.42310646]]) + >>> np.matlib.rand((2, 3)) + matrix([[0.9807642 , 0.68482974, 0.4809319 ], + [0.39211752, 0.34317802, 0.72904971]]) + + If the first argument is a tuple, other arguments are ignored: + + >>> np.matlib.rand((2, 3), 4) + matrix([[0.43857224, 0.0596779 , 0.39804426], + [0.73799541, 0.18249173, 0.17545176]]) + + """ + if isinstance(args[0], tuple): + args = args[0] + return asmatrix(np.random.rand(*args)) + +def randn(*args): + """ + Return a random matrix with data from the "standard normal" distribution. + + `randn` generates a matrix filled with random floats sampled from a + univariate "normal" (Gaussian) distribution of mean 0 and variance 1. + + Parameters + ---------- + \\*args : Arguments + Shape of the output. + If given as N integers, each integer specifies the size of one + dimension. If given as a tuple, this tuple gives the complete shape. + + Returns + ------- + Z : matrix of floats + A matrix of floating-point samples drawn from the standard normal + distribution. + + See Also + -------- + rand, numpy.random.RandomState.randn + + Notes + ----- + For random samples from the normal distribution with mean ``mu`` and + standard deviation ``sigma``, use:: + + sigma * np.matlib.randn(...) + mu + + Examples + -------- + >>> np.random.seed(123) + >>> import numpy.matlib + >>> np.matlib.randn(1) + matrix([[-1.0856306]]) + >>> np.matlib.randn(1, 2, 3) + matrix([[ 0.99734545, 0.2829785 , -1.50629471], + [-0.57860025, 1.65143654, -2.42667924]]) + + Two-by-four matrix of samples from the normal distribution with + mean 3 and standard deviation 2.5: + + >>> 2.5 * np.matlib.randn((2, 4)) + 3 + matrix([[1.92771843, 6.16484065, 0.83314899, 1.30278462], + [2.76322758, 6.72847407, 1.40274501, 1.8900451 ]]) + + """ + if isinstance(args[0], tuple): + args = args[0] + return asmatrix(np.random.randn(*args)) + +def repmat(a, m, n): + """ + Repeat a 0-D to 2-D array or matrix MxN times. + + Parameters + ---------- + a : array_like + The array or matrix to be repeated. + m, n : int + The number of times `a` is repeated along the first and second axes. + + Returns + ------- + out : ndarray + The result of repeating `a`. + + Examples + -------- + >>> import numpy.matlib + >>> a0 = np.array(1) + >>> np.matlib.repmat(a0, 2, 3) + array([[1, 1, 1], + [1, 1, 1]]) + + >>> a1 = np.arange(4) + >>> np.matlib.repmat(a1, 2, 2) + array([[0, 1, 2, 3, 0, 1, 2, 3], + [0, 1, 2, 3, 0, 1, 2, 3]]) + + >>> a2 = np.asmatrix(np.arange(6).reshape(2, 3)) + >>> np.matlib.repmat(a2, 2, 3) + matrix([[0, 1, 2, 0, 1, 2, 0, 1, 2], + [3, 4, 5, 3, 4, 5, 3, 4, 5], + [0, 1, 2, 0, 1, 2, 0, 1, 2], + [3, 4, 5, 3, 4, 5, 3, 4, 5]]) + + """ + a = asanyarray(a) + ndim = a.ndim + if ndim == 0: + origrows, origcols = (1, 1) + elif ndim == 1: + origrows, origcols = (1, a.shape[0]) + else: + origrows, origcols = a.shape + rows = origrows * m + cols = origcols * n + c = a.reshape(1, a.size).repeat(m, 0).reshape(rows, origcols).repeat(n, 0) + return c.reshape(rows, cols) diff --git a/lib/python3.12/site-packages/numpy/polynomial/__init__.py b/lib/python3.12/site-packages/numpy/polynomial/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e7baf2c683e27fca27f81e72c348fe8d225089 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/__init__.py @@ -0,0 +1,185 @@ +""" +A sub-package for efficiently dealing with polynomials. + +Within the documentation for this sub-package, a "finite power series," +i.e., a polynomial (also referred to simply as a "series") is represented +by a 1-D numpy array of the polynomial's coefficients, ordered from lowest +order term to highest. For example, array([1,2,3]) represents +``P_0 + 2*P_1 + 3*P_2``, where P_n is the n-th order basis polynomial +applicable to the specific module in question, e.g., `polynomial` (which +"wraps" the "standard" basis) or `chebyshev`. For optimal performance, +all operations on polynomials, including evaluation at an argument, are +implemented as operations on the coefficients. Additional (module-specific) +information can be found in the docstring for the module of interest. + +This package provides *convenience classes* for each of six different kinds +of polynomials: + + ======================== ================ + **Name** **Provides** + ======================== ================ + `~polynomial.Polynomial` Power series + `~chebyshev.Chebyshev` Chebyshev series + `~legendre.Legendre` Legendre series + `~laguerre.Laguerre` Laguerre series + `~hermite.Hermite` Hermite series + `~hermite_e.HermiteE` HermiteE series + ======================== ================ + +These *convenience classes* provide a consistent interface for creating, +manipulating, and fitting data with polynomials of different bases. +The convenience classes are the preferred interface for the `~numpy.polynomial` +package, and are available from the ``numpy.polynomial`` namespace. +This eliminates the need to navigate to the corresponding submodules, e.g. +``np.polynomial.Polynomial`` or ``np.polynomial.Chebyshev`` instead of +``np.polynomial.polynomial.Polynomial`` or +``np.polynomial.chebyshev.Chebyshev``, respectively. +The classes provide a more consistent and concise interface than the +type-specific functions defined in the submodules for each type of polynomial. +For example, to fit a Chebyshev polynomial with degree ``1`` to data given +by arrays ``xdata`` and ``ydata``, the +`~chebyshev.Chebyshev.fit` class method:: + + >>> from numpy.polynomial import Chebyshev + >>> c = Chebyshev.fit(xdata, ydata, deg=1) + +is preferred over the `chebyshev.chebfit` function from the +``np.polynomial.chebyshev`` module:: + + >>> from numpy.polynomial.chebyshev import chebfit + >>> c = chebfit(xdata, ydata, deg=1) + +See :doc:`routines.polynomials.classes` for more details. + +Convenience Classes +=================== + +The following lists the various constants and methods common to all of +the classes representing the various kinds of polynomials. In the following, +the term ``Poly`` represents any one of the convenience classes (e.g. +`~polynomial.Polynomial`, `~chebyshev.Chebyshev`, `~hermite.Hermite`, etc.) +while the lowercase ``p`` represents an **instance** of a polynomial class. + +Constants +--------- + +- ``Poly.domain`` -- Default domain +- ``Poly.window`` -- Default window +- ``Poly.basis_name`` -- String used to represent the basis +- ``Poly.maxpower`` -- Maximum value ``n`` such that ``p**n`` is allowed +- ``Poly.nickname`` -- String used in printing + +Creation +-------- + +Methods for creating polynomial instances. + +- ``Poly.basis(degree)`` -- Basis polynomial of given degree +- ``Poly.identity()`` -- ``p`` where ``p(x) = x`` for all ``x`` +- ``Poly.fit(x, y, deg)`` -- ``p`` of degree ``deg`` with coefficients + determined by the least-squares fit to the data ``x``, ``y`` +- ``Poly.fromroots(roots)`` -- ``p`` with specified roots +- ``p.copy()`` -- Create a copy of ``p`` + +Conversion +---------- + +Methods for converting a polynomial instance of one kind to another. + +- ``p.cast(Poly)`` -- Convert ``p`` to instance of kind ``Poly`` +- ``p.convert(Poly)`` -- Convert ``p`` to instance of kind ``Poly`` or map + between ``domain`` and ``window`` + +Calculus +-------- +- ``p.deriv()`` -- Take the derivative of ``p`` +- ``p.integ()`` -- Integrate ``p`` + +Validation +---------- +- ``Poly.has_samecoef(p1, p2)`` -- Check if coefficients match +- ``Poly.has_samedomain(p1, p2)`` -- Check if domains match +- ``Poly.has_sametype(p1, p2)`` -- Check if types match +- ``Poly.has_samewindow(p1, p2)`` -- Check if windows match + +Misc +---- +- ``p.linspace()`` -- Return ``x, p(x)`` at equally-spaced points in ``domain`` +- ``p.mapparms()`` -- Return the parameters for the linear mapping between + ``domain`` and ``window``. +- ``p.roots()`` -- Return the roots of `p`. +- ``p.trim()`` -- Remove trailing coefficients. +- ``p.cutdeg(degree)`` -- Truncate p to given degree +- ``p.truncate(size)`` -- Truncate p to given size + +""" +from .polynomial import Polynomial +from .chebyshev import Chebyshev +from .legendre import Legendre +from .hermite import Hermite +from .hermite_e import HermiteE +from .laguerre import Laguerre + +__all__ = [ + "set_default_printstyle", + "polynomial", "Polynomial", + "chebyshev", "Chebyshev", + "legendre", "Legendre", + "hermite", "Hermite", + "hermite_e", "HermiteE", + "laguerre", "Laguerre", +] + + +def set_default_printstyle(style): + """ + Set the default format for the string representation of polynomials. + + Values for ``style`` must be valid inputs to ``__format__``, i.e. 'ascii' + or 'unicode'. + + Parameters + ---------- + style : str + Format string for default printing style. Must be either 'ascii' or + 'unicode'. + + Notes + ----- + The default format depends on the platform: 'unicode' is used on + Unix-based systems and 'ascii' on Windows. This determination is based on + default font support for the unicode superscript and subscript ranges. + + Examples + -------- + >>> p = np.polynomial.Polynomial([1, 2, 3]) + >>> c = np.polynomial.Chebyshev([1, 2, 3]) + >>> np.polynomial.set_default_printstyle('unicode') + >>> print(p) + 1.0 + 2.0·x + 3.0·x² + >>> print(c) + 1.0 + 2.0·T₁(x) + 3.0·T₂(x) + >>> np.polynomial.set_default_printstyle('ascii') + >>> print(p) + 1.0 + 2.0 x + 3.0 x**2 + >>> print(c) + 1.0 + 2.0 T_1(x) + 3.0 T_2(x) + >>> # Formatting supersedes all class/package-level defaults + >>> print(f"{p:unicode}") + 1.0 + 2.0·x + 3.0·x² + """ + if style not in ('unicode', 'ascii'): + raise ValueError( + f"Unsupported format string '{style}'. Valid options are 'ascii' " + f"and 'unicode'" + ) + _use_unicode = True + if style == 'ascii': + _use_unicode = False + from ._polybase import ABCPolyBase + ABCPolyBase._use_unicode = _use_unicode + + +from numpy._pytesttester import PytestTester +test = PytestTester(__name__) +del PytestTester diff --git a/lib/python3.12/site-packages/numpy/polynomial/__init__.pyi b/lib/python3.12/site-packages/numpy/polynomial/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c9d1c27a96c2d8ccfeb9e378a2599c2e70003ee4 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/__init__.pyi @@ -0,0 +1,22 @@ +from numpy._pytesttester import PytestTester + +from numpy.polynomial import ( + chebyshev as chebyshev, + hermite as hermite, + hermite_e as hermite_e, + laguerre as laguerre, + legendre as legendre, + polynomial as polynomial, +) +from numpy.polynomial.chebyshev import Chebyshev as Chebyshev +from numpy.polynomial.hermite import Hermite as Hermite +from numpy.polynomial.hermite_e import HermiteE as HermiteE +from numpy.polynomial.laguerre import Laguerre as Laguerre +from numpy.polynomial.legendre import Legendre as Legendre +from numpy.polynomial.polynomial import Polynomial as Polynomial + +__all__: list[str] +__path__: list[str] +test: PytestTester + +def set_default_printstyle(style): ... diff --git a/lib/python3.12/site-packages/numpy/polynomial/_polybase.py b/lib/python3.12/site-packages/numpy/polynomial/_polybase.py new file mode 100644 index 0000000000000000000000000000000000000000..9730574cf22e22823aaa0c77be9e630425cb2f79 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/_polybase.py @@ -0,0 +1,1206 @@ +""" +Abstract base class for the various polynomial Classes. + +The ABCPolyBase class provides the methods needed to implement the common API +for the various polynomial classes. It operates as a mixin, but uses the +abc module from the stdlib, hence it is only available for Python >= 2.6. + +""" +import os +import abc +import numbers + +import numpy as np +from . import polyutils as pu + +__all__ = ['ABCPolyBase'] + +class ABCPolyBase(abc.ABC): + """An abstract base class for immutable series classes. + + ABCPolyBase provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' along with the + methods listed below. + + .. versionadded:: 1.9.0 + + Parameters + ---------- + coef : array_like + Series coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*P_0(x) + 2*P_1(x) + 3*P_2(x)``, where + ``P_i`` is the basis polynomials of degree ``i``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is the derived class domain. + window : (2,) array_like, optional + Window, see domain for its use. The default value is the + derived class window. + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + Attributes + ---------- + coef : (N,) ndarray + Series coefficients in order of increasing degree. + domain : (2,) ndarray + Domain that is mapped to window. + window : (2,) ndarray + Window that domain is mapped to. + symbol : str + Symbol representing the independent variable. + + Class Attributes + ---------------- + maxpower : int + Maximum power allowed, i.e., the largest number ``n`` such that + ``p(x)**n`` is allowed. This is to limit runaway polynomial size. + domain : (2,) ndarray + Default domain of the class. + window : (2,) ndarray + Default window of the class. + + """ + + # Not hashable + __hash__ = None + + # Opt out of numpy ufuncs and Python ops with ndarray subclasses. + __array_ufunc__ = None + + # Limit runaway size. T_n^m has degree n*m + maxpower = 100 + + # Unicode character mappings for improved __str__ + _superscript_mapping = str.maketrans({ + "0": "⁰", + "1": "¹", + "2": "²", + "3": "³", + "4": "⁴", + "5": "⁵", + "6": "⁶", + "7": "⁷", + "8": "⁸", + "9": "⁹" + }) + _subscript_mapping = str.maketrans({ + "0": "₀", + "1": "₁", + "2": "₂", + "3": "₃", + "4": "₄", + "5": "₅", + "6": "₆", + "7": "₇", + "8": "₈", + "9": "₉" + }) + # Some fonts don't support full unicode character ranges necessary for + # the full set of superscripts and subscripts, including common/default + # fonts in Windows shells/terminals. Therefore, default to ascii-only + # printing on windows. + _use_unicode = not os.name == 'nt' + + @property + def symbol(self): + return self._symbol + + @property + @abc.abstractmethod + def domain(self): + pass + + @property + @abc.abstractmethod + def window(self): + pass + + @property + @abc.abstractmethod + def basis_name(self): + pass + + @staticmethod + @abc.abstractmethod + def _add(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _sub(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _mul(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _div(c1, c2): + pass + + @staticmethod + @abc.abstractmethod + def _pow(c, pow, maxpower=None): + pass + + @staticmethod + @abc.abstractmethod + def _val(x, c): + pass + + @staticmethod + @abc.abstractmethod + def _int(c, m, k, lbnd, scl): + pass + + @staticmethod + @abc.abstractmethod + def _der(c, m, scl): + pass + + @staticmethod + @abc.abstractmethod + def _fit(x, y, deg, rcond, full): + pass + + @staticmethod + @abc.abstractmethod + def _line(off, scl): + pass + + @staticmethod + @abc.abstractmethod + def _roots(c): + pass + + @staticmethod + @abc.abstractmethod + def _fromroots(r): + pass + + def has_samecoef(self, other): + """Check if coefficients match. + + .. versionadded:: 1.6.0 + + Parameters + ---------- + other : class instance + The other class must have the ``coef`` attribute. + + Returns + ------- + bool : boolean + True if the coefficients are the same, False otherwise. + + """ + if len(self.coef) != len(other.coef): + return False + elif not np.all(self.coef == other.coef): + return False + else: + return True + + def has_samedomain(self, other): + """Check if domains match. + + .. versionadded:: 1.6.0 + + Parameters + ---------- + other : class instance + The other class must have the ``domain`` attribute. + + Returns + ------- + bool : boolean + True if the domains are the same, False otherwise. + + """ + return np.all(self.domain == other.domain) + + def has_samewindow(self, other): + """Check if windows match. + + .. versionadded:: 1.6.0 + + Parameters + ---------- + other : class instance + The other class must have the ``window`` attribute. + + Returns + ------- + bool : boolean + True if the windows are the same, False otherwise. + + """ + return np.all(self.window == other.window) + + def has_sametype(self, other): + """Check if types match. + + .. versionadded:: 1.7.0 + + Parameters + ---------- + other : object + Class instance. + + Returns + ------- + bool : boolean + True if other is same class as self + + """ + return isinstance(other, self.__class__) + + def _get_coefficients(self, other): + """Interpret other as polynomial coefficients. + + The `other` argument is checked to see if it is of the same + class as self with identical domain and window. If so, + return its coefficients, otherwise return `other`. + + .. versionadded:: 1.9.0 + + Parameters + ---------- + other : anything + Object to be checked. + + Returns + ------- + coef + The coefficients of`other` if it is a compatible instance, + of ABCPolyBase, otherwise `other`. + + Raises + ------ + TypeError + When `other` is an incompatible instance of ABCPolyBase. + + """ + if isinstance(other, ABCPolyBase): + if not isinstance(other, self.__class__): + raise TypeError("Polynomial types differ") + elif not np.all(self.domain == other.domain): + raise TypeError("Domains differ") + elif not np.all(self.window == other.window): + raise TypeError("Windows differ") + elif self.symbol != other.symbol: + raise ValueError("Polynomial symbols differ") + return other.coef + return other + + def __init__(self, coef, domain=None, window=None, symbol='x'): + [coef] = pu.as_series([coef], trim=False) + self.coef = coef + + if domain is not None: + [domain] = pu.as_series([domain], trim=False) + if len(domain) != 2: + raise ValueError("Domain has wrong number of elements.") + self.domain = domain + + if window is not None: + [window] = pu.as_series([window], trim=False) + if len(window) != 2: + raise ValueError("Window has wrong number of elements.") + self.window = window + + # Validation for symbol + try: + if not symbol.isidentifier(): + raise ValueError( + "Symbol string must be a valid Python identifier" + ) + # If a user passes in something other than a string, the above + # results in an AttributeError. Catch this and raise a more + # informative exception + except AttributeError: + raise TypeError("Symbol must be a non-empty string") + + self._symbol = symbol + + def __repr__(self): + coef = repr(self.coef)[6:-1] + domain = repr(self.domain)[6:-1] + window = repr(self.window)[6:-1] + name = self.__class__.__name__ + return (f"{name}({coef}, domain={domain}, window={window}, " + f"symbol='{self.symbol}')") + + def __format__(self, fmt_str): + if fmt_str == '': + return self.__str__() + if fmt_str not in ('ascii', 'unicode'): + raise ValueError( + f"Unsupported format string '{fmt_str}' passed to " + f"{self.__class__}.__format__. Valid options are " + f"'ascii' and 'unicode'" + ) + if fmt_str == 'ascii': + return self._generate_string(self._str_term_ascii) + return self._generate_string(self._str_term_unicode) + + def __str__(self): + if self._use_unicode: + return self._generate_string(self._str_term_unicode) + return self._generate_string(self._str_term_ascii) + + def _generate_string(self, term_method): + """ + Generate the full string representation of the polynomial, using + ``term_method`` to generate each polynomial term. + """ + # Get configuration for line breaks + linewidth = np.get_printoptions().get('linewidth', 75) + if linewidth < 1: + linewidth = 1 + out = pu.format_float(self.coef[0]) + for i, coef in enumerate(self.coef[1:]): + out += " " + power = str(i + 1) + # Polynomial coefficient + # The coefficient array can be an object array with elements that + # will raise a TypeError with >= 0 (e.g. strings or Python + # complex). In this case, represent the coefficient as-is. + try: + if coef >= 0: + next_term = f"+ " + pu.format_float(coef, parens=True) + else: + next_term = f"- " + pu.format_float(-coef, parens=True) + except TypeError: + next_term = f"+ {coef}" + # Polynomial term + next_term += term_method(power, self.symbol) + # Length of the current line with next term added + line_len = len(out.split('\n')[-1]) + len(next_term) + # If not the last term in the polynomial, it will be two + # characters longer due to the +/- with the next term + if i < len(self.coef[1:]) - 1: + line_len += 2 + # Handle linebreaking + if line_len >= linewidth: + next_term = next_term.replace(" ", "\n", 1) + out += next_term + return out + + @classmethod + def _str_term_unicode(cls, i, arg_str): + """ + String representation of single polynomial term using unicode + characters for superscripts and subscripts. + """ + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis_name, or override " + "_str_term_unicode(cls, i, arg_str)" + ) + return (f"·{cls.basis_name}{i.translate(cls._subscript_mapping)}" + f"({arg_str})") + + @classmethod + def _str_term_ascii(cls, i, arg_str): + """ + String representation of a single polynomial term using ** and _ to + represent superscripts and subscripts, respectively. + """ + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis_name, or override " + "_str_term_ascii(cls, i, arg_str)" + ) + return f" {cls.basis_name}_{i}({arg_str})" + + @classmethod + def _repr_latex_term(cls, i, arg_str, needs_parens): + if cls.basis_name is None: + raise NotImplementedError( + "Subclasses must define either a basis name, or override " + "_repr_latex_term(i, arg_str, needs_parens)") + # since we always add parens, we don't care if the expression needs them + return f"{{{cls.basis_name}}}_{{{i}}}({arg_str})" + + @staticmethod + def _repr_latex_scalar(x, parens=False): + # TODO: we're stuck with disabling math formatting until we handle + # exponents in this function + return r'\text{{{}}}'.format(pu.format_float(x, parens=parens)) + + def _repr_latex_(self): + # get the scaled argument string to the basis functions + off, scale = self.mapparms() + if off == 0 and scale == 1: + term = self.symbol + needs_parens = False + elif scale == 1: + term = f"{self._repr_latex_scalar(off)} + {self.symbol}" + needs_parens = True + elif off == 0: + term = f"{self._repr_latex_scalar(scale)}{self.symbol}" + needs_parens = True + else: + term = ( + f"{self._repr_latex_scalar(off)} + " + f"{self._repr_latex_scalar(scale)}{self.symbol}" + ) + needs_parens = True + + mute = r"\color{{LightGray}}{{{}}}".format + + parts = [] + for i, c in enumerate(self.coef): + # prevent duplication of + and - signs + if i == 0: + coef_str = f"{self._repr_latex_scalar(c)}" + elif not isinstance(c, numbers.Real): + coef_str = f" + ({self._repr_latex_scalar(c)})" + elif not np.signbit(c): + coef_str = f" + {self._repr_latex_scalar(c, parens=True)}" + else: + coef_str = f" - {self._repr_latex_scalar(-c, parens=True)}" + + # produce the string for the term + term_str = self._repr_latex_term(i, term, needs_parens) + if term_str == '1': + part = coef_str + else: + part = rf"{coef_str}\,{term_str}" + + if c == 0: + part = mute(part) + + parts.append(part) + + if parts: + body = ''.join(parts) + else: + # in case somehow there are no coefficients at all + body = '0' + + return rf"${self.symbol} \mapsto {body}$" + + + + # Pickle and copy + + def __getstate__(self): + ret = self.__dict__.copy() + ret['coef'] = self.coef.copy() + ret['domain'] = self.domain.copy() + ret['window'] = self.window.copy() + ret['symbol'] = self.symbol + return ret + + def __setstate__(self, dict): + self.__dict__ = dict + + # Call + + def __call__(self, arg): + off, scl = pu.mapparms(self.domain, self.window) + arg = off + scl*arg + return self._val(arg, self.coef) + + def __iter__(self): + return iter(self.coef) + + def __len__(self): + return len(self.coef) + + # Numeric properties. + + def __neg__(self): + return self.__class__( + -self.coef, self.domain, self.window, self.symbol + ) + + def __pos__(self): + return self + + def __add__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._add(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __sub__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._sub(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __mul__(self, other): + othercoef = self._get_coefficients(other) + try: + coef = self._mul(self.coef, othercoef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __truediv__(self, other): + # there is no true divide if the rhs is not a Number, although it + # could return the first n elements of an infinite series. + # It is hard to see where n would come from, though. + if not isinstance(other, numbers.Number) or isinstance(other, bool): + raise TypeError( + f"unsupported types for true division: " + f"'{type(self)}', '{type(other)}'" + ) + return self.__floordiv__(other) + + def __floordiv__(self, other): + res = self.__divmod__(other) + if res is NotImplemented: + return res + return res[0] + + def __mod__(self, other): + res = self.__divmod__(other) + if res is NotImplemented: + return res + return res[1] + + def __divmod__(self, other): + othercoef = self._get_coefficients(other) + try: + quo, rem = self._div(self.coef, othercoef) + except ZeroDivisionError: + raise + except Exception: + return NotImplemented + quo = self.__class__(quo, self.domain, self.window, self.symbol) + rem = self.__class__(rem, self.domain, self.window, self.symbol) + return quo, rem + + def __pow__(self, other): + coef = self._pow(self.coef, other, maxpower=self.maxpower) + res = self.__class__(coef, self.domain, self.window, self.symbol) + return res + + def __radd__(self, other): + try: + coef = self._add(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rsub__(self, other): + try: + coef = self._sub(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rmul__(self, other): + try: + coef = self._mul(other, self.coef) + except Exception: + return NotImplemented + return self.__class__(coef, self.domain, self.window, self.symbol) + + def __rdiv__(self, other): + # set to __floordiv__ /. + return self.__rfloordiv__(other) + + def __rtruediv__(self, other): + # An instance of ABCPolyBase is not considered a + # Number. + return NotImplemented + + def __rfloordiv__(self, other): + res = self.__rdivmod__(other) + if res is NotImplemented: + return res + return res[0] + + def __rmod__(self, other): + res = self.__rdivmod__(other) + if res is NotImplemented: + return res + return res[1] + + def __rdivmod__(self, other): + try: + quo, rem = self._div(other, self.coef) + except ZeroDivisionError: + raise + except Exception: + return NotImplemented + quo = self.__class__(quo, self.domain, self.window, self.symbol) + rem = self.__class__(rem, self.domain, self.window, self.symbol) + return quo, rem + + def __eq__(self, other): + res = (isinstance(other, self.__class__) and + np.all(self.domain == other.domain) and + np.all(self.window == other.window) and + (self.coef.shape == other.coef.shape) and + np.all(self.coef == other.coef) and + (self.symbol == other.symbol)) + return res + + def __ne__(self, other): + return not self.__eq__(other) + + # + # Extra methods. + # + + def copy(self): + """Return a copy. + + Returns + ------- + new_series : series + Copy of self. + + """ + return self.__class__(self.coef, self.domain, self.window, self.symbol) + + def degree(self): + """The degree of the series. + + .. versionadded:: 1.5.0 + + Returns + ------- + degree : int + Degree of the series, one less than the number of coefficients. + + Examples + -------- + + Create a polynomial object for ``1 + 7*x + 4*x**2``: + + >>> poly = np.polynomial.Polynomial([1, 7, 4]) + >>> print(poly) + 1.0 + 7.0·x + 4.0·x² + >>> poly.degree() + 2 + + Note that this method does not check for non-zero coefficients. + You must trim the polynomial to remove any trailing zeroes: + + >>> poly = np.polynomial.Polynomial([1, 7, 0]) + >>> print(poly) + 1.0 + 7.0·x + 0.0·x² + >>> poly.degree() + 2 + >>> poly.trim().degree() + 1 + + """ + return len(self) - 1 + + def cutdeg(self, deg): + """Truncate series to the given degree. + + Reduce the degree of the series to `deg` by discarding the + high order terms. If `deg` is greater than the current degree a + copy of the current series is returned. This can be useful in least + squares where the coefficients of the high degree terms may be very + small. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + deg : non-negative int + The series is reduced to degree `deg` by discarding the high + order terms. The value of `deg` must be a non-negative integer. + + Returns + ------- + new_series : series + New instance of series with reduced degree. + + """ + return self.truncate(deg + 1) + + def trim(self, tol=0): + """Remove trailing coefficients + + Remove trailing coefficients until a coefficient is reached whose + absolute value greater than `tol` or the beginning of the series is + reached. If all the coefficients would be removed the series is set + to ``[0]``. A new series instance is returned with the new + coefficients. The current instance remains unchanged. + + Parameters + ---------- + tol : non-negative number. + All trailing coefficients less than `tol` will be removed. + + Returns + ------- + new_series : series + New instance of series with trimmed coefficients. + + """ + coef = pu.trimcoef(self.coef, tol) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def truncate(self, size): + """Truncate series to length `size`. + + Reduce the series to length `size` by discarding the high + degree terms. The value of `size` must be a positive integer. This + can be useful in least squares where the coefficients of the + high degree terms may be very small. + + Parameters + ---------- + size : positive int + The series is reduced to length `size` by discarding the high + degree terms. The value of `size` must be a positive integer. + + Returns + ------- + new_series : series + New instance of series with truncated coefficients. + + """ + isize = int(size) + if isize != size or isize < 1: + raise ValueError("size must be a positive integer") + if isize >= len(self.coef): + coef = self.coef + else: + coef = self.coef[:isize] + return self.__class__(coef, self.domain, self.window, self.symbol) + + def convert(self, domain=None, kind=None, window=None): + """Convert series to a different kind and/or domain and/or window. + + Parameters + ---------- + domain : array_like, optional + The domain of the converted series. If the value is None, + the default domain of `kind` is used. + kind : class, optional + The polynomial series type class to which the current instance + should be converted. If kind is None, then the class of the + current instance is used. + window : array_like, optional + The window of the converted series. If the value is None, + the default window of `kind` is used. + + Returns + ------- + new_series : series + The returned class can be of different type than the current + instance and/or have a different domain and/or different + window. + + Notes + ----- + Conversion between domains and class types can result in + numerically ill defined series. + + """ + if kind is None: + kind = self.__class__ + if domain is None: + domain = kind.domain + if window is None: + window = kind.window + return self(kind.identity(domain, window=window, symbol=self.symbol)) + + def mapparms(self): + """Return the mapping parameters. + + The returned values define a linear map ``off + scl*x`` that is + applied to the input arguments before the series is evaluated. The + map depends on the ``domain`` and ``window``; if the current + ``domain`` is equal to the ``window`` the resulting map is the + identity. If the coefficients of the series instance are to be + used by themselves outside this class, then the linear function + must be substituted for the ``x`` in the standard representation of + the base polynomials. + + Returns + ------- + off, scl : float or complex + The mapping function is defined by ``off + scl*x``. + + Notes + ----- + If the current domain is the interval ``[l1, r1]`` and the window + is ``[l2, r2]``, then the linear mapping function ``L`` is + defined by the equations:: + + L(l1) = l2 + L(r1) = r2 + + """ + return pu.mapparms(self.domain, self.window) + + def integ(self, m=1, k=[], lbnd=None): + """Integrate. + + Return a series instance that is the definite integral of the + current series. + + Parameters + ---------- + m : non-negative int + The number of integrations to perform. + k : array_like + Integration constants. The first constant is applied to the + first integration, the second to the second, and so on. The + list of values must less than or equal to `m` in length and any + missing values are set to zero. + lbnd : Scalar + The lower bound of the definite integral. + + Returns + ------- + new_series : series + A new series representing the integral. The domain is the same + as the domain of the integrated series. + + """ + off, scl = self.mapparms() + if lbnd is None: + lbnd = 0 + else: + lbnd = off + scl*lbnd + coef = self._int(self.coef, m, k, lbnd, 1./scl) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def deriv(self, m=1): + """Differentiate. + + Return a series instance of that is the derivative of the current + series. + + Parameters + ---------- + m : non-negative int + Find the derivative of order `m`. + + Returns + ------- + new_series : series + A new series representing the derivative. The domain is the same + as the domain of the differentiated series. + + """ + off, scl = self.mapparms() + coef = self._der(self.coef, m, scl) + return self.__class__(coef, self.domain, self.window, self.symbol) + + def roots(self): + """Return the roots of the series polynomial. + + Compute the roots for the series. Note that the accuracy of the + roots decreases the further outside the `domain` they lie. + + Returns + ------- + roots : ndarray + Array containing the roots of the series. + + """ + roots = self._roots(self.coef) + return pu.mapdomain(roots, self.window, self.domain) + + def linspace(self, n=100, domain=None): + """Return x, y values at equally spaced points in domain. + + Returns the x, y values at `n` linearly spaced points across the + domain. Here y is the value of the polynomial at the points x. By + default the domain is the same as that of the series instance. + This method is intended mostly as a plotting aid. + + .. versionadded:: 1.5.0 + + Parameters + ---------- + n : int, optional + Number of point pairs to return. The default value is 100. + domain : {None, array_like}, optional + If not None, the specified domain is used instead of that of + the calling instance. It should be of the form ``[beg,end]``. + The default is None which case the class domain is used. + + Returns + ------- + x, y : ndarray + x is equal to linspace(self.domain[0], self.domain[1], n) and + y is the series evaluated at element of x. + + """ + if domain is None: + domain = self.domain + x = np.linspace(domain[0], domain[1], n) + y = self(x) + return x, y + + @classmethod + def fit(cls, x, y, deg, domain=None, rcond=None, full=False, w=None, + window=None, symbol='x'): + """Least squares fit to data. + + Return a series instance that is the least squares fit to the data + `y` sampled at `x`. The domain of the returned instance can be + specified and this will often result in a superior fit with less + chance of ill conditioning. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) + y-coordinates of the M sample points ``(x[i], y[i])``. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + domain : {None, [beg, end], []}, optional + Domain to use for the returned series. If ``None``, + then a minimal domain that covers the points `x` is chosen. If + ``[]`` the class domain is used. The default value was the + class domain in NumPy 1.4 and ``None`` in later versions. + The ``[]`` option was added in numpy 1.5.0. + rcond : float, optional + Relative condition number of the fit. Singular values smaller + than this relative to the largest singular value will be + ignored. The default value is len(x)*eps, where eps is the + relative precision of the float type, about 2e-16 in most + cases. + full : bool, optional + Switch determining nature of return value. When it is False + (the default) just the coefficients are returned, when True + diagnostic information from the singular value decomposition is + also returned. + w : array_like, shape (M,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have + the same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + .. versionadded:: 1.5.0 + window : {[beg, end]}, optional + Window to use for the returned series. The default + value is the default class domain + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + A series that represents the least squares fit to the data and + has the domain and window specified in the call. If the + coefficients for the unscaled and unshifted basis polynomials are + of interest, do ``new_series.convert().coef``. + + [resid, rank, sv, rcond] : list + These values are only returned if ``full == True`` + + - resid -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - sv -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `linalg.lstsq`. + + """ + if domain is None: + domain = pu.getdomain(x) + elif type(domain) is list and len(domain) == 0: + domain = cls.domain + + if window is None: + window = cls.window + + xnew = pu.mapdomain(x, domain, window) + res = cls._fit(xnew, y, deg, w=w, rcond=rcond, full=full) + if full: + [coef, status] = res + return ( + cls(coef, domain=domain, window=window, symbol=symbol), status + ) + else: + coef = res + return cls(coef, domain=domain, window=window, symbol=symbol) + + @classmethod + def fromroots(cls, roots, domain=[], window=None, symbol='x'): + """Return series instance that has the specified roots. + + Returns a series representing the product + ``(x - r[0])*(x - r[1])*...*(x - r[n-1])``, where ``r`` is a + list of roots. + + Parameters + ---------- + roots : array_like + List of roots. + domain : {[], None, array_like}, optional + Domain for the resulting series. If None the domain is the + interval from the smallest root to the largest. If [] the + domain is the class domain. The default is []. + window : {None, array_like}, optional + Window for the returned series. If None the class window is + used. The default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + Series with the specified roots. + + """ + [roots] = pu.as_series([roots], trim=False) + if domain is None: + domain = pu.getdomain(roots) + elif type(domain) is list and len(domain) == 0: + domain = cls.domain + + if window is None: + window = cls.window + + deg = len(roots) + off, scl = pu.mapparms(domain, window) + rnew = off + scl*roots + coef = cls._fromroots(rnew) / scl**deg + return cls(coef, domain=domain, window=window, symbol=symbol) + + @classmethod + def identity(cls, domain=None, window=None, symbol='x'): + """Identity function. + + If ``p`` is the returned series, then ``p(x) == x`` for all + values of x. + + Parameters + ---------- + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + Series of representing the identity. + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + off, scl = pu.mapparms(window, domain) + coef = cls._line(off, scl) + return cls(coef, domain, window, symbol) + + @classmethod + def basis(cls, deg, domain=None, window=None, symbol='x'): + """Series basis polynomial of degree `deg`. + + Returns the series representing the basis polynomial of degree `deg`. + + .. versionadded:: 1.7.0 + + Parameters + ---------- + deg : int + Degree of the basis polynomial for the series. Must be >= 0. + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + symbol : str, optional + Symbol representing the independent variable. Default is 'x'. + + Returns + ------- + new_series : series + A series with the coefficient of the `deg` term set to one and + all others zero. + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + ideg = int(deg) + + if ideg != deg or ideg < 0: + raise ValueError("deg must be non-negative integer") + return cls([0]*ideg + [1], domain, window, symbol) + + @classmethod + def cast(cls, series, domain=None, window=None): + """Convert series to series of this class. + + The `series` is expected to be an instance of some polynomial + series of one of the types supported by by the numpy.polynomial + module, but could be some other class that supports the convert + method. + + .. versionadded:: 1.7.0 + + Parameters + ---------- + series : series + The series instance to be converted. + domain : {None, array_like}, optional + If given, the array must be of the form ``[beg, end]``, where + ``beg`` and ``end`` are the endpoints of the domain. If None is + given then the class domain is used. The default is None. + window : {None, array_like}, optional + If given, the resulting array must be if the form + ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of + the window. If None is given then the class window is used. The + default is None. + + Returns + ------- + new_series : series + A series of the same kind as the calling class and equal to + `series` when evaluated. + + See Also + -------- + convert : similar instance method + + """ + if domain is None: + domain = cls.domain + if window is None: + window = cls.window + return series.convert(domain, cls, window) diff --git a/lib/python3.12/site-packages/numpy/polynomial/_polybase.pyi b/lib/python3.12/site-packages/numpy/polynomial/_polybase.pyi new file mode 100644 index 0000000000000000000000000000000000000000..25c740dbedd02ca6c3f6e1beb155876a967cb57c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/_polybase.pyi @@ -0,0 +1,71 @@ +import abc +from typing import Any, ClassVar + +__all__: list[str] + +class ABCPolyBase(abc.ABC): + __hash__: ClassVar[None] # type: ignore[assignment] + __array_ufunc__: ClassVar[None] + maxpower: ClassVar[int] + coef: Any + @property + def symbol(self) -> str: ... + @property + @abc.abstractmethod + def domain(self): ... + @property + @abc.abstractmethod + def window(self): ... + @property + @abc.abstractmethod + def basis_name(self): ... + def has_samecoef(self, other): ... + def has_samedomain(self, other): ... + def has_samewindow(self, other): ... + def has_sametype(self, other): ... + def __init__(self, coef, domain=..., window=..., symbol: str = ...) -> None: ... + def __format__(self, fmt_str): ... + def __call__(self, arg): ... + def __iter__(self): ... + def __len__(self): ... + def __neg__(self): ... + def __pos__(self): ... + def __add__(self, other): ... + def __sub__(self, other): ... + def __mul__(self, other): ... + def __truediv__(self, other): ... + def __floordiv__(self, other): ... + def __mod__(self, other): ... + def __divmod__(self, other): ... + def __pow__(self, other): ... + def __radd__(self, other): ... + def __rsub__(self, other): ... + def __rmul__(self, other): ... + def __rdiv__(self, other): ... + def __rtruediv__(self, other): ... + def __rfloordiv__(self, other): ... + def __rmod__(self, other): ... + def __rdivmod__(self, other): ... + def __eq__(self, other): ... + def __ne__(self, other): ... + def copy(self): ... + def degree(self): ... + def cutdeg(self, deg): ... + def trim(self, tol=...): ... + def truncate(self, size): ... + def convert(self, domain=..., kind=..., window=...): ... + def mapparms(self): ... + def integ(self, m=..., k = ..., lbnd=...): ... + def deriv(self, m=...): ... + def roots(self): ... + def linspace(self, n=..., domain=...): ... + @classmethod + def fit(cls, x, y, deg, domain=..., rcond=..., full=..., w=..., window=...): ... + @classmethod + def fromroots(cls, roots, domain = ..., window=...): ... + @classmethod + def identity(cls, domain=..., window=...): ... + @classmethod + def basis(cls, deg, domain=..., window=...): ... + @classmethod + def cast(cls, series, domain=..., window=...): ... diff --git a/lib/python3.12/site-packages/numpy/polynomial/chebyshev.py b/lib/python3.12/site-packages/numpy/polynomial/chebyshev.py new file mode 100644 index 0000000000000000000000000000000000000000..efbe13e0cadb27e29bea430a858dea5110621a0c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/chebyshev.py @@ -0,0 +1,2082 @@ +""" +==================================================== +Chebyshev Series (:mod:`numpy.polynomial.chebyshev`) +==================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Chebyshev series, including a `Chebyshev` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- + +.. autosummary:: + :toctree: generated/ + + Chebyshev + + +Constants +--------- + +.. autosummary:: + :toctree: generated/ + + chebdomain + chebzero + chebone + chebx + +Arithmetic +---------- + +.. autosummary:: + :toctree: generated/ + + chebadd + chebsub + chebmulx + chebmul + chebdiv + chebpow + chebval + chebval2d + chebval3d + chebgrid2d + chebgrid3d + +Calculus +-------- + +.. autosummary:: + :toctree: generated/ + + chebder + chebint + +Misc Functions +-------------- + +.. autosummary:: + :toctree: generated/ + + chebfromroots + chebroots + chebvander + chebvander2d + chebvander3d + chebgauss + chebweight + chebcompanion + chebfit + chebpts1 + chebpts2 + chebtrim + chebline + cheb2poly + poly2cheb + chebinterpolate + +See also +-------- +`numpy.polynomial` + +Notes +----- +The implementations of multiplication, division, integration, and +differentiation use the algebraic identities [1]_: + +.. math:: + T_n(x) = \\frac{z^n + z^{-n}}{2} \\\\ + z\\frac{dx}{dz} = \\frac{z - z^{-1}}{2}. + +where + +.. math:: x = \\frac{z + z^{-1}}{2}. + +These identities allow a Chebyshev series to be expressed as a finite, +symmetric Laurent series. In this module, this sort of Laurent series +is referred to as a "z-series." + +References +---------- +.. [1] A. T. Benjamin, et al., "Combinatorial Trigonometry with Chebyshev + Polynomials," *Journal of Statistical Planning and Inference 14*, 2008 + (https://web.archive.org/web/20080221202153/https://www.math.hmc.edu/~benjamin/papers/CombTrig.pdf, pg. 4) + +""" +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'chebzero', 'chebone', 'chebx', 'chebdomain', 'chebline', 'chebadd', + 'chebsub', 'chebmulx', 'chebmul', 'chebdiv', 'chebpow', 'chebval', + 'chebder', 'chebint', 'cheb2poly', 'poly2cheb', 'chebfromroots', + 'chebvander', 'chebfit', 'chebtrim', 'chebroots', 'chebpts1', + 'chebpts2', 'Chebyshev', 'chebval2d', 'chebval3d', 'chebgrid2d', + 'chebgrid3d', 'chebvander2d', 'chebvander3d', 'chebcompanion', + 'chebgauss', 'chebweight', 'chebinterpolate'] + +chebtrim = pu.trimcoef + +# +# A collection of functions for manipulating z-series. These are private +# functions and do minimal error checking. +# + +def _cseries_to_zseries(c): + """Convert Chebyshev series to z-series. + + Convert a Chebyshev series to the equivalent z-series. The result is + never an empty array. The dtype of the return is the same as that of + the input. No checks are run on the arguments as this routine is for + internal use. + + Parameters + ---------- + c : 1-D ndarray + Chebyshev coefficients, ordered from low to high + + Returns + ------- + zs : 1-D ndarray + Odd length symmetric z-series, ordered from low to high. + + """ + n = c.size + zs = np.zeros(2*n-1, dtype=c.dtype) + zs[n-1:] = c/2 + return zs + zs[::-1] + + +def _zseries_to_cseries(zs): + """Convert z-series to a Chebyshev series. + + Convert a z series to the equivalent Chebyshev series. The result is + never an empty array. The dtype of the return is the same as that of + the input. No checks are run on the arguments as this routine is for + internal use. + + Parameters + ---------- + zs : 1-D ndarray + Odd length symmetric z-series, ordered from low to high. + + Returns + ------- + c : 1-D ndarray + Chebyshev coefficients, ordered from low to high. + + """ + n = (zs.size + 1)//2 + c = zs[n-1:].copy() + c[1:n] *= 2 + return c + + +def _zseries_mul(z1, z2): + """Multiply two z-series. + + Multiply two z-series to produce a z-series. + + Parameters + ---------- + z1, z2 : 1-D ndarray + The arrays must be 1-D but this is not checked. + + Returns + ------- + product : 1-D ndarray + The product z-series. + + Notes + ----- + This is simply convolution. If symmetric/anti-symmetric z-series are + denoted by S/A then the following rules apply: + + S*S, A*A -> S + S*A, A*S -> A + + """ + return np.convolve(z1, z2) + + +def _zseries_div(z1, z2): + """Divide the first z-series by the second. + + Divide `z1` by `z2` and return the quotient and remainder as z-series. + Warning: this implementation only applies when both z1 and z2 have the + same symmetry, which is sufficient for present purposes. + + Parameters + ---------- + z1, z2 : 1-D ndarray + The arrays must be 1-D and have the same symmetry, but this is not + checked. + + Returns + ------- + + (quotient, remainder) : 1-D ndarrays + Quotient and remainder as z-series. + + Notes + ----- + This is not the same as polynomial division on account of the desired form + of the remainder. If symmetric/anti-symmetric z-series are denoted by S/A + then the following rules apply: + + S/S -> S,S + A/A -> S,A + + The restriction to types of the same symmetry could be fixed but seems like + unneeded generality. There is no natural form for the remainder in the case + where there is no symmetry. + + """ + z1 = z1.copy() + z2 = z2.copy() + lc1 = len(z1) + lc2 = len(z2) + if lc2 == 1: + z1 /= z2 + return z1, z1[:1]*0 + elif lc1 < lc2: + return z1[:1]*0, z1 + else: + dlen = lc1 - lc2 + scl = z2[0] + z2 /= scl + quo = np.empty(dlen + 1, dtype=z1.dtype) + i = 0 + j = dlen + while i < j: + r = z1[i] + quo[i] = z1[i] + quo[dlen - i] = r + tmp = r*z2 + z1[i:i+lc2] -= tmp + z1[j:j+lc2] -= tmp + i += 1 + j -= 1 + r = z1[i] + quo[i] = r + tmp = r*z2 + z1[i:i+lc2] -= tmp + quo /= scl + rem = z1[i+1:i-1+lc2].copy() + return quo, rem + + +def _zseries_der(zs): + """Differentiate a z-series. + + The derivative is with respect to x, not z. This is achieved using the + chain rule and the value of dx/dz given in the module notes. + + Parameters + ---------- + zs : z-series + The z-series to differentiate. + + Returns + ------- + derivative : z-series + The derivative + + Notes + ----- + The zseries for x (ns) has been multiplied by two in order to avoid + using floats that are incompatible with Decimal and likely other + specialized scalar types. This scaling has been compensated by + multiplying the value of zs by two also so that the two cancels in the + division. + + """ + n = len(zs)//2 + ns = np.array([-1, 0, 1], dtype=zs.dtype) + zs *= np.arange(-n, n+1)*2 + d, r = _zseries_div(zs, ns) + return d + + +def _zseries_int(zs): + """Integrate a z-series. + + The integral is with respect to x, not z. This is achieved by a change + of variable using dx/dz given in the module notes. + + Parameters + ---------- + zs : z-series + The z-series to integrate + + Returns + ------- + integral : z-series + The indefinite integral + + Notes + ----- + The zseries for x (ns) has been multiplied by two in order to avoid + using floats that are incompatible with Decimal and likely other + specialized scalar types. This scaling has been compensated by + dividing the resulting zs by two. + + """ + n = 1 + len(zs)//2 + ns = np.array([-1, 0, 1], dtype=zs.dtype) + zs = _zseries_mul(zs, ns) + div = np.arange(-n, n+1)*2 + zs[:n] /= div[:n] + zs[n+1:] /= div[n+1:] + zs[n] = 0 + return zs + +# +# Chebyshev series functions +# + + +def poly2cheb(pol): + """ + Convert a polynomial to a Chebyshev series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Chebyshev series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Chebyshev + series. + + See Also + -------- + cheb2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> p = P.Polynomial(range(4)) + >>> p + Polynomial([0., 1., 2., 3.], domain=[-1, 1], window=[-1, 1]) + >>> c = p.convert(kind=P.Chebyshev) + >>> c + Chebyshev([1. , 3.25, 1. , 0.75], domain=[-1., 1.], window=[-1., 1.]) + >>> P.chebyshev.poly2cheb(range(4)) + array([1. , 3.25, 1. , 0.75]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = chebadd(chebmulx(res), pol[i]) + return res + + +def cheb2poly(c): + """ + Convert a Chebyshev series to a polynomial. + + Convert an array representing the coefficients of a Chebyshev series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Chebyshev series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2cheb + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> c = P.Chebyshev(range(4)) + >>> c + Chebyshev([0., 1., 2., 3.], domain=[-1, 1], window=[-1, 1]) + >>> p = c.convert(kind=P.Polynomial) + >>> p + Polynomial([-2., -8., 4., 12.], domain=[-1., 1.], window=[-1., 1.]) + >>> P.chebyshev.cheb2poly(range(4)) + array([-2., -8., 4., 12.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n < 3: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1) + c1 = polyadd(tmp, polymulx(c1)*2) + return polyadd(c0, polymulx(c1)) + + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Chebyshev default domain. +chebdomain = np.array([-1, 1]) + +# Chebyshev coefficients representing zero. +chebzero = np.array([0]) + +# Chebyshev coefficients representing one. +chebone = np.array([1]) + +# Chebyshev coefficients representing the identity x. +chebx = np.array([0, 1]) + + +def chebline(off, scl): + """ + Chebyshev series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Chebyshev series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebline(3,2) + array([3, 2]) + >>> C.chebval(-3, C.chebline(3,2)) # should be -3 + -3.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def chebfromroots(roots): + """ + Generate a Chebyshev series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Chebyshev form, where the `r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * T_1(x) + ... + c_n * T_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Chebyshev form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebfromroots((-1,0,1)) # x^3 - x relative to the standard basis + array([ 0. , -0.25, 0. , 0.25]) + >>> j = complex(0,1) + >>> C.chebfromroots((-j,j)) # x^2 + 1 relative to the standard basis + array([1.5+0.j, 0. +0.j, 0.5+0.j]) + + """ + return pu._fromroots(chebline, chebmul, roots) + + +def chebadd(c1, c2): + """ + Add one Chebyshev series to another. + + Returns the sum of two Chebyshev series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Chebyshev series of their sum. + + See Also + -------- + chebsub, chebmulx, chebmul, chebdiv, chebpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Chebyshev series + is a Chebyshev series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebadd(c1,c2) + array([4., 4., 4.]) + + """ + return pu._add(c1, c2) + + +def chebsub(c1, c2): + """ + Subtract one Chebyshev series from another. + + Returns the difference of two Chebyshev series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Chebyshev series coefficients representing their difference. + + See Also + -------- + chebadd, chebmulx, chebmul, chebdiv, chebpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Chebyshev + series is a Chebyshev series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebsub(c1,c2) + array([-2., 0., 2.]) + >>> C.chebsub(c2,c1) # -C.chebsub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def chebmulx(c): + """Multiply a Chebyshev series by x. + + Multiply the polynomial `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + Notes + ----- + + .. versionadded:: 1.5.0 + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> C.chebmulx([1,2,3]) + array([1. , 2.5, 1. , 1.5]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + if len(c) > 1: + tmp = c[1:]/2 + prd[2:] = tmp + prd[0:-2] += tmp + return prd + + +def chebmul(c1, c2): + """ + Multiply one Chebyshev series by another. + + Returns the product of two Chebyshev series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Chebyshev series coefficients representing their product. + + See Also + -------- + chebadd, chebsub, chebmulx, chebdiv, chebpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Chebyshev polynomial basis set. Thus, to express + the product as a C-series, it is typically necessary to "reproject" + the product onto said basis set, which typically produces + "unintuitive live" (but correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebmul(c1,c2) # multiplication requires "reprojection" + array([ 6.5, 12. , 12. , 4. , 1.5]) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + z1 = _cseries_to_zseries(c1) + z2 = _cseries_to_zseries(c2) + prd = _zseries_mul(z1, z2) + ret = _zseries_to_cseries(prd) + return pu.trimseq(ret) + + +def chebdiv(c1, c2): + """ + Divide one Chebyshev series by another. + + Returns the quotient-with-remainder of two Chebyshev series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``T_0 + 2*T_1 + 3*T_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Chebyshev series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Chebyshev series coefficients representing the quotient and + remainder. + + See Also + -------- + chebadd, chebsub, chebmulx, chebmul, chebpow + + Notes + ----- + In general, the (polynomial) division of one C-series by another + results in quotient and remainder terms that are not in the Chebyshev + polynomial basis set. Thus, to express these results as C-series, it + is typically necessary to "reproject" the results onto said basis + set, which typically produces "unintuitive" (but correct) results; + see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> C.chebdiv(c1,c2) # quotient "intuitive," remainder not + (array([3.]), array([-8., -4.])) + >>> c2 = (0,1,2,3) + >>> C.chebdiv(c2,c1) # neither "intuitive" + (array([0., 2.]), array([-2., -4.])) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError() + + # note: this is more efficient than `pu._div(chebmul, c1, c2)` + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + z1 = _cseries_to_zseries(c1) + z2 = _cseries_to_zseries(c2) + quo, rem = _zseries_div(z1, z2) + quo = pu.trimseq(_zseries_to_cseries(quo)) + rem = pu.trimseq(_zseries_to_cseries(rem)) + return quo, rem + + +def chebpow(c, pow, maxpower=16): + """Raise a Chebyshev series to a power. + + Returns the Chebyshev series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``T_0 + 2*T_1 + 3*T_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Chebyshev series of power. + + See Also + -------- + chebadd, chebsub, chebmulx, chebmul, chebdiv + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> C.chebpow([1, 2, 3, 4], 2) + array([15.5, 22. , 16. , ..., 12.5, 12. , 8. ]) + + """ + # note: this is more efficient than `pu._pow(chebmul, c1, c2)`, as it + # avoids converting between z and c series repeatedly + + # c is a trimmed copy + [c] = pu.as_series([c]) + power = int(pow) + if power != pow or power < 0: + raise ValueError("Power must be a non-negative integer.") + elif maxpower is not None and power > maxpower: + raise ValueError("Power is too large") + elif power == 0: + return np.array([1], dtype=c.dtype) + elif power == 1: + return c + else: + # This can be made more efficient by using powers of two + # in the usual way. + zs = _cseries_to_zseries(c) + prd = zs + for i in range(2, power + 1): + prd = np.convolve(prd, zs) + return _zseries_to_cseries(prd) + + +def chebder(c, m=1, scl=1, axis=0): + """ + Differentiate a Chebyshev series. + + Returns the Chebyshev series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*T_0 + 2*T_1 + 3*T_2`` + while [[1,2],[1,2]] represents ``1*T_0(x)*T_0(y) + 1*T_1(x)*T_0(y) + + 2*T_0(x)*T_1(y) + 2*T_1(x)*T_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Chebyshev series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Chebyshev series of the derivative. + + See Also + -------- + chebint + + Notes + ----- + In general, the result of differentiating a C-series needs to be + "reprojected" onto the C-series basis set. Thus, typically, the + result of this function is "unintuitive," albeit correct; see Examples + section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c = (1,2,3,4) + >>> C.chebder(c) + array([14., 12., 24.]) + >>> C.chebder(c,3) + array([96.]) + >>> C.chebder(c,scl=-1) + array([-14., -12., -24.]) + >>> C.chebder(c,2,-1) + array([12., 96.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 2, -1): + der[j - 1] = (2*j)*c[j] + c[j - 2] += (j*c[j])/(j - 2) + if n > 1: + der[1] = 4*c[2] + der[0] = c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def chebint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Chebyshev series. + + Returns the Chebyshev series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``T_0 + 2*T_1 + 3*T_2`` while [[1,2],[1,2]] + represents ``1*T_0(x)*T_0(y) + 1*T_1(x)*T_0(y) + 2*T_0(x)*T_1(y) + + 2*T_1(x)*T_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Chebyshev series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at zero + is the first value in the list, the value of the second integral + at zero is the second value, etc. If ``k == []`` (the default), + all constants are set to zero. If ``m == 1``, a single scalar can + be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + C-series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 1``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + chebder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a`- perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import chebyshev as C + >>> c = (1,2,3) + >>> C.chebint(c) + array([ 0.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,3) + array([ 0.03125 , -0.1875 , 0.04166667, -0.05208333, 0.01041667, # may vary + 0.00625 ]) + >>> C.chebint(c, k=3) + array([ 3.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,lbnd=-2) + array([ 8.5, -0.5, 0.5, 0.5]) + >>> C.chebint(c,scl=-2) + array([-1., 1., -1., -1.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + if n > 1: + tmp[2] = c[1]/4 + for j in range(2, n): + tmp[j + 1] = c[j]/(2*(j + 1)) + tmp[j - 1] -= c[j]/(2*(j - 1)) + tmp[0] += k[i] - chebval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def chebval(x, c, tensor=True): + """ + Evaluate a Chebyshev series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * T_0(x) + c_1 * T_1(x) + ... + c_n * T_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + chebval2d, chebgrid2d, chebval3d, chebgrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + x2 = 2*x + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + c0 = c[-i] - c1 + c1 = tmp + c1*x2 + return c0 + c1*x + + +def chebval2d(x, y, c): + """ + Evaluate a 2-D Chebyshev series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * T_i(x) * T_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than 2 the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points formed + from pairs of corresponding values from `x` and `y`. + + See Also + -------- + chebval, chebgrid2d, chebval3d, chebgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(chebval, c, x, y) + + +def chebgrid2d(x, y, c): + """ + Evaluate a 2-D Chebyshev series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * T_i(a) * T_j(b), + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in `c[i,j]`. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + chebval, chebval2d, chebval3d, chebgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(chebval, c, x, y) + + +def chebval3d(x, y, z, c): + """ + Evaluate a 3-D Chebyshev series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * T_i(x) * T_j(y) * T_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + chebval, chebval2d, chebgrid2d, chebgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(chebval, c, x, y, z) + + +def chebgrid3d(x, y, z, c): + """ + Evaluate a 3-D Chebyshev series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * T_i(a) * T_j(b) * T_k(c) + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + chebval, chebval2d, chebgrid2d, chebval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(chebval, c, x, y, z) + + +def chebvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = T_i(x), + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Chebyshev polynomial. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + matrix ``V = chebvander(x, n)``, then ``np.dot(V, c)`` and + ``chebval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Chebyshev series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Chebyshev polynomial. The dtype will be the same as + the converted `x`. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + # Use forward recursion to generate the entries. + v[0] = x*0 + 1 + if ideg > 0: + x2 = 2*x + v[1] = x + for i in range(2, ideg + 1): + v[i] = v[i-1]*x2 - v[i-2] + return np.moveaxis(v, 0, -1) + + +def chebvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = T_i(x) * T_j(y), + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the degrees of + the Chebyshev polynomials. + + If ``V = chebvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``chebval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Chebyshev + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + chebvander, chebvander3d, chebval2d, chebval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((chebvander, chebvander), (x, y), deg) + + +def chebvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = T_i(x)*T_j(y)*T_k(z), + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the degrees of the Chebyshev polynomials. + + If ``V = chebvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``chebval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Chebyshev + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + chebvander, chebvander3d, chebval2d, chebval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((chebvander, chebvander, chebvander), (x, y, z), deg) + + +def chebfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Chebyshev series to data. + + Return the coefficients of a Chebyshev series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * T_1(x) + ... + c_n * T_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer, + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + .. versionadded:: 1.5.0 + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Chebyshev coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + chebval : Evaluates a Chebyshev series. + chebvander : Vandermonde matrix of Chebyshev series. + chebweight : Chebyshev weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Chebyshev series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where :math:`w_j` are the weights. This problem is solved by setting up + as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `RankWarning` will be issued. This means that the + coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Chebyshev series are usually better conditioned than fits + using power series, but much can depend on the distribution of the + sample points and the smoothness of the data. If the quality of the fit + is inadequate splines may be a good alternative. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + + """ + return pu._fit(chebvander, x, y, deg, rcond, full, w) + + +def chebcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is a Chebyshev basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Chebyshev series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.array([1.] + [np.sqrt(.5)]*(n-1)) + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[0] = np.sqrt(.5) + top[1:] = 1/2 + bot[...] = top + mat[:, -1] -= (c[:-1]/c[-1])*(scl/scl[-1])*.5 + return mat + + +def chebroots(c): + """ + Compute the roots of a Chebyshev series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * T_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Chebyshev series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> import numpy.polynomial.chebyshev as cheb + >>> cheb.chebroots((-1, 1,-1, 1)) # T3 - T2 + T1 - T0 has real roots + array([ -5.00000000e-01, 2.60860684e-17, 1.00000000e+00]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = chebcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def chebinterpolate(func, deg, args=()): + """Interpolate a function at the Chebyshev points of the first kind. + + Returns the Chebyshev series that interpolates `func` at the Chebyshev + points of the first kind in the interval [-1, 1]. The interpolating + series tends to a minmax approximation to `func` with increasing `deg` + if the function is continuous in the interval. + + .. versionadded:: 1.14.0 + + Parameters + ---------- + func : function + The function to be approximated. It must be a function of a single + variable of the form ``f(x, a, b, c...)``, where ``a, b, c...`` are + extra arguments passed in the `args` parameter. + deg : int + Degree of the interpolating polynomial + args : tuple, optional + Extra arguments to be used in the function call. Default is no extra + arguments. + + Returns + ------- + coef : ndarray, shape (deg + 1,) + Chebyshev coefficients of the interpolating series ordered from low to + high. + + Examples + -------- + >>> import numpy.polynomial.chebyshev as C + >>> C.chebfromfunction(lambda x: np.tanh(x) + 0.5, 8) + array([ 5.00000000e-01, 8.11675684e-01, -9.86864911e-17, + -5.42457905e-02, -2.71387850e-16, 4.51658839e-03, + 2.46716228e-17, -3.79694221e-04, -3.26899002e-16]) + + Notes + ----- + + The Chebyshev polynomials used in the interpolation are orthogonal when + sampled at the Chebyshev points of the first kind. If it is desired to + constrain some of the coefficients they can simply be set to the desired + value after the interpolation, no new interpolation or fit is needed. This + is especially useful if it is known apriori that some of coefficients are + zero. For instance, if the function is even then the coefficients of the + terms of odd degree in the result can be set to zero. + + """ + deg = np.asarray(deg) + + # check arguments. + if deg.ndim > 0 or deg.dtype.kind not in 'iu' or deg.size == 0: + raise TypeError("deg must be an int") + if deg < 0: + raise ValueError("expected deg >= 0") + + order = deg + 1 + xcheb = chebpts1(order) + yfunc = func(xcheb, *args) + m = chebvander(xcheb, deg) + c = np.dot(m.T, yfunc) + c[0] /= order + c[1:] /= 0.5*order + + return c + + +def chebgauss(deg): + """ + Gauss-Chebyshev quadrature. + + Computes the sample points and weights for Gauss-Chebyshev quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-1, 1]` with + the weight function :math:`f(x) = 1/\\sqrt{1 - x^2}`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + + .. versionadded:: 1.7.0 + + The results have only been tested up to degree 100, higher degrees may + be problematic. For Gauss-Chebyshev there are closed form solutions for + the sample points and weights. If n = `deg`, then + + .. math:: x_i = \\cos(\\pi (2 i - 1) / (2 n)) + + .. math:: w_i = \\pi / n + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + x = np.cos(np.pi * np.arange(1, 2*ideg, 2) / (2.0*ideg)) + w = np.ones(ideg)*(np.pi/ideg) + + return x, w + + +def chebweight(x): + """ + The weight function of the Chebyshev polynomials. + + The weight function is :math:`1/\\sqrt{1 - x^2}` and the interval of + integration is :math:`[-1, 1]`. The Chebyshev polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + w = 1./(np.sqrt(1. + x) * np.sqrt(1. - x)) + return w + + +def chebpts1(npts): + """ + Chebyshev points of the first kind. + + The Chebyshev points of the first kind are the points ``cos(x)``, + where ``x = [pi*(k + .5)/npts for k in range(npts)]``. + + Parameters + ---------- + npts : int + Number of sample points desired. + + Returns + ------- + pts : ndarray + The Chebyshev points of the first kind. + + See Also + -------- + chebpts2 + + Notes + ----- + + .. versionadded:: 1.5.0 + + """ + _npts = int(npts) + if _npts != npts: + raise ValueError("npts must be integer") + if _npts < 1: + raise ValueError("npts must be >= 1") + + x = 0.5 * np.pi / _npts * np.arange(-_npts+1, _npts+1, 2) + return np.sin(x) + + +def chebpts2(npts): + """ + Chebyshev points of the second kind. + + The Chebyshev points of the second kind are the points ``cos(x)``, + where ``x = [pi*k/(npts - 1) for k in range(npts)]`` sorted in ascending + order. + + Parameters + ---------- + npts : int + Number of sample points desired. + + Returns + ------- + pts : ndarray + The Chebyshev points of the second kind. + + Notes + ----- + + .. versionadded:: 1.5.0 + + """ + _npts = int(npts) + if _npts != npts: + raise ValueError("npts must be integer") + if _npts < 2: + raise ValueError("npts must be >= 2") + + x = np.linspace(-np.pi, 0, _npts) + return np.cos(x) + + +# +# Chebyshev series class +# + +class Chebyshev(ABCPolyBase): + """A Chebyshev series class. + + The Chebyshev class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + methods listed below. + + Parameters + ---------- + coef : array_like + Chebyshev coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*T_0(x) + 2*T_1(x) + 3*T_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(chebadd) + _sub = staticmethod(chebsub) + _mul = staticmethod(chebmul) + _div = staticmethod(chebdiv) + _pow = staticmethod(chebpow) + _val = staticmethod(chebval) + _int = staticmethod(chebint) + _der = staticmethod(chebder) + _fit = staticmethod(chebfit) + _line = staticmethod(chebline) + _roots = staticmethod(chebroots) + _fromroots = staticmethod(chebfromroots) + + @classmethod + def interpolate(cls, func, deg, domain=None, args=()): + """Interpolate a function at the Chebyshev points of the first kind. + + Returns the series that interpolates `func` at the Chebyshev points of + the first kind scaled and shifted to the `domain`. The resulting series + tends to a minmax approximation of `func` when the function is + continuous in the domain. + + .. versionadded:: 1.14.0 + + Parameters + ---------- + func : function + The function to be interpolated. It must be a function of a single + variable of the form ``f(x, a, b, c...)``, where ``a, b, c...`` are + extra arguments passed in the `args` parameter. + deg : int + Degree of the interpolating polynomial. + domain : {None, [beg, end]}, optional + Domain over which `func` is interpolated. The default is None, in + which case the domain is [-1, 1]. + args : tuple, optional + Extra arguments to be used in the function call. Default is no + extra arguments. + + Returns + ------- + polynomial : Chebyshev instance + Interpolating Chebyshev instance. + + Notes + ----- + See `numpy.polynomial.chebfromfunction` for more details. + + """ + if domain is None: + domain = cls.domain + xfunc = lambda x: func(pu.mapdomain(x, cls.window, domain), *args) + coef = chebinterpolate(xfunc, deg) + return cls(coef, domain=domain) + + # Virtual properties + domain = np.array(chebdomain) + window = np.array(chebdomain) + basis_name = 'T' diff --git a/lib/python3.12/site-packages/numpy/polynomial/chebyshev.pyi b/lib/python3.12/site-packages/numpy/polynomial/chebyshev.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e8113dbae780263de1bd99ae841df16a4646d761 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/chebyshev.pyi @@ -0,0 +1,51 @@ +from typing import Any + +from numpy import ndarray, dtype, int_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +chebtrim = trimcoef + +def poly2cheb(pol): ... +def cheb2poly(c): ... + +chebdomain: ndarray[Any, dtype[int_]] +chebzero: ndarray[Any, dtype[int_]] +chebone: ndarray[Any, dtype[int_]] +chebx: ndarray[Any, dtype[int_]] + +def chebline(off, scl): ... +def chebfromroots(roots): ... +def chebadd(c1, c2): ... +def chebsub(c1, c2): ... +def chebmulx(c): ... +def chebmul(c1, c2): ... +def chebdiv(c1, c2): ... +def chebpow(c, pow, maxpower=...): ... +def chebder(c, m=..., scl=..., axis=...): ... +def chebint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ... +def chebval(x, c, tensor=...): ... +def chebval2d(x, y, c): ... +def chebgrid2d(x, y, c): ... +def chebval3d(x, y, z, c): ... +def chebgrid3d(x, y, z, c): ... +def chebvander(x, deg): ... +def chebvander2d(x, y, deg): ... +def chebvander3d(x, y, z, deg): ... +def chebfit(x, y, deg, rcond=..., full=..., w=...): ... +def chebcompanion(c): ... +def chebroots(c): ... +def chebinterpolate(func, deg, args = ...): ... +def chebgauss(deg): ... +def chebweight(x): ... +def chebpts1(npts): ... +def chebpts2(npts): ... + +class Chebyshev(ABCPolyBase): + @classmethod + def interpolate(cls, func, deg, domain=..., args = ...): ... + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/hermite.py b/lib/python3.12/site-packages/numpy/polynomial/hermite.py new file mode 100644 index 0000000000000000000000000000000000000000..210df25f5ca3ace7aaa8c7614936e305097a6195 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/hermite.py @@ -0,0 +1,1703 @@ +""" +============================================================== +Hermite Series, "Physicists" (:mod:`numpy.polynomial.hermite`) +============================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Hermite series, including a `Hermite` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Hermite + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + hermdomain + hermzero + hermone + hermx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + hermadd + hermsub + hermmulx + hermmul + hermdiv + hermpow + hermval + hermval2d + hermval3d + hermgrid2d + hermgrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + hermder + hermint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + hermfromroots + hermroots + hermvander + hermvander2d + hermvander3d + hermgauss + hermweight + hermcompanion + hermfit + hermtrim + hermline + herm2poly + poly2herm + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'hermzero', 'hermone', 'hermx', 'hermdomain', 'hermline', 'hermadd', + 'hermsub', 'hermmulx', 'hermmul', 'hermdiv', 'hermpow', 'hermval', + 'hermder', 'hermint', 'herm2poly', 'poly2herm', 'hermfromroots', + 'hermvander', 'hermfit', 'hermtrim', 'hermroots', 'Hermite', + 'hermval2d', 'hermval3d', 'hermgrid2d', 'hermgrid3d', 'hermvander2d', + 'hermvander3d', 'hermcompanion', 'hermgauss', 'hermweight'] + +hermtrim = pu.trimcoef + + +def poly2herm(pol): + """ + poly2herm(pol) + + Convert a polynomial to a Hermite series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Hermite series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Hermite + series. + + See Also + -------- + herm2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite import poly2herm + >>> poly2herm(np.arange(4)) + array([1. , 2.75 , 0.5 , 0.375]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = hermadd(hermmulx(res), pol[i]) + return res + + +def herm2poly(c): + """ + Convert a Hermite series to a polynomial. + + Convert an array representing the coefficients of a Hermite series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Hermite series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2herm + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite import herm2poly + >>> herm2poly([ 1. , 2.75 , 0.5 , 0.375]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + if n == 2: + c[1] *= 2 + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1*(2*(i - 1))) + c1 = polyadd(tmp, polymulx(c1)*2) + return polyadd(c0, polymulx(c1)*2) + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Hermite +hermdomain = np.array([-1, 1]) + +# Hermite coefficients representing zero. +hermzero = np.array([0]) + +# Hermite coefficients representing one. +hermone = np.array([1]) + +# Hermite coefficients representing the identity x. +hermx = np.array([0, 1/2]) + + +def hermline(off, scl): + """ + Hermite series whose graph is a straight line. + + + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Hermite series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial.hermite import hermline, hermval + >>> hermval(0,hermline(3, 2)) + 3.0 + >>> hermval(1,hermline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off, scl/2]) + else: + return np.array([off]) + + +def hermfromroots(roots): + """ + Generate a Hermite series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Hermite form, where the `r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * H_1(x) + ... + c_n * H_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Hermite form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> from numpy.polynomial.hermite import hermfromroots, hermval + >>> coef = hermfromroots((-1, 0, 1)) + >>> hermval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = hermfromroots((-1j, 1j)) + >>> hermval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(hermline, hermmul, roots) + + +def hermadd(c1, c2): + """ + Add one Hermite series to another. + + Returns the sum of two Hermite series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Hermite series of their sum. + + See Also + -------- + hermsub, hermmulx, hermmul, hermdiv, hermpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Hermite series + is a Hermite series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite import hermadd + >>> hermadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + """ + return pu._add(c1, c2) + + +def hermsub(c1, c2): + """ + Subtract one Hermite series from another. + + Returns the difference of two Hermite series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their difference. + + See Also + -------- + hermadd, hermmulx, hermmul, hermdiv, hermpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Hermite + series is a Hermite series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite import hermsub + >>> hermsub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def hermmulx(c): + """Multiply a Hermite series by x. + + Multiply the Hermite series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + hermadd, hermsub, hermmul, hermdiv, hermpow + + Notes + ----- + The multiplication uses the recursion relationship for Hermite + polynomials in the form + + .. math:: + + xP_i(x) = (P_{i + 1}(x)/2 + i*P_{i - 1}(x)) + + Examples + -------- + >>> from numpy.polynomial.hermite import hermmulx + >>> hermmulx([1, 2, 3]) + array([2. , 6.5, 1. , 1.5]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0]/2 + for i in range(1, len(c)): + prd[i + 1] = c[i]/2 + prd[i - 1] += c[i]*i + return prd + + +def hermmul(c1, c2): + """ + Multiply one Hermite series by another. + + Returns the product of two Hermite series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their product. + + See Also + -------- + hermadd, hermsub, hermmulx, hermdiv, hermpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Hermite polynomial basis set. Thus, to express + the product as a Hermite series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermmul + >>> hermmul([1, 2, 3], [0, 1, 2]) + array([52., 29., 52., 7., 6.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = hermsub(c[-i]*xs, c1*(2*(nd - 1))) + c1 = hermadd(tmp, hermmulx(c1)*2) + return hermadd(c0, hermmulx(c1)*2) + + +def hermdiv(c1, c2): + """ + Divide one Hermite series by another. + + Returns the quotient-with-remainder of two Hermite series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Hermite series coefficients representing the quotient and + remainder. + + See Also + -------- + hermadd, hermsub, hermmulx, hermmul, hermpow + + Notes + ----- + In general, the (polynomial) division of one Hermite series by another + results in quotient and remainder terms that are not in the Hermite + polynomial basis set. Thus, to express these results as a Hermite + series, it is necessary to "reproject" the results onto the Hermite + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermdiv + >>> hermdiv([ 52., 29., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> hermdiv([ 54., 31., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([2., 2.])) + >>> hermdiv([ 53., 30., 52., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 1.])) + + """ + return pu._div(hermmul, c1, c2) + + +def hermpow(c, pow, maxpower=16): + """Raise a Hermite series to a power. + + Returns the Hermite series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Hermite series of power. + + See Also + -------- + hermadd, hermsub, hermmulx, hermmul, hermdiv + + Examples + -------- + >>> from numpy.polynomial.hermite import hermpow + >>> hermpow([1, 2, 3], 2) + array([81., 52., 82., 12., 9.]) + + """ + return pu._pow(hermmul, c, pow, maxpower) + + +def hermder(c, m=1, scl=1, axis=0): + """ + Differentiate a Hermite series. + + Returns the Hermite series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*H_0 + 2*H_1 + 3*H_2`` + while [[1,2],[1,2]] represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + + 2*H_0(x)*H_1(y) + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite series coefficients. If `c` is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Hermite series of the derivative. + + See Also + -------- + hermint + + Notes + ----- + In general, the result of differentiating a Hermite series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermder + >>> hermder([ 1. , 0.5, 0.5, 0.5]) + array([1., 2., 3.]) + >>> hermder([-0.5, 1./2., 1./8., 1./12., 1./16.], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 0, -1): + der[j - 1] = (2*j)*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Hermite series. + + Returns the Hermite series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``H_0 + 2*H_1 + 3*H_2`` while [[1,2],[1,2]] + represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + 2*H_0(x)*H_1(y) + + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + Hermite series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + hermder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermint + >>> hermint([1,2,3]) # integrate once, value 0 at 0. + array([1. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], m=2) # integrate twice, value & deriv 0 at 0 + array([-0.5 , 0.5 , 0.125 , 0.08333333, 0.0625 ]) # may vary + >>> hermint([1,2,3], k=1) # integrate once, value 1 at 0. + array([2. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], lbnd=-1) # integrate once, value 0 at -1 + array([-2. , 0.5, 0.5, 0.5]) + >>> hermint([1,2,3], m=2, k=[1,2], lbnd=-1) + array([ 1.66666667, -0.5 , 0.125 , 0.08333333, 0.0625 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0]/2 + for j in range(1, n): + tmp[j + 1] = c[j]/(2*(j + 1)) + tmp[0] += k[i] - hermval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermval(x, c, tensor=True): + """ + Evaluate an Hermite series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * H_0(x) + c_1 * H_1(x) + ... + c_n * H_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + hermval2d, hermgrid2d, hermval3d, hermgrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermval + >>> coef = [1,2,3] + >>> hermval(1, coef) + 11.0 + >>> hermval([[1,2],[3,4]], coef) + array([[ 11., 51.], + [115., 203.]]) + + """ + c = np.array(c, ndmin=1, copy=False) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + x2 = x*2 + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - c1*(2*(nd - 1)) + c1 = tmp + c1*x2 + return c0 + c1*x2 + + +def hermval2d(x, y, c): + """ + Evaluate a 2-D Hermite series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * H_i(x) * H_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + hermval, hermgrid2d, hermval3d, hermgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(hermval, c, x, y) + + +def hermgrid2d(x, y, c): + """ + Evaluate a 2-D Hermite series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * H_i(a) * H_j(b) + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermval, hermval2d, hermval3d, hermgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(hermval, c, x, y) + + +def hermval3d(x, y, z, c): + """ + Evaluate a 3-D Hermite series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * H_i(x) * H_j(y) * H_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + hermval, hermval2d, hermgrid2d, hermgrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(hermval, c, x, y, z) + + +def hermgrid3d(x, y, z, c): + """ + Evaluate a 3-D Hermite series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * H_i(a) * H_j(b) * H_k(c) + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermval, hermval2d, hermgrid2d, hermval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(hermval, c, x, y, z) + + +def hermvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = H_i(x), + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Hermite polynomial. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + array ``V = hermvander(x, n)``, then ``np.dot(V, c)`` and + ``hermval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Hermite series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Hermite polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermvander + >>> x = np.array([-1, 0, 1]) + >>> hermvander(x, 3) + array([[ 1., -2., 2., 4.], + [ 1., 0., -2., -0.], + [ 1., 2., 2., -4.]]) + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + x2 = x*2 + v[1] = x2 + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x2 - v[i-2]*(2*(i - 1))) + return np.moveaxis(v, 0, -1) + + +def hermvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = H_i(x) * H_j(y), + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the degrees of + the Hermite polynomials. + + If ``V = hermvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``hermval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Hermite + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + hermvander, hermvander3d, hermval2d, hermval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((hermvander, hermvander), (x, y), deg) + + +def hermvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = H_i(x)*H_j(y)*H_k(z), + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the degrees of the Hermite polynomials. + + If ``V = hermvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``hermval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Hermite + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + hermvander, hermvander3d, hermval2d, hermval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((hermvander, hermvander, hermvander), (x, y, z), deg) + + +def hermfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Hermite series to data. + + Return the coefficients of a Hermite series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * H_1(x) + ... + c_n * H_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Hermite coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.polynomial.polyfit + numpy.polynomial.hermite_e.hermefit + hermval : Evaluates a Hermite series. + hermvander : Vandermonde matrix of Hermite series. + hermweight : Hermite weight function + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Hermite series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `RankWarning` will be issued. This means that the + coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Hermite series are probably most useful when the data can be + approximated by ``sqrt(w(x)) * p(x)``, where `w(x)` is the Hermite + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `hermweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> from numpy.polynomial.hermite import hermfit, hermval + >>> x = np.linspace(-10, 10) + >>> err = np.random.randn(len(x))/10 + >>> y = hermval(x, [1, 2, 3]) + err + >>> hermfit(x, y, 2) + array([1.0218, 1.9986, 2.9999]) # may vary + + """ + return pu._fit(hermvander, x, y, deg, rcond, full, w) + + +def hermcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an Hermite basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-.5*c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.hstack((1., 1./np.sqrt(2.*np.arange(n - 1, 0, -1)))) + scl = np.multiply.accumulate(scl)[::-1] + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.sqrt(.5*np.arange(1, n)) + bot[...] = top + mat[:, -1] -= scl*c[:-1]/(2.0*c[-1]) + return mat + + +def hermroots(c): + """ + Compute the roots of a Hermite series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * H_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Hermite series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.hermite import hermroots, hermfromroots + >>> coef = hermfromroots([-1, 0, 1]) + >>> coef + array([0. , 0.25 , 0. , 0.125]) + >>> hermroots(coef) + array([-1.00000000e+00, -1.38777878e-17, 1.00000000e+00]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-.5*c[0]/c[1]]) + + # rotated companion matrix reduces error + m = hermcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def _normed_hermite_n(x, n): + """ + Evaluate a normalized Hermite polynomial. + + Compute the value of the normalized Hermite polynomial of degree ``n`` + at the points ``x``. + + + Parameters + ---------- + x : ndarray of double. + Points at which to evaluate the function + n : int + Degree of the normalized Hermite function to be evaluated. + + Returns + ------- + values : ndarray + The shape of the return value is described above. + + Notes + ----- + .. versionadded:: 1.10.0 + + This function is needed for finding the Gauss points and integration + weights for high degrees. The values of the standard Hermite functions + overflow when n >= 207. + + """ + if n == 0: + return np.full(x.shape, 1/np.sqrt(np.sqrt(np.pi))) + + c0 = 0. + c1 = 1./np.sqrt(np.sqrt(np.pi)) + nd = float(n) + for i in range(n - 1): + tmp = c0 + c0 = -c1*np.sqrt((nd - 1.)/nd) + c1 = tmp + c1*x*np.sqrt(2./nd) + nd = nd - 1.0 + return c0 + c1*x*np.sqrt(2) + + +def hermgauss(deg): + """ + Gauss-Hermite quadrature. + + Computes the sample points and weights for Gauss-Hermite quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-\\inf, \\inf]` + with the weight function :math:`f(x) = \\exp(-x^2)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + + .. versionadded:: 1.7.0 + + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (H'_n(x_k) * H_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`H_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1], dtype=np.float64) + m = hermcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = _normed_hermite_n(x, ideg) + df = _normed_hermite_n(x, ideg - 1) * np.sqrt(2*ideg) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = _normed_hermite_n(x, ideg - 1) + fm /= np.abs(fm).max() + w = 1/(fm * fm) + + # for Hermite we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= np.sqrt(np.pi) / w.sum() + + return x, w + + +def hermweight(x): + """ + Weight function of the Hermite polynomials. + + The weight function is :math:`\\exp(-x^2)` and the interval of + integration is :math:`[-\\inf, \\inf]`. the Hermite polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + w = np.exp(-x**2) + return w + + +# +# Hermite series class +# + +class Hermite(ABCPolyBase): + """An Hermite series class. + + The Hermite class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed in the `ABCPolyBase` documentation. + + Parameters + ---------- + coef : array_like + Hermite coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*H_0(x) + 2*H_1(X) + 3*H_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(hermadd) + _sub = staticmethod(hermsub) + _mul = staticmethod(hermmul) + _div = staticmethod(hermdiv) + _pow = staticmethod(hermpow) + _val = staticmethod(hermval) + _int = staticmethod(hermint) + _der = staticmethod(hermder) + _fit = staticmethod(hermfit) + _line = staticmethod(hermline) + _roots = staticmethod(hermroots) + _fromroots = staticmethod(hermfromroots) + + # Virtual properties + domain = np.array(hermdomain) + window = np.array(hermdomain) + basis_name = 'H' diff --git a/lib/python3.12/site-packages/numpy/polynomial/hermite.pyi b/lib/python3.12/site-packages/numpy/polynomial/hermite.pyi new file mode 100644 index 0000000000000000000000000000000000000000..0d3556d696410689b4614138ad4cf1f6c2283a9c --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/hermite.pyi @@ -0,0 +1,46 @@ +from typing import Any + +from numpy import ndarray, dtype, int_, float_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +hermtrim = trimcoef + +def poly2herm(pol): ... +def herm2poly(c): ... + +hermdomain: ndarray[Any, dtype[int_]] +hermzero: ndarray[Any, dtype[int_]] +hermone: ndarray[Any, dtype[int_]] +hermx: ndarray[Any, dtype[float_]] + +def hermline(off, scl): ... +def hermfromroots(roots): ... +def hermadd(c1, c2): ... +def hermsub(c1, c2): ... +def hermmulx(c): ... +def hermmul(c1, c2): ... +def hermdiv(c1, c2): ... +def hermpow(c, pow, maxpower=...): ... +def hermder(c, m=..., scl=..., axis=...): ... +def hermint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ... +def hermval(x, c, tensor=...): ... +def hermval2d(x, y, c): ... +def hermgrid2d(x, y, c): ... +def hermval3d(x, y, z, c): ... +def hermgrid3d(x, y, z, c): ... +def hermvander(x, deg): ... +def hermvander2d(x, y, deg): ... +def hermvander3d(x, y, z, deg): ... +def hermfit(x, y, deg, rcond=..., full=..., w=...): ... +def hermcompanion(c): ... +def hermroots(c): ... +def hermgauss(deg): ... +def hermweight(x): ... + +class Hermite(ABCPolyBase): + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/hermite_e.py b/lib/python3.12/site-packages/numpy/polynomial/hermite_e.py new file mode 100644 index 0000000000000000000000000000000000000000..bdf29405bee7788d5ca6a8677b8402b9a7af393e --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/hermite_e.py @@ -0,0 +1,1695 @@ +""" +=================================================================== +HermiteE Series, "Probabilists" (:mod:`numpy.polynomial.hermite_e`) +=================================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Hermite_e series, including a `HermiteE` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + HermiteE + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + hermedomain + hermezero + hermeone + hermex + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + hermeadd + hermesub + hermemulx + hermemul + hermediv + hermepow + hermeval + hermeval2d + hermeval3d + hermegrid2d + hermegrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + hermeder + hermeint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + hermefromroots + hermeroots + hermevander + hermevander2d + hermevander3d + hermegauss + hermeweight + hermecompanion + hermefit + hermetrim + hermeline + herme2poly + poly2herme + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'hermezero', 'hermeone', 'hermex', 'hermedomain', 'hermeline', + 'hermeadd', 'hermesub', 'hermemulx', 'hermemul', 'hermediv', + 'hermepow', 'hermeval', 'hermeder', 'hermeint', 'herme2poly', + 'poly2herme', 'hermefromroots', 'hermevander', 'hermefit', 'hermetrim', + 'hermeroots', 'HermiteE', 'hermeval2d', 'hermeval3d', 'hermegrid2d', + 'hermegrid3d', 'hermevander2d', 'hermevander3d', 'hermecompanion', + 'hermegauss', 'hermeweight'] + +hermetrim = pu.trimcoef + + +def poly2herme(pol): + """ + poly2herme(pol) + + Convert a polynomial to a Hermite series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Hermite series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Hermite + series. + + See Also + -------- + herme2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import poly2herme + >>> poly2herme(np.arange(4)) + array([ 2., 10., 2., 3.]) + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = hermeadd(hermemulx(res), pol[i]) + return res + + +def herme2poly(c): + """ + Convert a Hermite series to a polynomial. + + Convert an array representing the coefficients of a Hermite series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Hermite series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2herme + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import herme2poly + >>> herme2poly([ 2., 10., 2., 3.]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + if n == 2: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], c1*(i - 1)) + c1 = polyadd(tmp, polymulx(c1)) + return polyadd(c0, polymulx(c1)) + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Hermite +hermedomain = np.array([-1, 1]) + +# Hermite coefficients representing zero. +hermezero = np.array([0]) + +# Hermite coefficients representing one. +hermeone = np.array([1]) + +# Hermite coefficients representing the identity x. +hermex = np.array([0, 1]) + + +def hermeline(off, scl): + """ + Hermite series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Hermite series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeline + >>> from numpy.polynomial.hermite_e import hermeline, hermeval + >>> hermeval(0,hermeline(3, 2)) + 3.0 + >>> hermeval(1,hermeline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def hermefromroots(roots): + """ + Generate a HermiteE series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in HermiteE form, where the `r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * He_1(x) + ... + c_n * He_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in HermiteE form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.chebyshev.chebfromroots + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermefromroots, hermeval + >>> coef = hermefromroots((-1, 0, 1)) + >>> hermeval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = hermefromroots((-1j, 1j)) + >>> hermeval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(hermeline, hermemul, roots) + + +def hermeadd(c1, c2): + """ + Add one Hermite series to another. + + Returns the sum of two Hermite series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Hermite series of their sum. + + See Also + -------- + hermesub, hermemulx, hermemul, hermediv, hermepow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Hermite series + is a Hermite series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeadd + >>> hermeadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + """ + return pu._add(c1, c2) + + +def hermesub(c1, c2): + """ + Subtract one Hermite series from another. + + Returns the difference of two Hermite series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their difference. + + See Also + -------- + hermeadd, hermemulx, hermemul, hermediv, hermepow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Hermite + series is a Hermite series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermesub + >>> hermesub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def hermemulx(c): + """Multiply a Hermite series by x. + + Multiply the Hermite series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + Notes + ----- + The multiplication uses the recursion relationship for Hermite + polynomials in the form + + .. math:: + + xP_i(x) = (P_{i + 1}(x) + iP_{i - 1}(x))) + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermemulx + >>> hermemulx([1, 2, 3]) + array([2., 7., 2., 3.]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + for i in range(1, len(c)): + prd[i + 1] = c[i] + prd[i - 1] += c[i]*i + return prd + + +def hermemul(c1, c2): + """ + Multiply one Hermite series by another. + + Returns the product of two Hermite series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Hermite series coefficients representing their product. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermediv, hermepow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Hermite polynomial basis set. Thus, to express + the product as a Hermite series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermemul + >>> hermemul([1, 2, 3], [0, 1, 2]) + array([14., 15., 28., 7., 6.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = hermesub(c[-i]*xs, c1*(nd - 1)) + c1 = hermeadd(tmp, hermemulx(c1)) + return hermeadd(c0, hermemulx(c1)) + + +def hermediv(c1, c2): + """ + Divide one Hermite series by another. + + Returns the quotient-with-remainder of two Hermite series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Hermite series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Hermite series coefficients representing the quotient and + remainder. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermemul, hermepow + + Notes + ----- + In general, the (polynomial) division of one Hermite series by another + results in quotient and remainder terms that are not in the Hermite + polynomial basis set. Thus, to express these results as a Hermite + series, it is necessary to "reproject" the results onto the Hermite + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermediv + >>> hermediv([ 14., 15., 28., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> hermediv([ 15., 17., 28., 7., 6.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 2.])) + + """ + return pu._div(hermemul, c1, c2) + + +def hermepow(c, pow, maxpower=16): + """Raise a Hermite series to a power. + + Returns the Hermite series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Hermite series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Hermite series of power. + + See Also + -------- + hermeadd, hermesub, hermemulx, hermemul, hermediv + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermepow + >>> hermepow([1, 2, 3], 2) + array([23., 28., 46., 12., 9.]) + + """ + return pu._pow(hermemul, c, pow, maxpower) + + +def hermeder(c, m=1, scl=1, axis=0): + """ + Differentiate a Hermite_e series. + + Returns the series coefficients `c` differentiated `m` times along + `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*He_0 + 2*He_1 + 3*He_2`` + while [[1,2],[1,2]] represents ``1*He_0(x)*He_0(y) + 1*He_1(x)*He_0(y) + + 2*He_0(x)*He_1(y) + 2*He_1(x)*He_1(y)`` if axis=0 is ``x`` and axis=1 + is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite_e series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Hermite series of the derivative. + + See Also + -------- + hermeint + + Notes + ----- + In general, the result of differentiating a Hermite series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeder + >>> hermeder([ 1., 1., 1., 1.]) + array([1., 2., 3.]) + >>> hermeder([-0.25, 1., 1./2., 1./3., 1./4 ], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + return c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 0, -1): + der[j - 1] = j*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermeint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Hermite_e series. + + Returns the Hermite_e series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``H_0 + 2*H_1 + 3*H_2`` while [[1,2],[1,2]] + represents ``1*H_0(x)*H_0(y) + 1*H_1(x)*H_0(y) + 2*H_0(x)*H_1(y) + + 2*H_1(x)*H_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Hermite_e series coefficients. If c is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + Hermite_e series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + hermeder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeint + >>> hermeint([1, 2, 3]) # integrate once, value 0 at 0. + array([1., 1., 1., 1.]) + >>> hermeint([1, 2, 3], m=2) # integrate twice, value & deriv 0 at 0 + array([-0.25 , 1. , 0.5 , 0.33333333, 0.25 ]) # may vary + >>> hermeint([1, 2, 3], k=1) # integrate once, value 1 at 0. + array([2., 1., 1., 1.]) + >>> hermeint([1, 2, 3], lbnd=-1) # integrate once, value 0 at -1 + array([-1., 1., 1., 1.]) + >>> hermeint([1, 2, 3], m=2, k=[1, 2], lbnd=-1) + array([ 1.83333333, 0. , 0.5 , 0.33333333, 0.25 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + for j in range(1, n): + tmp[j + 1] = c[j]/(j + 1) + tmp[0] += k[i] - hermeval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def hermeval(x, c, tensor=True): + """ + Evaluate an HermiteE series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * He_0(x) + c_1 * He_1(x) + ... + c_n * He_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + hermeval2d, hermegrid2d, hermeval3d, hermegrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeval + >>> coef = [1,2,3] + >>> hermeval(1, coef) + 3.0 + >>> hermeval([[1,2],[3,4]], coef) + array([[ 3., 14.], + [31., 54.]]) + + """ + c = np.array(c, ndmin=1, copy=False) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - c1*(nd - 1) + c1 = tmp + c1*x + return c0 + c1*x + + +def hermeval2d(x, y, c): + """ + Evaluate a 2-D HermiteE series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * He_i(x) * He_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + hermeval, hermegrid2d, hermeval3d, hermegrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(hermeval, c, x, y) + + +def hermegrid2d(x, y, c): + """ + Evaluate a 2-D HermiteE series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * H_i(a) * H_j(b) + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermeval, hermeval2d, hermeval3d, hermegrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(hermeval, c, x, y) + + +def hermeval3d(x, y, z, c): + """ + Evaluate a 3-D Hermite_e series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * He_i(x) * He_j(y) * He_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + hermeval, hermeval2d, hermegrid2d, hermegrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(hermeval, c, x, y, z) + + +def hermegrid3d(x, y, z, c): + """ + Evaluate a 3-D HermiteE series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * He_i(a) * He_j(b) * He_k(c) + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + hermeval, hermeval2d, hermegrid2d, hermeval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(hermeval, c, x, y, z) + + +def hermevander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = He_i(x), + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the degree of the HermiteE polynomial. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + array ``V = hermevander(x, n)``, then ``np.dot(V, c)`` and + ``hermeval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of HermiteE series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding HermiteE polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermevander + >>> x = np.array([-1, 0, 1]) + >>> hermevander(x, 3) + array([[ 1., -1., 0., 2.], + [ 1., 0., -1., -0.], + [ 1., 1., 0., -2.]]) + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x - v[i-2]*(i - 1)) + return np.moveaxis(v, 0, -1) + + +def hermevander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = He_i(x) * He_j(y), + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the degrees of + the HermiteE polynomials. + + If ``V = hermevander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``hermeval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D HermiteE + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + hermevander, hermevander3d, hermeval2d, hermeval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((hermevander, hermevander), (x, y), deg) + + +def hermevander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then Hehe pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = He_i(x)*He_j(y)*He_k(z), + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the degrees of the HermiteE polynomials. + + If ``V = hermevander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``hermeval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D HermiteE + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + hermevander, hermevander3d, hermeval2d, hermeval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((hermevander, hermevander, hermevander), (x, y, z), deg) + + +def hermefit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Hermite series to data. + + Return the coefficients of a HermiteE series of degree `deg` that is + the least squares fit to the data values `y` given at points `x`. If + `y` is 1-D the returned coefficients will also be 1-D. If `y` is 2-D + multiple fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * He_1(x) + ... + c_n * He_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Hermite coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column k of `y` are in column + `k`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full = False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.polynomial.polyfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.laguerre.lagfit + hermeval : Evaluates a Hermite series. + hermevander : pseudo Vandermonde matrix of Hermite series. + hermeweight : HermiteE weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the HermiteE series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the pseudo Vandermonde matrix of `x`, the elements of `c` + are the coefficients to be solved for, and the elements of `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `RankWarning` will be issued. This means that the + coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using HermiteE series are probably most useful when the data can + be approximated by ``sqrt(w(x)) * p(x)``, where `w(x)` is the HermiteE + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `hermeweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermefit, hermeval + >>> x = np.linspace(-10, 10) + >>> np.random.seed(123) + >>> err = np.random.randn(len(x))/10 + >>> y = hermeval(x, [1, 2, 3]) + err + >>> hermefit(x, y, 2) + array([ 1.01690445, 1.99951418, 2.99948696]) # may vary + + """ + return pu._fit(hermevander, x, y, deg, rcond, full, w) + + +def hermecompanion(c): + """ + Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an HermiteE basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of HermiteE series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = np.hstack((1., 1./np.sqrt(np.arange(n - 1, 0, -1)))) + scl = np.multiply.accumulate(scl)[::-1] + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.sqrt(np.arange(1, n)) + bot[...] = top + mat[:, -1] -= scl*c[:-1]/c[-1] + return mat + + +def hermeroots(c): + """ + Compute the roots of a HermiteE series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * He_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.chebyshev.chebroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The HermiteE series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.hermite_e import hermeroots, hermefromroots + >>> coef = hermefromroots([-1, 0, 1]) + >>> coef + array([0., 2., 0., 1.]) + >>> hermeroots(coef) + array([-1., 0., 1.]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = hermecompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def _normed_hermite_e_n(x, n): + """ + Evaluate a normalized HermiteE polynomial. + + Compute the value of the normalized HermiteE polynomial of degree ``n`` + at the points ``x``. + + + Parameters + ---------- + x : ndarray of double. + Points at which to evaluate the function + n : int + Degree of the normalized HermiteE function to be evaluated. + + Returns + ------- + values : ndarray + The shape of the return value is described above. + + Notes + ----- + .. versionadded:: 1.10.0 + + This function is needed for finding the Gauss points and integration + weights for high degrees. The values of the standard HermiteE functions + overflow when n >= 207. + + """ + if n == 0: + return np.full(x.shape, 1/np.sqrt(np.sqrt(2*np.pi))) + + c0 = 0. + c1 = 1./np.sqrt(np.sqrt(2*np.pi)) + nd = float(n) + for i in range(n - 1): + tmp = c0 + c0 = -c1*np.sqrt((nd - 1.)/nd) + c1 = tmp + c1*x*np.sqrt(1./nd) + nd = nd - 1.0 + return c0 + c1*x + + +def hermegauss(deg): + """ + Gauss-HermiteE quadrature. + + Computes the sample points and weights for Gauss-HermiteE quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-\\inf, \\inf]` + with the weight function :math:`f(x) = \\exp(-x^2/2)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + + .. versionadded:: 1.7.0 + + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (He'_n(x_k) * He_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`He_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = hermecompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = _normed_hermite_e_n(x, ideg) + df = _normed_hermite_e_n(x, ideg - 1) * np.sqrt(ideg) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = _normed_hermite_e_n(x, ideg - 1) + fm /= np.abs(fm).max() + w = 1/(fm * fm) + + # for Hermite_e we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= np.sqrt(2*np.pi) / w.sum() + + return x, w + + +def hermeweight(x): + """Weight function of the Hermite_e polynomials. + + The weight function is :math:`\\exp(-x^2/2)` and the interval of + integration is :math:`[-\\inf, \\inf]`. the HermiteE polynomials are + orthogonal, but not normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + w = np.exp(-.5*x**2) + return w + + +# +# HermiteE series class +# + +class HermiteE(ABCPolyBase): + """An HermiteE series class. + + The HermiteE class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed in the `ABCPolyBase` documentation. + + Parameters + ---------- + coef : array_like + HermiteE coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*He_0(x) + 2*He_1(X) + 3*He_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(hermeadd) + _sub = staticmethod(hermesub) + _mul = staticmethod(hermemul) + _div = staticmethod(hermediv) + _pow = staticmethod(hermepow) + _val = staticmethod(hermeval) + _int = staticmethod(hermeint) + _der = staticmethod(hermeder) + _fit = staticmethod(hermefit) + _line = staticmethod(hermeline) + _roots = staticmethod(hermeroots) + _fromroots = staticmethod(hermefromroots) + + # Virtual properties + domain = np.array(hermedomain) + window = np.array(hermedomain) + basis_name = 'He' diff --git a/lib/python3.12/site-packages/numpy/polynomial/hermite_e.pyi b/lib/python3.12/site-packages/numpy/polynomial/hermite_e.pyi new file mode 100644 index 0000000000000000000000000000000000000000..0b7152a253b654da2c069711a1bfdbd4e084cf6f --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/hermite_e.pyi @@ -0,0 +1,46 @@ +from typing import Any + +from numpy import ndarray, dtype, int_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +hermetrim = trimcoef + +def poly2herme(pol): ... +def herme2poly(c): ... + +hermedomain: ndarray[Any, dtype[int_]] +hermezero: ndarray[Any, dtype[int_]] +hermeone: ndarray[Any, dtype[int_]] +hermex: ndarray[Any, dtype[int_]] + +def hermeline(off, scl): ... +def hermefromroots(roots): ... +def hermeadd(c1, c2): ... +def hermesub(c1, c2): ... +def hermemulx(c): ... +def hermemul(c1, c2): ... +def hermediv(c1, c2): ... +def hermepow(c, pow, maxpower=...): ... +def hermeder(c, m=..., scl=..., axis=...): ... +def hermeint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ... +def hermeval(x, c, tensor=...): ... +def hermeval2d(x, y, c): ... +def hermegrid2d(x, y, c): ... +def hermeval3d(x, y, z, c): ... +def hermegrid3d(x, y, z, c): ... +def hermevander(x, deg): ... +def hermevander2d(x, y, deg): ... +def hermevander3d(x, y, z, deg): ... +def hermefit(x, y, deg, rcond=..., full=..., w=...): ... +def hermecompanion(c): ... +def hermeroots(c): ... +def hermegauss(deg): ... +def hermeweight(x): ... + +class HermiteE(ABCPolyBase): + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/laguerre.py b/lib/python3.12/site-packages/numpy/polynomial/laguerre.py new file mode 100644 index 0000000000000000000000000000000000000000..925d4898ec07673f221937fff1082711a9851df9 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/laguerre.py @@ -0,0 +1,1651 @@ +""" +================================================== +Laguerre Series (:mod:`numpy.polynomial.laguerre`) +================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Laguerre series, including a `Laguerre` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Laguerre + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + lagdomain + lagzero + lagone + lagx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + lagadd + lagsub + lagmulx + lagmul + lagdiv + lagpow + lagval + lagval2d + lagval3d + laggrid2d + laggrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + lagder + lagint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + lagfromroots + lagroots + lagvander + lagvander2d + lagvander3d + laggauss + lagweight + lagcompanion + lagfit + lagtrim + lagline + lag2poly + poly2lag + +See also +-------- +`numpy.polynomial` + +""" +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'lagzero', 'lagone', 'lagx', 'lagdomain', 'lagline', 'lagadd', + 'lagsub', 'lagmulx', 'lagmul', 'lagdiv', 'lagpow', 'lagval', 'lagder', + 'lagint', 'lag2poly', 'poly2lag', 'lagfromroots', 'lagvander', + 'lagfit', 'lagtrim', 'lagroots', 'Laguerre', 'lagval2d', 'lagval3d', + 'laggrid2d', 'laggrid3d', 'lagvander2d', 'lagvander3d', 'lagcompanion', + 'laggauss', 'lagweight'] + +lagtrim = pu.trimcoef + + +def poly2lag(pol): + """ + poly2lag(pol) + + Convert a polynomial to a Laguerre series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Laguerre series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Laguerre + series. + + See Also + -------- + lag2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.laguerre import poly2lag + >>> poly2lag(np.arange(4)) + array([ 23., -63., 58., -18.]) + + """ + [pol] = pu.as_series([pol]) + res = 0 + for p in pol[::-1]: + res = lagadd(lagmulx(res), p) + return res + + +def lag2poly(c): + """ + Convert a Laguerre series to a polynomial. + + Convert an array representing the coefficients of a Laguerre series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Laguerre series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2lag + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lag2poly + >>> lag2poly([ 23., -63., 58., -18.]) + array([0., 1., 2., 3.]) + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n == 1: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], (c1*(i - 1))/i) + c1 = polyadd(tmp, polysub((2*i - 1)*c1, polymulx(c1))/i) + return polyadd(c0, polysub(c1, polymulx(c1))) + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Laguerre +lagdomain = np.array([0, 1]) + +# Laguerre coefficients representing zero. +lagzero = np.array([0]) + +# Laguerre coefficients representing one. +lagone = np.array([1]) + +# Laguerre coefficients representing the identity x. +lagx = np.array([1, -1]) + + +def lagline(off, scl): + """ + Laguerre series whose graph is a straight line. + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Laguerre series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagline, lagval + >>> lagval(0,lagline(3, 2)) + 3.0 + >>> lagval(1,lagline(3, 2)) + 5.0 + + """ + if scl != 0: + return np.array([off + scl, -scl]) + else: + return np.array([off]) + + +def lagfromroots(roots): + """ + Generate a Laguerre series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Laguerre form, where the `r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Laguerre form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagfromroots, lagval + >>> coef = lagfromroots((-1, 0, 1)) + >>> lagval((-1, 0, 1), coef) + array([0., 0., 0.]) + >>> coef = lagfromroots((-1j, 1j)) + >>> lagval((-1j, 1j), coef) + array([0.+0.j, 0.+0.j]) + + """ + return pu._fromroots(lagline, lagmul, roots) + + +def lagadd(c1, c2): + """ + Add one Laguerre series to another. + + Returns the sum of two Laguerre series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Laguerre series of their sum. + + See Also + -------- + lagsub, lagmulx, lagmul, lagdiv, lagpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Laguerre series + is a Laguerre series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagadd + >>> lagadd([1, 2, 3], [1, 2, 3, 4]) + array([2., 4., 6., 4.]) + + + """ + return pu._add(c1, c2) + + +def lagsub(c1, c2): + """ + Subtract one Laguerre series from another. + + Returns the difference of two Laguerre series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Laguerre series coefficients representing their difference. + + See Also + -------- + lagadd, lagmulx, lagmul, lagdiv, lagpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Laguerre + series is a Laguerre series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagsub + >>> lagsub([1, 2, 3, 4], [1, 2, 3]) + array([0., 0., 0., 4.]) + + """ + return pu._sub(c1, c2) + + +def lagmulx(c): + """Multiply a Laguerre series by x. + + Multiply the Laguerre series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + lagadd, lagsub, lagmul, lagdiv, lagpow + + Notes + ----- + The multiplication uses the recursion relationship for Laguerre + polynomials in the form + + .. math:: + + xP_i(x) = (-(i + 1)*P_{i + 1}(x) + (2i + 1)P_{i}(x) - iP_{i - 1}(x)) + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagmulx + >>> lagmulx([1, 2, 3]) + array([-1., -1., 11., -9.]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0] + prd[1] = -c[0] + for i in range(1, len(c)): + prd[i + 1] = -c[i]*(i + 1) + prd[i] += c[i]*(2*i + 1) + prd[i - 1] -= c[i]*i + return prd + + +def lagmul(c1, c2): + """ + Multiply one Laguerre series by another. + + Returns the product of two Laguerre series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Laguerre series coefficients representing their product. + + See Also + -------- + lagadd, lagsub, lagmulx, lagdiv, lagpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Laguerre polynomial basis set. Thus, to express + the product as a Laguerre series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagmul + >>> lagmul([1, 2, 3], [0, 1, 2]) + array([ 8., -13., 38., -51., 36.]) + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = lagsub(c[-i]*xs, (c1*(nd - 1))/nd) + c1 = lagadd(tmp, lagsub((2*nd - 1)*c1, lagmulx(c1))/nd) + return lagadd(c0, lagsub(c1, lagmulx(c1))) + + +def lagdiv(c1, c2): + """ + Divide one Laguerre series by another. + + Returns the quotient-with-remainder of two Laguerre series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Laguerre series coefficients ordered from low to + high. + + Returns + ------- + [quo, rem] : ndarrays + Of Laguerre series coefficients representing the quotient and + remainder. + + See Also + -------- + lagadd, lagsub, lagmulx, lagmul, lagpow + + Notes + ----- + In general, the (polynomial) division of one Laguerre series by another + results in quotient and remainder terms that are not in the Laguerre + polynomial basis set. Thus, to express these results as a Laguerre + series, it is necessary to "reproject" the results onto the Laguerre + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagdiv + >>> lagdiv([ 8., -13., 38., -51., 36.], [0, 1, 2]) + (array([1., 2., 3.]), array([0.])) + >>> lagdiv([ 9., -12., 38., -51., 36.], [0, 1, 2]) + (array([1., 2., 3.]), array([1., 1.])) + + """ + return pu._div(lagmul, c1, c2) + + +def lagpow(c, pow, maxpower=16): + """Raise a Laguerre series to a power. + + Returns the Laguerre series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Laguerre series of power. + + See Also + -------- + lagadd, lagsub, lagmulx, lagmul, lagdiv + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagpow + >>> lagpow([1, 2, 3], 2) + array([ 14., -16., 56., -72., 54.]) + + """ + return pu._pow(lagmul, c, pow, maxpower) + + +def lagder(c, m=1, scl=1, axis=0): + """ + Differentiate a Laguerre series. + + Returns the Laguerre series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*L_0 + 2*L_1 + 3*L_2`` + while [[1,2],[1,2]] represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + + 2*L_0(x)*L_1(y) + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Laguerre series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Laguerre series of the derivative. + + See Also + -------- + lagint + + Notes + ----- + In general, the result of differentiating a Laguerre series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagder + >>> lagder([ 1., 1., 1., -3.]) + array([1., 2., 3.]) + >>> lagder([ 1., 0., 0., -4., 3.], m=2) + array([1., 2., 3.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 1, -1): + der[j - 1] = -c[j] + c[j - 1] += c[j] + der[0] = -c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def lagint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Laguerre series. + + Returns the Laguerre series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``L_0 + 2*L_1 + 3*L_2`` while [[1,2],[1,2]] + represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + 2*L_0(x)*L_1(y) + + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + + Parameters + ---------- + c : array_like + Array of Laguerre series coefficients. If `c` is multidimensional + the different axis correspond to different variables with the + degree in each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + Laguerre series coefficients of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + lagder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagint + >>> lagint([1,2,3]) + array([ 1., 1., 1., -3.]) + >>> lagint([1,2,3], m=2) + array([ 1., 0., 0., -4., 3.]) + >>> lagint([1,2,3], k=1) + array([ 2., 1., 1., -3.]) + >>> lagint([1,2,3], lbnd=-1) + array([11.5, 1. , 1. , -3. ]) + >>> lagint([1,2], m=2, k=[1,2], lbnd=-1) + array([ 11.16666667, -5. , -3. , 2. ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0] + tmp[1] = -c[0] + for j in range(1, n): + tmp[j] += c[j] + tmp[j + 1] = -c[j] + tmp[0] += k[i] - lagval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def lagval(x, c, tensor=True): + """ + Evaluate a Laguerre series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * L_0(x) + c_1 * L_1(x) + ... + c_n * L_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + lagval2d, laggrid2d, lagval3d, laggrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagval + >>> coef = [1,2,3] + >>> lagval(1, coef) + -0.5 + >>> lagval([[1,2],[3,4]], coef) + array([[-0.5, -4. ], + [-4.5, -2. ]]) + + """ + c = np.array(c, ndmin=1, copy=False) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - (c1*(nd - 1))/nd + c1 = tmp + (c1*((2*nd - 1) - x))/nd + return c0 + c1*(1 - x) + + +def lagval2d(x, y, c): + """ + Evaluate a 2-D Laguerre series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * L_i(x) * L_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + lagval, laggrid2d, lagval3d, laggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(lagval, c, x, y) + + +def laggrid2d(x, y, c): + """ + Evaluate a 2-D Laguerre series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * L_i(a) * L_j(b) + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in `c[i,j]`. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + lagval, lagval2d, lagval3d, laggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(lagval, c, x, y) + + +def lagval3d(x, y, z, c): + """ + Evaluate a 3-D Laguerre series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * L_i(x) * L_j(y) * L_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + lagval, lagval2d, laggrid2d, laggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(lagval, c, x, y, z) + + +def laggrid3d(x, y, z, c): + """ + Evaluate a 3-D Laguerre series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * L_i(a) * L_j(b) * L_k(c) + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + lagval, lagval2d, laggrid2d, lagval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(lagval, c, x, y, z) + + +def lagvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = L_i(x) + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Laguerre polynomial. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + array ``V = lagvander(x, n)``, then ``np.dot(V, c)`` and + ``lagval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Laguerre series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Laguerre polynomial. The dtype will be the same as + the converted `x`. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagvander + >>> x = np.array([0, 1, 2]) + >>> lagvander(x, 3) + array([[ 1. , 1. , 1. , 1. ], + [ 1. , 0. , -0.5 , -0.66666667], + [ 1. , -1. , -1. , -0.33333333]]) + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = 1 - x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*(2*i - 1 - x) - v[i-2]*(i - 1))/i + return np.moveaxis(v, 0, -1) + + +def lagvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = L_i(x) * L_j(y), + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the degrees of + the Laguerre polynomials. + + If ``V = lagvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``lagval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Laguerre + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + lagvander, lagvander3d, lagval2d, lagval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((lagvander, lagvander), (x, y), deg) + + +def lagvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = L_i(x)*L_j(y)*L_k(z), + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the degrees of the Laguerre polynomials. + + If ``V = lagvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``lagval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Laguerre + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + lagvander, lagvander3d, lagval2d, lagval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((lagvander, lagvander, lagvander), (x, y, z), deg) + + +def lagfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Laguerre series to data. + + Return the coefficients of a Laguerre series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x), + + where ``n`` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Laguerre coefficients ordered from low to high. If `y` was 2-D, + the coefficients for the data in column *k* of `y` are in column + *k*. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.legendre.legfit + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + lagval : Evaluates a Laguerre series. + lagvander : pseudo Vandermonde matrix of Laguerre series. + lagweight : Laguerre weight function. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Laguerre series ``p`` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where ``V`` is the weighted pseudo Vandermonde matrix of `x`, ``c`` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of ``V``. + + If some of the singular values of `V` are so small that they are + neglected, then a `RankWarning` will be issued. This means that the + coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Laguerre series are probably most useful when the data can + be approximated by ``sqrt(w(x)) * p(x)``, where ``w(x)`` is the Laguerre + weight. In that case the weight ``sqrt(w(x[i]))`` should be used + together with data values ``y[i]/sqrt(w(x[i]))``. The weight function is + available as `lagweight`. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagfit, lagval + >>> x = np.linspace(0, 10) + >>> err = np.random.randn(len(x))/10 + >>> y = lagval(x, [1, 2, 3]) + err + >>> lagfit(x, y, 2) + array([ 0.96971004, 2.00193749, 3.00288744]) # may vary + + """ + return pu._fit(lagvander, x, y, deg, rcond, full, w) + + +def lagcompanion(c): + """ + Return the companion matrix of c. + + The usual companion matrix of the Laguerre polynomials is already + symmetric when `c` is a basis Laguerre polynomial, so no scaling is + applied. + + Parameters + ---------- + c : array_like + 1-D array of Laguerre series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[1 + c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + top = mat.reshape(-1)[1::n+1] + mid = mat.reshape(-1)[0::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = -np.arange(1, n) + mid[...] = 2.*np.arange(n) + 1. + bot[...] = top + mat[:, -1] += (c[:-1]/c[-1])*n + return mat + + +def lagroots(c): + """ + Compute the roots of a Laguerre series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * L_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.legendre.legroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + The Laguerre series basis polynomials aren't powers of `x` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> from numpy.polynomial.laguerre import lagroots, lagfromroots + >>> coef = lagfromroots([0, 1, 2]) + >>> coef + array([ 2., -8., 12., -6.]) + >>> lagroots(coef) + array([-4.4408921e-16, 1.0000000e+00, 2.0000000e+00]) + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) <= 1: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([1 + c[0]/c[1]]) + + # rotated companion matrix reduces error + m = lagcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def laggauss(deg): + """ + Gauss-Laguerre quadrature. + + Computes the sample points and weights for Gauss-Laguerre quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[0, \\inf]` + with the weight function :math:`f(x) = \\exp(-x)`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + + .. versionadded:: 1.7.0 + + The results have only been tested up to degree 100 higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (L'_n(x_k) * L_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`L_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = lagcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = lagval(x, c) + df = lagval(x, lagder(c)) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = lagval(x, c[1:]) + fm /= np.abs(fm).max() + df /= np.abs(df).max() + w = 1/(fm * df) + + # scale w to get the right value, 1 in this case + w /= w.sum() + + return x, w + + +def lagweight(x): + """Weight function of the Laguerre polynomials. + + The weight function is :math:`exp(-x)` and the interval of integration + is :math:`[0, \\inf]`. The Laguerre polynomials are orthogonal, but not + normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + w = np.exp(-x) + return w + +# +# Laguerre series class +# + +class Laguerre(ABCPolyBase): + """A Laguerre series class. + + The Laguerre class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed in the `ABCPolyBase` documentation. + + Parameters + ---------- + coef : array_like + Laguerre coefficients in order of increasing degree, i.e, + ``(1, 2, 3)`` gives ``1*L_0(x) + 2*L_1(X) + 3*L_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [0, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [0, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(lagadd) + _sub = staticmethod(lagsub) + _mul = staticmethod(lagmul) + _div = staticmethod(lagdiv) + _pow = staticmethod(lagpow) + _val = staticmethod(lagval) + _int = staticmethod(lagint) + _der = staticmethod(lagder) + _fit = staticmethod(lagfit) + _line = staticmethod(lagline) + _roots = staticmethod(lagroots) + _fromroots = staticmethod(lagfromroots) + + # Virtual properties + domain = np.array(lagdomain) + window = np.array(lagdomain) + basis_name = 'L' diff --git a/lib/python3.12/site-packages/numpy/polynomial/laguerre.pyi b/lib/python3.12/site-packages/numpy/polynomial/laguerre.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e546bc20a54c0e522cd7ea851ad8e8a42d895980 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/laguerre.pyi @@ -0,0 +1,46 @@ +from typing import Any + +from numpy import ndarray, dtype, int_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +lagtrim = trimcoef + +def poly2lag(pol): ... +def lag2poly(c): ... + +lagdomain: ndarray[Any, dtype[int_]] +lagzero: ndarray[Any, dtype[int_]] +lagone: ndarray[Any, dtype[int_]] +lagx: ndarray[Any, dtype[int_]] + +def lagline(off, scl): ... +def lagfromroots(roots): ... +def lagadd(c1, c2): ... +def lagsub(c1, c2): ... +def lagmulx(c): ... +def lagmul(c1, c2): ... +def lagdiv(c1, c2): ... +def lagpow(c, pow, maxpower=...): ... +def lagder(c, m=..., scl=..., axis=...): ... +def lagint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ... +def lagval(x, c, tensor=...): ... +def lagval2d(x, y, c): ... +def laggrid2d(x, y, c): ... +def lagval3d(x, y, z, c): ... +def laggrid3d(x, y, z, c): ... +def lagvander(x, deg): ... +def lagvander2d(x, y, deg): ... +def lagvander3d(x, y, z, deg): ... +def lagfit(x, y, deg, rcond=..., full=..., w=...): ... +def lagcompanion(c): ... +def lagroots(c): ... +def laggauss(deg): ... +def lagweight(x): ... + +class Laguerre(ABCPolyBase): + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/legendre.py b/lib/python3.12/site-packages/numpy/polynomial/legendre.py new file mode 100644 index 0000000000000000000000000000000000000000..8e9c19d94ff60c7d314231e8bfbc1c200f12653e --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/legendre.py @@ -0,0 +1,1664 @@ +""" +================================================== +Legendre Series (:mod:`numpy.polynomial.legendre`) +================================================== + +This module provides a number of objects (mostly functions) useful for +dealing with Legendre series, including a `Legendre` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with such polynomials is in the +docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Legendre + +Constants +--------- + +.. autosummary:: + :toctree: generated/ + + legdomain + legzero + legone + legx + +Arithmetic +---------- + +.. autosummary:: + :toctree: generated/ + + legadd + legsub + legmulx + legmul + legdiv + legpow + legval + legval2d + legval3d + leggrid2d + leggrid3d + +Calculus +-------- + +.. autosummary:: + :toctree: generated/ + + legder + legint + +Misc Functions +-------------- + +.. autosummary:: + :toctree: generated/ + + legfromroots + legroots + legvander + legvander2d + legvander3d + leggauss + legweight + legcompanion + legfit + legtrim + legline + leg2poly + poly2leg + +See also +-------- +numpy.polynomial + +""" +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +__all__ = [ + 'legzero', 'legone', 'legx', 'legdomain', 'legline', 'legadd', + 'legsub', 'legmulx', 'legmul', 'legdiv', 'legpow', 'legval', 'legder', + 'legint', 'leg2poly', 'poly2leg', 'legfromroots', 'legvander', + 'legfit', 'legtrim', 'legroots', 'Legendre', 'legval2d', 'legval3d', + 'leggrid2d', 'leggrid3d', 'legvander2d', 'legvander3d', 'legcompanion', + 'leggauss', 'legweight'] + +legtrim = pu.trimcoef + + +def poly2leg(pol): + """ + Convert a polynomial to a Legendre series. + + Convert an array representing the coefficients of a polynomial (relative + to the "standard" basis) ordered from lowest degree to highest, to an + array of the coefficients of the equivalent Legendre series, ordered + from lowest to highest degree. + + Parameters + ---------- + pol : array_like + 1-D array containing the polynomial coefficients + + Returns + ------- + c : ndarray + 1-D array containing the coefficients of the equivalent Legendre + series. + + See Also + -------- + leg2poly + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> p = P.Polynomial(np.arange(4)) + >>> p + Polynomial([0., 1., 2., 3.], domain=[-1, 1], window=[-1, 1]) + >>> c = P.Legendre(P.legendre.poly2leg(p.coef)) + >>> c + Legendre([ 1. , 3.25, 1. , 0.75], domain=[-1, 1], window=[-1, 1]) # may vary + + """ + [pol] = pu.as_series([pol]) + deg = len(pol) - 1 + res = 0 + for i in range(deg, -1, -1): + res = legadd(legmulx(res), pol[i]) + return res + + +def leg2poly(c): + """ + Convert a Legendre series to a polynomial. + + Convert an array representing the coefficients of a Legendre series, + ordered from lowest degree to highest, to an array of the coefficients + of the equivalent polynomial (relative to the "standard" basis) ordered + from lowest to highest degree. + + Parameters + ---------- + c : array_like + 1-D array containing the Legendre series coefficients, ordered + from lowest order term to highest. + + Returns + ------- + pol : ndarray + 1-D array containing the coefficients of the equivalent polynomial + (relative to the "standard" basis) ordered from lowest order term + to highest. + + See Also + -------- + poly2leg + + Notes + ----- + The easy way to do conversions between polynomial basis sets + is to use the convert method of a class instance. + + Examples + -------- + >>> from numpy import polynomial as P + >>> c = P.Legendre(range(4)) + >>> c + Legendre([0., 1., 2., 3.], domain=[-1, 1], window=[-1, 1]) + >>> p = c.convert(kind=P.Polynomial) + >>> p + Polynomial([-1. , -3.5, 3. , 7.5], domain=[-1., 1.], window=[-1., 1.]) + >>> P.legendre.leg2poly(range(4)) + array([-1. , -3.5, 3. , 7.5]) + + + """ + from .polynomial import polyadd, polysub, polymulx + + [c] = pu.as_series([c]) + n = len(c) + if n < 3: + return c + else: + c0 = c[-2] + c1 = c[-1] + # i is the current degree of c1 + for i in range(n - 1, 1, -1): + tmp = c0 + c0 = polysub(c[i - 2], (c1*(i - 1))/i) + c1 = polyadd(tmp, (polymulx(c1)*(2*i - 1))/i) + return polyadd(c0, polymulx(c1)) + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Legendre +legdomain = np.array([-1, 1]) + +# Legendre coefficients representing zero. +legzero = np.array([0]) + +# Legendre coefficients representing one. +legone = np.array([1]) + +# Legendre coefficients representing the identity x. +legx = np.array([0, 1]) + + +def legline(off, scl): + """ + Legendre series whose graph is a straight line. + + + + Parameters + ---------- + off, scl : scalars + The specified line is given by ``off + scl*x``. + + Returns + ------- + y : ndarray + This module's representation of the Legendre series for + ``off + scl*x``. + + See Also + -------- + numpy.polynomial.polynomial.polyline + numpy.polynomial.chebyshev.chebline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> import numpy.polynomial.legendre as L + >>> L.legline(3,2) + array([3, 2]) + >>> L.legval(-3, L.legline(3,2)) # should be -3 + -3.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def legfromroots(roots): + """ + Generate a Legendre series with given roots. + + The function returns the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + in Legendre form, where the `r_n` are the roots specified in `roots`. + If a zero has multiplicity n, then it must appear in `roots` n times. + For instance, if 2 is a root of multiplicity three and 3 is a root of + multiplicity 2, then `roots` looks something like [2, 2, 2, 3, 3]. The + roots can appear in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x) + + The coefficient of the last term is not generally 1 for monic + polynomials in Legendre form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of coefficients. If all roots are real then `out` is a + real array, if some of the roots are complex, then `out` is complex + even if all the coefficients in the result are real (see Examples + below). + + See Also + -------- + numpy.polynomial.polynomial.polyfromroots + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Examples + -------- + >>> import numpy.polynomial.legendre as L + >>> L.legfromroots((-1,0,1)) # x^3 - x relative to the standard basis + array([ 0. , -0.4, 0. , 0.4]) + >>> j = complex(0,1) + >>> L.legfromroots((-j,j)) # x^2 + 1 relative to the standard basis + array([ 1.33333333+0.j, 0.00000000+0.j, 0.66666667+0.j]) # may vary + + """ + return pu._fromroots(legline, legmul, roots) + + +def legadd(c1, c2): + """ + Add one Legendre series to another. + + Returns the sum of two Legendre series `c1` + `c2`. The arguments + are sequences of coefficients ordered from lowest order term to + highest, i.e., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the Legendre series of their sum. + + See Also + -------- + legsub, legmulx, legmul, legdiv, legpow + + Notes + ----- + Unlike multiplication, division, etc., the sum of two Legendre series + is a Legendre series (without having to "reproject" the result onto + the basis set) so addition, just like that of "standard" polynomials, + is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legadd(c1,c2) + array([4., 4., 4.]) + + """ + return pu._add(c1, c2) + + +def legsub(c1, c2): + """ + Subtract one Legendre series from another. + + Returns the difference of two Legendre series `c1` - `c2`. The + sequences of coefficients are from lowest order term to highest, i.e., + [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Legendre series coefficients representing their difference. + + See Also + -------- + legadd, legmulx, legmul, legdiv, legpow + + Notes + ----- + Unlike multiplication, division, etc., the difference of two Legendre + series is a Legendre series (without having to "reproject" the result + onto the basis set) so subtraction, just like that of "standard" + polynomials, is simply "component-wise." + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legsub(c1,c2) + array([-2., 0., 2.]) + >>> L.legsub(c2,c1) # -C.legsub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def legmulx(c): + """Multiply a Legendre series by x. + + Multiply the Legendre series `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + legadd, legmul, legdiv, legpow + + Notes + ----- + The multiplication uses the recursion relationship for Legendre + polynomials in the form + + .. math:: + + xP_i(x) = ((i + 1)*P_{i + 1}(x) + i*P_{i - 1}(x))/(2i + 1) + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> L.legmulx([1,2,3]) + array([ 0.66666667, 2.2, 1.33333333, 1.8]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1] = c[0] + for i in range(1, len(c)): + j = i + 1 + k = i - 1 + s = i + j + prd[j] = (c[i]*j)/s + prd[k] += (c[i]*i)/s + return prd + + +def legmul(c1, c2): + """ + Multiply one Legendre series by another. + + Returns the product of two Legendre series `c1` * `c2`. The arguments + are sequences of coefficients, from lowest order "term" to highest, + e.g., [1,2,3] represents the series ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of Legendre series coefficients representing their product. + + See Also + -------- + legadd, legsub, legmulx, legdiv, legpow + + Notes + ----- + In general, the (polynomial) product of two C-series results in terms + that are not in the Legendre polynomial basis set. Thus, to express + the product as a Legendre series, it is necessary to "reproject" the + product onto said basis set, which may produce "unintuitive" (but + correct) results; see Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2) + >>> L.legmul(c1,c2) # multiplication requires "reprojection" + array([ 4.33333333, 10.4 , 11.66666667, 3.6 ]) # may vary + + """ + # s1, s2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + + if len(c1) > len(c2): + c = c2 + xs = c1 + else: + c = c1 + xs = c2 + + if len(c) == 1: + c0 = c[0]*xs + c1 = 0 + elif len(c) == 2: + c0 = c[0]*xs + c1 = c[1]*xs + else: + nd = len(c) + c0 = c[-2]*xs + c1 = c[-1]*xs + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = legsub(c[-i]*xs, (c1*(nd - 1))/nd) + c1 = legadd(tmp, (legmulx(c1)*(2*nd - 1))/nd) + return legadd(c0, legmulx(c1)) + + +def legdiv(c1, c2): + """ + Divide one Legendre series by another. + + Returns the quotient-with-remainder of two Legendre series + `c1` / `c2`. The arguments are sequences of coefficients from lowest + order "term" to highest, e.g., [1,2,3] represents the series + ``P_0 + 2*P_1 + 3*P_2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of Legendre series coefficients ordered from low to + high. + + Returns + ------- + quo, rem : ndarrays + Of Legendre series coefficients representing the quotient and + remainder. + + See Also + -------- + legadd, legsub, legmulx, legmul, legpow + + Notes + ----- + In general, the (polynomial) division of one Legendre series by another + results in quotient and remainder terms that are not in the Legendre + polynomial basis set. Thus, to express these results as a Legendre + series, it is necessary to "reproject" the results onto the Legendre + basis set, which may produce "unintuitive" (but correct) results; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> L.legdiv(c1,c2) # quotient "intuitive," remainder not + (array([3.]), array([-8., -4.])) + >>> c2 = (0,1,2,3) + >>> L.legdiv(c2,c1) # neither "intuitive" + (array([-0.07407407, 1.66666667]), array([-1.03703704, -2.51851852])) # may vary + + """ + return pu._div(legmul, c1, c2) + + +def legpow(c, pow, maxpower=16): + """Raise a Legendre series to a power. + + Returns the Legendre series `c` raised to the power `pow`. The + argument `c` is a sequence of coefficients ordered from low to high. + i.e., [1,2,3] is the series ``P_0 + 2*P_1 + 3*P_2.`` + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to + high. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Legendre series of power. + + See Also + -------- + legadd, legsub, legmulx, legmul, legdiv + + """ + return pu._pow(legmul, c, pow, maxpower) + + +def legder(c, m=1, scl=1, axis=0): + """ + Differentiate a Legendre series. + + Returns the Legendre series coefficients `c` differentiated `m` times + along `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The argument + `c` is an array of coefficients from low to high degree along each + axis, e.g., [1,2,3] represents the series ``1*L_0 + 2*L_1 + 3*L_2`` + while [[1,2],[1,2]] represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + + 2*L_0(x)*L_1(y) + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + Array of Legendre series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change of + variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Legendre series of the derivative. + + See Also + -------- + legint + + Notes + ----- + In general, the result of differentiating a Legendre series does not + resemble the same operation on a power series. Thus the result of this + function may be "unintuitive," albeit correct; see Examples section + below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c = (1,2,3,4) + >>> L.legder(c) + array([ 6., 9., 20.]) + >>> L.legder(c, 3) + array([60.]) + >>> L.legder(c, scl=-1) + array([ -6., -9., -20.]) + >>> L.legder(c, 2,-1) + array([ 9., 60.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=c.dtype) + for j in range(n, 2, -1): + der[j - 1] = (2*j - 1)*c[j] + c[j - 2] += c[j] + if n > 1: + der[1] = 3*c[2] + der[0] = c[1] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def legint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a Legendre series. + + Returns the Legendre series coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients from low to high degree along each axis, e.g., [1,2,3] + represents the series ``L_0 + 2*L_1 + 3*L_2`` while [[1,2],[1,2]] + represents ``1*L_0(x)*L_0(y) + 1*L_1(x)*L_0(y) + 2*L_0(x)*L_1(y) + + 2*L_1(x)*L_1(y)`` if axis=0 is ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of Legendre series coefficients. If c is multidimensional the + different axis correspond to different variables with the degree in + each axis given by the corresponding index. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at + ``lbnd`` is the first value in the list, the value of the second + integral at ``lbnd`` is the second value, etc. If ``k == []`` (the + default), all constants are set to zero. If ``m == 1``, a single + scalar can be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + Legendre series coefficient array of the integral. + + Raises + ------ + ValueError + If ``m < 0``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + legder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. + Why is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Also note that, in general, the result of integrating a C-series needs + to be "reprojected" onto the C-series basis set. Thus, typically, + the result of this function is "unintuitive," albeit correct; see + Examples section below. + + Examples + -------- + >>> from numpy.polynomial import legendre as L + >>> c = (1,2,3) + >>> L.legint(c) + array([ 0.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, 3) + array([ 1.66666667e-02, -1.78571429e-02, 4.76190476e-02, # may vary + -1.73472348e-18, 1.90476190e-02, 9.52380952e-03]) + >>> L.legint(c, k=3) + array([ 3.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, lbnd=-2) + array([ 7.33333333, 0.4 , 0.66666667, 0.6 ]) # may vary + >>> L.legint(c, scl=2) + array([ 0.66666667, 0.8 , 1.33333333, 1.2 ]) # may vary + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + k = list(k) + [0]*(cnt - len(k)) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=c.dtype) + tmp[0] = c[0]*0 + tmp[1] = c[0] + if n > 1: + tmp[2] = c[1]/3 + for j in range(2, n): + t = c[j]/(2*j + 1) + tmp[j + 1] = t + tmp[j - 1] -= t + tmp[0] += k[i] - legval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def legval(x, c, tensor=True): + """ + Evaluate a Legendre series at points x. + + If `c` is of length `n + 1`, this function returns the value: + + .. math:: p(x) = c_0 * L_0(x) + c_1 * L_1(x) + ... + c_n * L_n(x) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, algebra_like + The shape of the return value is described above. + + See Also + -------- + legval2d, leggrid2d, legval3d, leggrid3d + + Notes + ----- + The evaluation uses Clenshaw recursion, aka synthetic division. + + """ + c = np.array(c, ndmin=1, copy=False) + if c.dtype.char in '?bBhHiIlLqQpP': + c = c.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + if len(c) == 1: + c0 = c[0] + c1 = 0 + elif len(c) == 2: + c0 = c[0] + c1 = c[1] + else: + nd = len(c) + c0 = c[-2] + c1 = c[-1] + for i in range(3, len(c) + 1): + tmp = c0 + nd = nd - 1 + c0 = c[-i] - (c1*(nd - 1))/nd + c1 = tmp + (c1*x*(2*nd - 1))/nd + return c0 + c1*x + + +def legval2d(x, y, c): + """ + Evaluate a 2-D Legendre series at points (x, y). + + This function returns the values: + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * L_i(x) * L_j(y) + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` is a 1-D array a one is implicitly appended to its shape to make + it 2-D. The shape of the result will be c.shape[2:] + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and if it isn't an ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in ``c[i,j]``. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Legendre series at points formed + from pairs of corresponding values from `x` and `y`. + + See Also + -------- + legval, leggrid2d, legval3d, leggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(legval, c, x, y) + + +def leggrid2d(x, y, c): + """ + Evaluate a 2-D Legendre series on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * L_i(a) * L_j(b) + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j is contained in `c[i,j]`. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional Chebyshev series at points in the + Cartesian product of `x` and `y`. + + See Also + -------- + legval, legval2d, legval3d, leggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(legval, c, x, y) + + +def legval3d(x, y, z, c): + """ + Evaluate a 3-D Legendre series at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * L_i(x) * L_j(y) * L_k(z) + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + legval, legval2d, leggrid2d, leggrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(legval, c, x, y, z) + + +def leggrid3d(x, y, z, c): + """ + Evaluate a 3-D Legendre series on the Cartesian product of x, y, and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * L_i(a) * L_j(b) * L_k(c) + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + legval, legval2d, leggrid2d, legval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(legval, c, x, y, z) + + +def legvander(x, deg): + """Pseudo-Vandermonde matrix of given degree. + + Returns the pseudo-Vandermonde matrix of degree `deg` and sample points + `x`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., i] = L_i(x) + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the degree of the Legendre polynomial. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + array ``V = legvander(x, n)``, then ``np.dot(V, c)`` and + ``legval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of Legendre series of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray + The pseudo-Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where The last index is the degree of the + corresponding Legendre polynomial. The dtype will be the same as + the converted `x`. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + # Use forward recursion to generate the entries. This is not as accurate + # as reverse recursion in this application but it is more efficient. + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = (v[i-1]*x*(2*i - 1) - v[i-2]*(i - 1))/i + return np.moveaxis(v, 0, -1) + + +def legvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = L_i(x) * L_j(y), + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the degrees of + the Legendre polynomials. + + If ``V = legvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``legval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D Legendre + series of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + legvander, legvander3d, legval2d, legval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((legvander, legvander), (x, y), deg) + + +def legvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = L_i(x)*L_j(y)*L_k(z), + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the degrees of the Legendre polynomials. + + If ``V = legvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``legval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D Legendre + series of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg[1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + legvander, legvander3d, legval2d, legval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((legvander, legvander, legvander), (x, y, z), deg) + + +def legfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least squares fit of Legendre series to data. + + Return the coefficients of a Legendre series of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * L_1(x) + ... + c_n * L_n(x), + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (M,) + x-coordinates of the M sample points ``(x[i], y[i])``. + y : array_like, shape (M,) or (M, K) + y-coordinates of the sample points. Several data sets of sample + points sharing the same x-coordinates can be fitted at once by + passing in a 2D-array that contains one dataset per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller than + this relative to the largest singular value will be ignored. The + default value is len(x)*eps, where eps is the relative precision of + the float type, about 2e-16 in most cases. + full : bool, optional + Switch determining nature of return value. When it is False (the + default) just the coefficients are returned, when True diagnostic + information from the singular value decomposition is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + .. versionadded:: 1.5.0 + + Returns + ------- + coef : ndarray, shape (M,) or (M, K) + Legendre coefficients ordered from low to high. If `y` was + 2-D, the coefficients for the data in column k of `y` are in + column `k`. If `deg` is specified as a list, coefficients for + terms not included in the fit are set equal to zero in the + returned `coef`. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Warns + ----- + RankWarning + The rank of the coefficient matrix in the least-squares fit is + deficient. The warning is only raised if ``full == False``. The + warnings can be turned off by + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.polynomial.polyfit + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + legval : Evaluates a Legendre series. + legvander : Vandermonde matrix of Legendre series. + legweight : Legendre weight function (= 1). + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the Legendre series `p` that + minimizes the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where :math:`w_j` are the weights. This problem is solved by setting up + as the (typically) overdetermined matrix equation + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected, then a `RankWarning` will be issued. This means that the + coefficient values may be poorly determined. Using a lower order fit + will usually get rid of the warning. The `rcond` parameter can also be + set to a value smaller than its default, but the resulting fit may be + spurious and have large contributions from roundoff error. + + Fits using Legendre series are usually better conditioned than fits + using power series, but much can depend on the distribution of the + sample points and the smoothness of the data. If the quality of the fit + is inadequate splines may be a good alternative. + + References + ---------- + .. [1] Wikipedia, "Curve fitting", + https://en.wikipedia.org/wiki/Curve_fitting + + Examples + -------- + + """ + return pu._fit(legvander, x, y, deg, rcond, full, w) + + +def legcompanion(c): + """Return the scaled companion matrix of c. + + The basis polynomials are scaled so that the companion matrix is + symmetric when `c` is an Legendre basis polynomial. This provides + better eigenvalue estimates than the unscaled case and for basis + polynomials the eigenvalues are guaranteed to be real if + `numpy.linalg.eigvalsh` is used to obtain them. + + Parameters + ---------- + c : array_like + 1-D array of Legendre series coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Scaled companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + scl = 1./np.sqrt(2*np.arange(n) + 1) + top = mat.reshape(-1)[1::n+1] + bot = mat.reshape(-1)[n::n+1] + top[...] = np.arange(1, n)*scl[:n-1]*scl[1:n] + bot[...] = top + mat[:, -1] -= (c[:-1]/c[-1])*(scl/scl[-1])*(n/(2*n - 1)) + return mat + + +def legroots(c): + """ + Compute the roots of a Legendre series. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * L_i(x). + + Parameters + ---------- + c : 1-D array_like + 1-D array of coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the series. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.polynomial.polyroots + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the series for such values. + Roots with multiplicity greater than 1 will also show larger errors as + the value of the series near such points is relatively insensitive to + errors in the roots. Isolated roots near the origin can be improved by + a few iterations of Newton's method. + + The Legendre series basis polynomials aren't powers of ``x`` so the + results of this function may seem unintuitive. + + Examples + -------- + >>> import numpy.polynomial.legendre as leg + >>> leg.legroots((1, 2, 3, 4)) # 4L_3 + 3L_2 + 2L_1 + 1L_0, all real roots + array([-0.85099543, -0.11407192, 0.51506735]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = legcompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +def leggauss(deg): + """ + Gauss-Legendre quadrature. + + Computes the sample points and weights for Gauss-Legendre quadrature. + These sample points and weights will correctly integrate polynomials of + degree :math:`2*deg - 1` or less over the interval :math:`[-1, 1]` with + the weight function :math:`f(x) = 1`. + + Parameters + ---------- + deg : int + Number of sample points and weights. It must be >= 1. + + Returns + ------- + x : ndarray + 1-D ndarray containing the sample points. + y : ndarray + 1-D ndarray containing the weights. + + Notes + ----- + + .. versionadded:: 1.7.0 + + The results have only been tested up to degree 100, higher degrees may + be problematic. The weights are determined by using the fact that + + .. math:: w_k = c / (L'_n(x_k) * L_{n-1}(x_k)) + + where :math:`c` is a constant independent of :math:`k` and :math:`x_k` + is the k'th root of :math:`L_n`, and then scaling the results to get + the right value when integrating 1. + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg <= 0: + raise ValueError("deg must be a positive integer") + + # first approximation of roots. We use the fact that the companion + # matrix is symmetric in this case in order to obtain better zeros. + c = np.array([0]*deg + [1]) + m = legcompanion(c) + x = la.eigvalsh(m) + + # improve roots by one application of Newton + dy = legval(x, c) + df = legval(x, legder(c)) + x -= dy/df + + # compute the weights. We scale the factor to avoid possible numerical + # overflow. + fm = legval(x, c[1:]) + fm /= np.abs(fm).max() + df /= np.abs(df).max() + w = 1/(fm * df) + + # for Legendre we can also symmetrize + w = (w + w[::-1])/2 + x = (x - x[::-1])/2 + + # scale w to get the right value + w *= 2. / w.sum() + + return x, w + + +def legweight(x): + """ + Weight function of the Legendre polynomials. + + The weight function is :math:`1` and the interval of integration is + :math:`[-1, 1]`. The Legendre polynomials are orthogonal, but not + normalized, with respect to this weight function. + + Parameters + ---------- + x : array_like + Values at which the weight function will be computed. + + Returns + ------- + w : ndarray + The weight function at `x`. + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + w = x*0.0 + 1.0 + return w + +# +# Legendre series class +# + +class Legendre(ABCPolyBase): + """A Legendre series class. + + The Legendre class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed in the `ABCPolyBase` documentation. + + Parameters + ---------- + coef : array_like + Legendre coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` gives ``1*P_0(x) + 2*P_1(x) + 3*P_2(x)``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(legadd) + _sub = staticmethod(legsub) + _mul = staticmethod(legmul) + _div = staticmethod(legdiv) + _pow = staticmethod(legpow) + _val = staticmethod(legval) + _int = staticmethod(legint) + _der = staticmethod(legder) + _fit = staticmethod(legfit) + _line = staticmethod(legline) + _roots = staticmethod(legroots) + _fromroots = staticmethod(legfromroots) + + # Virtual properties + domain = np.array(legdomain) + window = np.array(legdomain) + basis_name = 'P' diff --git a/lib/python3.12/site-packages/numpy/polynomial/legendre.pyi b/lib/python3.12/site-packages/numpy/polynomial/legendre.pyi new file mode 100644 index 0000000000000000000000000000000000000000..63a1c3f3a1f89c2c2da61e385f7dba1e7be16c06 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/legendre.pyi @@ -0,0 +1,46 @@ +from typing import Any + +from numpy import ndarray, dtype, int_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +legtrim = trimcoef + +def poly2leg(pol): ... +def leg2poly(c): ... + +legdomain: ndarray[Any, dtype[int_]] +legzero: ndarray[Any, dtype[int_]] +legone: ndarray[Any, dtype[int_]] +legx: ndarray[Any, dtype[int_]] + +def legline(off, scl): ... +def legfromroots(roots): ... +def legadd(c1, c2): ... +def legsub(c1, c2): ... +def legmulx(c): ... +def legmul(c1, c2): ... +def legdiv(c1, c2): ... +def legpow(c, pow, maxpower=...): ... +def legder(c, m=..., scl=..., axis=...): ... +def legint(c, m=..., k = ..., lbnd=..., scl=..., axis=...): ... +def legval(x, c, tensor=...): ... +def legval2d(x, y, c): ... +def leggrid2d(x, y, c): ... +def legval3d(x, y, z, c): ... +def leggrid3d(x, y, z, c): ... +def legvander(x, deg): ... +def legvander2d(x, y, deg): ... +def legvander3d(x, y, z, deg): ... +def legfit(x, y, deg, rcond=..., full=..., w=...): ... +def legcompanion(c): ... +def legroots(c): ... +def leggauss(deg): ... +def legweight(x): ... + +class Legendre(ABCPolyBase): + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/polynomial.py b/lib/python3.12/site-packages/numpy/polynomial/polynomial.py new file mode 100644 index 0000000000000000000000000000000000000000..ceadff0bf4ed32f8bbbb9f208bf4d84946efe195 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/polynomial.py @@ -0,0 +1,1542 @@ +""" +================================================= +Power Series (:mod:`numpy.polynomial.polynomial`) +================================================= + +This module provides a number of objects (mostly functions) useful for +dealing with polynomials, including a `Polynomial` class that +encapsulates the usual arithmetic operations. (General information +on how this module represents and works with polynomial objects is in +the docstring for its "parent" sub-package, `numpy.polynomial`). + +Classes +------- +.. autosummary:: + :toctree: generated/ + + Polynomial + +Constants +--------- +.. autosummary:: + :toctree: generated/ + + polydomain + polyzero + polyone + polyx + +Arithmetic +---------- +.. autosummary:: + :toctree: generated/ + + polyadd + polysub + polymulx + polymul + polydiv + polypow + polyval + polyval2d + polyval3d + polygrid2d + polygrid3d + +Calculus +-------- +.. autosummary:: + :toctree: generated/ + + polyder + polyint + +Misc Functions +-------------- +.. autosummary:: + :toctree: generated/ + + polyfromroots + polyroots + polyvalfromroots + polyvander + polyvander2d + polyvander3d + polycompanion + polyfit + polytrim + polyline + +See Also +-------- +`numpy.polynomial` + +""" +__all__ = [ + 'polyzero', 'polyone', 'polyx', 'polydomain', 'polyline', 'polyadd', + 'polysub', 'polymulx', 'polymul', 'polydiv', 'polypow', 'polyval', + 'polyvalfromroots', 'polyder', 'polyint', 'polyfromroots', 'polyvander', + 'polyfit', 'polytrim', 'polyroots', 'Polynomial', 'polyval2d', 'polyval3d', + 'polygrid2d', 'polygrid3d', 'polyvander2d', 'polyvander3d'] + +import numpy as np +import numpy.linalg as la +from numpy.core.multiarray import normalize_axis_index + +from . import polyutils as pu +from ._polybase import ABCPolyBase + +polytrim = pu.trimcoef + +# +# These are constant arrays are of integer type so as to be compatible +# with the widest range of other types, such as Decimal. +# + +# Polynomial default domain. +polydomain = np.array([-1, 1]) + +# Polynomial coefficients representing zero. +polyzero = np.array([0]) + +# Polynomial coefficients representing one. +polyone = np.array([1]) + +# Polynomial coefficients representing the identity x. +polyx = np.array([0, 1]) + +# +# Polynomial series functions +# + + +def polyline(off, scl): + """ + Returns an array representing a linear polynomial. + + Parameters + ---------- + off, scl : scalars + The "y-intercept" and "slope" of the line, respectively. + + Returns + ------- + y : ndarray + This module's representation of the linear polynomial ``off + + scl*x``. + + See Also + -------- + numpy.polynomial.chebyshev.chebline + numpy.polynomial.legendre.legline + numpy.polynomial.laguerre.lagline + numpy.polynomial.hermite.hermline + numpy.polynomial.hermite_e.hermeline + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polyline(1,-1) + array([ 1, -1]) + >>> P.polyval(1, P.polyline(1,-1)) # should be 0 + 0.0 + + """ + if scl != 0: + return np.array([off, scl]) + else: + return np.array([off]) + + +def polyfromroots(roots): + """ + Generate a monic polynomial with given roots. + + Return the coefficients of the polynomial + + .. math:: p(x) = (x - r_0) * (x - r_1) * ... * (x - r_n), + + where the ``r_n`` are the roots specified in `roots`. If a zero has + multiplicity n, then it must appear in `roots` n times. For instance, + if 2 is a root of multiplicity three and 3 is a root of multiplicity 2, + then `roots` looks something like [2, 2, 2, 3, 3]. The roots can appear + in any order. + + If the returned coefficients are `c`, then + + .. math:: p(x) = c_0 + c_1 * x + ... + x^n + + The coefficient of the last term is 1 for monic polynomials in this + form. + + Parameters + ---------- + roots : array_like + Sequence containing the roots. + + Returns + ------- + out : ndarray + 1-D array of the polynomial's coefficients If all the roots are + real, then `out` is also real, otherwise it is complex. (see + Examples below). + + See Also + -------- + numpy.polynomial.chebyshev.chebfromroots + numpy.polynomial.legendre.legfromroots + numpy.polynomial.laguerre.lagfromroots + numpy.polynomial.hermite.hermfromroots + numpy.polynomial.hermite_e.hermefromroots + + Notes + ----- + The coefficients are determined by multiplying together linear factors + of the form ``(x - r_i)``, i.e. + + .. math:: p(x) = (x - r_0) (x - r_1) ... (x - r_n) + + where ``n == len(roots) - 1``; note that this implies that ``1`` is always + returned for :math:`a_n`. + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polyfromroots((-1,0,1)) # x(x - 1)(x + 1) = x^3 - x + array([ 0., -1., 0., 1.]) + >>> j = complex(0,1) + >>> P.polyfromroots((-j,j)) # complex returned, though values are real + array([1.+0.j, 0.+0.j, 1.+0.j]) + + """ + return pu._fromroots(polyline, polymul, roots) + + +def polyadd(c1, c2): + """ + Add one polynomial to another. + + Returns the sum of two polynomials `c1` + `c2`. The arguments are + sequences of coefficients from lowest order term to highest, i.e., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to high. + + Returns + ------- + out : ndarray + The coefficient array representing their sum. + + See Also + -------- + polysub, polymulx, polymul, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> sum = P.polyadd(c1,c2); sum + array([4., 4., 4.]) + >>> P.polyval(2, sum) # 4 + 4(2) + 4(2**2) + 28.0 + + """ + return pu._add(c1, c2) + + +def polysub(c1, c2): + """ + Subtract one polynomial from another. + + Returns the difference of two polynomials `c1` - `c2`. The arguments + are sequences of coefficients from lowest order term to highest, i.e., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Of coefficients representing their difference. + + See Also + -------- + polyadd, polymulx, polymul, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> P.polysub(c1,c2) + array([-2., 0., 2.]) + >>> P.polysub(c2,c1) # -P.polysub(c1,c2) + array([ 2., 0., -2.]) + + """ + return pu._sub(c1, c2) + + +def polymulx(c): + """Multiply a polynomial by x. + + Multiply the polynomial `c` by x, where x is the independent + variable. + + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients ordered from low to + high. + + Returns + ------- + out : ndarray + Array representing the result of the multiplication. + + See Also + -------- + polyadd, polysub, polymul, polydiv, polypow + + Notes + ----- + + .. versionadded:: 1.5.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + # The zero series needs special treatment + if len(c) == 1 and c[0] == 0: + return c + + prd = np.empty(len(c) + 1, dtype=c.dtype) + prd[0] = c[0]*0 + prd[1:] = c + return prd + + +def polymul(c1, c2): + """ + Multiply one polynomial by another. + + Returns the product of two polynomials `c1` * `c2`. The arguments are + sequences of coefficients, from lowest order term to highest, e.g., + [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2.`` + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of coefficients representing a polynomial, relative to the + "standard" basis, and ordered from lowest order term to highest. + + Returns + ------- + out : ndarray + Of the coefficients of their product. + + See Also + -------- + polyadd, polysub, polymulx, polydiv, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> P.polymul(c1,c2) + array([ 3., 8., 14., 8., 3.]) + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + ret = np.convolve(c1, c2) + return pu.trimseq(ret) + + +def polydiv(c1, c2): + """ + Divide one polynomial by another. + + Returns the quotient-with-remainder of two polynomials `c1` / `c2`. + The arguments are sequences of coefficients, from lowest order term + to highest, e.g., [1,2,3] represents ``1 + 2*x + 3*x**2``. + + Parameters + ---------- + c1, c2 : array_like + 1-D arrays of polynomial coefficients ordered from low to high. + + Returns + ------- + [quo, rem] : ndarrays + Of coefficient series representing the quotient and remainder. + + See Also + -------- + polyadd, polysub, polymulx, polymul, polypow + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c1 = (1,2,3) + >>> c2 = (3,2,1) + >>> P.polydiv(c1,c2) + (array([3.]), array([-8., -4.])) + >>> P.polydiv(c2,c1) + (array([ 0.33333333]), array([ 2.66666667, 1.33333333])) # may vary + + """ + # c1, c2 are trimmed copies + [c1, c2] = pu.as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError() + + # note: this is more efficient than `pu._div(polymul, c1, c2)` + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + dlen = lc1 - lc2 + scl = c2[-1] + c2 = c2[:-1]/scl + i = dlen + j = lc1 - 1 + while i >= 0: + c1[i:j] -= c2*c1[j] + i -= 1 + j -= 1 + return c1[j+1:]/scl, pu.trimseq(c1[:j+1]) + + +def polypow(c, pow, maxpower=None): + """Raise a polynomial to a power. + + Returns the polynomial `c` raised to the power `pow`. The argument + `c` is a sequence of coefficients ordered from low to high. i.e., + [1,2,3] is the series ``1 + 2*x + 3*x**2.`` + + Parameters + ---------- + c : array_like + 1-D array of array of series coefficients ordered from low to + high degree. + pow : integer + Power to which the series will be raised + maxpower : integer, optional + Maximum power allowed. This is mainly to limit growth of the series + to unmanageable size. Default is 16 + + Returns + ------- + coef : ndarray + Power series of power. + + See Also + -------- + polyadd, polysub, polymulx, polymul, polydiv + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> P.polypow([1,2,3], 2) + array([ 1., 4., 10., 12., 9.]) + + """ + # note: this is more efficient than `pu._pow(polymul, c1, c2)`, as it + # avoids calling `as_series` repeatedly + return pu._pow(np.convolve, c, pow, maxpower) + + +def polyder(c, m=1, scl=1, axis=0): + """ + Differentiate a polynomial. + + Returns the polynomial coefficients `c` differentiated `m` times along + `axis`. At each iteration the result is multiplied by `scl` (the + scaling factor is for use in a linear change of variable). The + argument `c` is an array of coefficients from low to high degree along + each axis, e.g., [1,2,3] represents the polynomial ``1 + 2*x + 3*x**2`` + while [[1,2],[1,2]] represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is + ``x`` and axis=1 is ``y``. + + Parameters + ---------- + c : array_like + Array of polynomial coefficients. If c is multidimensional the + different axis correspond to different variables with the degree + in each axis given by the corresponding index. + m : int, optional + Number of derivatives taken, must be non-negative. (Default: 1) + scl : scalar, optional + Each differentiation is multiplied by `scl`. The end result is + multiplication by ``scl**m``. This is for use in a linear change + of variable. (Default: 1) + axis : int, optional + Axis over which the derivative is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + der : ndarray + Polynomial coefficients of the derivative. + + See Also + -------- + polyint + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1,2,3,4) # 1 + 2x + 3x**2 + 4x**3 + >>> P.polyder(c) # (d/dx)(c) = 2 + 6x + 12x**2 + array([ 2., 6., 12.]) + >>> P.polyder(c,3) # (d**3/dx**3)(c) = 24 + array([24.]) + >>> P.polyder(c,scl=-1) # (d/d(-x))(c) = -2 - 6x - 12x**2 + array([ -2., -6., -12.]) + >>> P.polyder(c,2,-1) # (d**2/d(-x)**2)(c) = 6 + 24x + array([ 6., 24.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype fails with NA + c = c + 0.0 + cdt = c.dtype + cnt = pu._deprecate_as_int(m, "the order of derivation") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of derivation must be non-negative") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + c = np.moveaxis(c, iaxis, 0) + n = len(c) + if cnt >= n: + c = c[:1]*0 + else: + for i in range(cnt): + n = n - 1 + c *= scl + der = np.empty((n,) + c.shape[1:], dtype=cdt) + for j in range(n, 0, -1): + der[j - 1] = j*c[j] + c = der + c = np.moveaxis(c, 0, iaxis) + return c + + +def polyint(c, m=1, k=[], lbnd=0, scl=1, axis=0): + """ + Integrate a polynomial. + + Returns the polynomial coefficients `c` integrated `m` times from + `lbnd` along `axis`. At each iteration the resulting series is + **multiplied** by `scl` and an integration constant, `k`, is added. + The scaling factor is for use in a linear change of variable. ("Buyer + beware": note that, depending on what one is doing, one may want `scl` + to be the reciprocal of what one might expect; for more information, + see the Notes section below.) The argument `c` is an array of + coefficients, from low to high degree along each axis, e.g., [1,2,3] + represents the polynomial ``1 + 2*x + 3*x**2`` while [[1,2],[1,2]] + represents ``1 + 1*x + 2*y + 2*x*y`` if axis=0 is ``x`` and axis=1 is + ``y``. + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients, ordered from low to high. + m : int, optional + Order of integration, must be positive. (Default: 1) + k : {[], list, scalar}, optional + Integration constant(s). The value of the first integral at zero + is the first value in the list, the value of the second integral + at zero is the second value, etc. If ``k == []`` (the default), + all constants are set to zero. If ``m == 1``, a single scalar can + be given instead of a list. + lbnd : scalar, optional + The lower bound of the integral. (Default: 0) + scl : scalar, optional + Following each integration the result is *multiplied* by `scl` + before the integration constant is added. (Default: 1) + axis : int, optional + Axis over which the integral is taken. (Default: 0). + + .. versionadded:: 1.7.0 + + Returns + ------- + S : ndarray + Coefficient array of the integral. + + Raises + ------ + ValueError + If ``m < 1``, ``len(k) > m``, ``np.ndim(lbnd) != 0``, or + ``np.ndim(scl) != 0``. + + See Also + -------- + polyder + + Notes + ----- + Note that the result of each integration is *multiplied* by `scl`. Why + is this important to note? Say one is making a linear change of + variable :math:`u = ax + b` in an integral relative to `x`. Then + :math:`dx = du/a`, so one will need to set `scl` equal to + :math:`1/a` - perhaps not what one would have first thought. + + Examples + -------- + >>> from numpy.polynomial import polynomial as P + >>> c = (1,2,3) + >>> P.polyint(c) # should return array([0, 1, 1, 1]) + array([0., 1., 1., 1.]) + >>> P.polyint(c,3) # should return array([0, 0, 0, 1/6, 1/12, 1/20]) + array([ 0. , 0. , 0. , 0.16666667, 0.08333333, # may vary + 0.05 ]) + >>> P.polyint(c,k=3) # should return array([3, 1, 1, 1]) + array([3., 1., 1., 1.]) + >>> P.polyint(c,lbnd=-2) # should return array([6, 1, 1, 1]) + array([6., 1., 1., 1.]) + >>> P.polyint(c,scl=-2) # should return array([0, -2, -2, -2]) + array([ 0., -2., -2., -2.]) + + """ + c = np.array(c, ndmin=1, copy=True) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype doesn't preserve mask attribute. + c = c + 0.0 + cdt = c.dtype + if not np.iterable(k): + k = [k] + cnt = pu._deprecate_as_int(m, "the order of integration") + iaxis = pu._deprecate_as_int(axis, "the axis") + if cnt < 0: + raise ValueError("The order of integration must be non-negative") + if len(k) > cnt: + raise ValueError("Too many integration constants") + if np.ndim(lbnd) != 0: + raise ValueError("lbnd must be a scalar.") + if np.ndim(scl) != 0: + raise ValueError("scl must be a scalar.") + iaxis = normalize_axis_index(iaxis, c.ndim) + + if cnt == 0: + return c + + k = list(k) + [0]*(cnt - len(k)) + c = np.moveaxis(c, iaxis, 0) + for i in range(cnt): + n = len(c) + c *= scl + if n == 1 and np.all(c[0] == 0): + c[0] += k[i] + else: + tmp = np.empty((n + 1,) + c.shape[1:], dtype=cdt) + tmp[0] = c[0]*0 + tmp[1] = c[0] + for j in range(1, n): + tmp[j + 1] = c[j]/(j + 1) + tmp[0] += k[i] - polyval(lbnd, tmp) + c = tmp + c = np.moveaxis(c, 0, iaxis) + return c + + +def polyval(x, c, tensor=True): + """ + Evaluate a polynomial at points x. + + If `c` is of length `n + 1`, this function returns the value + + .. math:: p(x) = c_0 + c_1 * x + ... + c_n * x^n + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `c`. + + If `c` is a 1-D array, then `p(x)` will have the same shape as `x`. If + `c` is multidimensional, then the shape of the result depends on the + value of `tensor`. If `tensor` is true the shape will be c.shape[1:] + + x.shape. If `tensor` is false the shape will be c.shape[1:]. Note that + scalars have shape (,). + + Trailing zeros in the coefficients will be used in the evaluation, so + they should be avoided if efficiency is a concern. + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `c`. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree n are contained in c[n]. If `c` is multidimensional the + remaining indices enumerate multiple polynomials. In the two + dimensional case the coefficients may be thought of as stored in + the columns of `c`. + tensor : boolean, optional + If True, the shape of the coefficient array is extended with ones + on the right, one for each dimension of `x`. Scalars have dimension 0 + for this action. The result is that every column of coefficients in + `c` is evaluated for every element of `x`. If False, `x` is broadcast + over the columns of `c` for the evaluation. This keyword is useful + when `c` is multidimensional. The default value is True. + + .. versionadded:: 1.7.0 + + Returns + ------- + values : ndarray, compatible object + The shape of the returned array is described above. + + See Also + -------- + polyval2d, polygrid2d, polyval3d, polygrid3d + + Notes + ----- + The evaluation uses Horner's method. + + Examples + -------- + >>> from numpy.polynomial.polynomial import polyval + >>> polyval(1, [1,2,3]) + 6.0 + >>> a = np.arange(4).reshape(2,2) + >>> a + array([[0, 1], + [2, 3]]) + >>> polyval(a, [1,2,3]) + array([[ 1., 6.], + [17., 34.]]) + >>> coef = np.arange(4).reshape(2,2) # multidimensional coefficients + >>> coef + array([[0, 1], + [2, 3]]) + >>> polyval([1,2], coef, tensor=True) + array([[2., 4.], + [4., 7.]]) + >>> polyval([1,2], coef, tensor=False) + array([2., 7.]) + + """ + c = np.array(c, ndmin=1, copy=False) + if c.dtype.char in '?bBhHiIlLqQpP': + # astype fails with NA + c = c + 0.0 + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray) and tensor: + c = c.reshape(c.shape + (1,)*x.ndim) + + c0 = c[-1] + x*0 + for i in range(2, len(c) + 1): + c0 = c[-i] + c0*x + return c0 + + +def polyvalfromroots(x, r, tensor=True): + """ + Evaluate a polynomial specified by its roots at points x. + + If `r` is of length `N`, this function returns the value + + .. math:: p(x) = \\prod_{n=1}^{N} (x - r_n) + + The parameter `x` is converted to an array only if it is a tuple or a + list, otherwise it is treated as a scalar. In either case, either `x` + or its elements must support multiplication and addition both with + themselves and with the elements of `r`. + + If `r` is a 1-D array, then `p(x)` will have the same shape as `x`. If `r` + is multidimensional, then the shape of the result depends on the value of + `tensor`. If `tensor` is ``True`` the shape will be r.shape[1:] + x.shape; + that is, each polynomial is evaluated at every value of `x`. If `tensor` is + ``False``, the shape will be r.shape[1:]; that is, each polynomial is + evaluated only for the corresponding broadcast value of `x`. Note that + scalars have shape (,). + + .. versionadded:: 1.12 + + Parameters + ---------- + x : array_like, compatible object + If `x` is a list or tuple, it is converted to an ndarray, otherwise + it is left unchanged and treated as a scalar. In either case, `x` + or its elements must support addition and multiplication with + with themselves and with the elements of `r`. + r : array_like + Array of roots. If `r` is multidimensional the first index is the + root index, while the remaining indices enumerate multiple + polynomials. For instance, in the two dimensional case the roots + of each polynomial may be thought of as stored in the columns of `r`. + tensor : boolean, optional + If True, the shape of the roots array is extended with ones on the + right, one for each dimension of `x`. Scalars have dimension 0 for this + action. The result is that every column of coefficients in `r` is + evaluated for every element of `x`. If False, `x` is broadcast over the + columns of `r` for the evaluation. This keyword is useful when `r` is + multidimensional. The default value is True. + + Returns + ------- + values : ndarray, compatible object + The shape of the returned array is described above. + + See Also + -------- + polyroots, polyfromroots, polyval + + Examples + -------- + >>> from numpy.polynomial.polynomial import polyvalfromroots + >>> polyvalfromroots(1, [1,2,3]) + 0.0 + >>> a = np.arange(4).reshape(2,2) + >>> a + array([[0, 1], + [2, 3]]) + >>> polyvalfromroots(a, [-1, 0, 1]) + array([[-0., 0.], + [ 6., 24.]]) + >>> r = np.arange(-2, 2).reshape(2,2) # multidimensional coefficients + >>> r # each column of r defines one polynomial + array([[-2, -1], + [ 0, 1]]) + >>> b = [-2, 1] + >>> polyvalfromroots(b, r, tensor=True) + array([[-0., 3.], + [ 3., 0.]]) + >>> polyvalfromroots(b, r, tensor=False) + array([-0., 0.]) + """ + r = np.array(r, ndmin=1, copy=False) + if r.dtype.char in '?bBhHiIlLqQpP': + r = r.astype(np.double) + if isinstance(x, (tuple, list)): + x = np.asarray(x) + if isinstance(x, np.ndarray): + if tensor: + r = r.reshape(r.shape + (1,)*x.ndim) + elif x.ndim >= r.ndim: + raise ValueError("x.ndim must be < r.ndim when tensor == False") + return np.prod(x - r, axis=0) + + +def polyval2d(x, y, c): + """ + Evaluate a 2-D polynomial at points (x, y). + + This function returns the value + + .. math:: p(x,y) = \\sum_{i,j} c_{i,j} * x^i * y^j + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars and they + must have the same shape after conversion. In either case, either `x` + and `y` or their elements must support multiplication and addition both + with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points `(x, y)`, + where `x` and `y` must have the same shape. If `x` or `y` is a list + or tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term + of multi-degree i,j is contained in `c[i,j]`. If `c` has + dimension greater than two the remaining indices enumerate multiple + sets of coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points formed with + pairs of corresponding values from `x` and `y`. + + See Also + -------- + polyval, polygrid2d, polyval3d, polygrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(polyval, c, x, y) + + +def polygrid2d(x, y, c): + """ + Evaluate a 2-D polynomial on the Cartesian product of x and y. + + This function returns the values: + + .. math:: p(a,b) = \\sum_{i,j} c_{i,j} * a^i * b^j + + where the points `(a, b)` consist of all pairs formed by taking + `a` from `x` and `b` from `y`. The resulting points form a grid with + `x` in the first dimension and `y` in the second. + + The parameters `x` and `y` are converted to arrays only if they are + tuples or a lists, otherwise they are treated as a scalars. In either + case, either `x` and `y` or their elements must support multiplication + and addition both with themselves and with the elements of `c`. + + If `c` has fewer than two dimensions, ones are implicitly appended to + its shape to make it 2-D. The shape of the result will be c.shape[2:] + + x.shape + y.shape. + + Parameters + ---------- + x, y : array_like, compatible objects + The two dimensional series is evaluated at the points in the + Cartesian product of `x` and `y`. If `x` or `y` is a list or + tuple, it is first converted to an ndarray, otherwise it is left + unchanged and, if it isn't an ndarray, it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + polyval, polyval2d, polyval3d, polygrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(polyval, c, x, y) + + +def polyval3d(x, y, z, c): + """ + Evaluate a 3-D polynomial at points (x, y, z). + + This function returns the values: + + .. math:: p(x,y,z) = \\sum_{i,j,k} c_{i,j,k} * x^i * y^j * z^k + + The parameters `x`, `y`, and `z` are converted to arrays only if + they are tuples or a lists, otherwise they are treated as a scalars and + they must have the same shape after conversion. In either case, either + `x`, `y`, and `z` or their elements must support multiplication and + addition both with themselves and with the elements of `c`. + + If `c` has fewer than 3 dimensions, ones are implicitly appended to its + shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape. + + Parameters + ---------- + x, y, z : array_like, compatible object + The three dimensional series is evaluated at the points + `(x, y, z)`, where `x`, `y`, and `z` must have the same shape. If + any of `x`, `y`, or `z` is a list or tuple, it is first converted + to an ndarray, otherwise it is left unchanged and if it isn't an + ndarray it is treated as a scalar. + c : array_like + Array of coefficients ordered so that the coefficient of the term of + multi-degree i,j,k is contained in ``c[i,j,k]``. If `c` has dimension + greater than 3 the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the multidimensional polynomial on points formed with + triples of corresponding values from `x`, `y`, and `z`. + + See Also + -------- + polyval, polyval2d, polygrid2d, polygrid3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._valnd(polyval, c, x, y, z) + + +def polygrid3d(x, y, z, c): + """ + Evaluate a 3-D polynomial on the Cartesian product of x, y and z. + + This function returns the values: + + .. math:: p(a,b,c) = \\sum_{i,j,k} c_{i,j,k} * a^i * b^j * c^k + + where the points `(a, b, c)` consist of all triples formed by taking + `a` from `x`, `b` from `y`, and `c` from `z`. The resulting points form + a grid with `x` in the first dimension, `y` in the second, and `z` in + the third. + + The parameters `x`, `y`, and `z` are converted to arrays only if they + are tuples or a lists, otherwise they are treated as a scalars. In + either case, either `x`, `y`, and `z` or their elements must support + multiplication and addition both with themselves and with the elements + of `c`. + + If `c` has fewer than three dimensions, ones are implicitly appended to + its shape to make it 3-D. The shape of the result will be c.shape[3:] + + x.shape + y.shape + z.shape. + + Parameters + ---------- + x, y, z : array_like, compatible objects + The three dimensional series is evaluated at the points in the + Cartesian product of `x`, `y`, and `z`. If `x`,`y`, or `z` is a + list or tuple, it is first converted to an ndarray, otherwise it is + left unchanged and, if it isn't an ndarray, it is treated as a + scalar. + c : array_like + Array of coefficients ordered so that the coefficients for terms of + degree i,j are contained in ``c[i,j]``. If `c` has dimension + greater than two the remaining indices enumerate multiple sets of + coefficients. + + Returns + ------- + values : ndarray, compatible object + The values of the two dimensional polynomial at points in the Cartesian + product of `x` and `y`. + + See Also + -------- + polyval, polyval2d, polygrid2d, polyval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._gridnd(polyval, c, x, y, z) + + +def polyvander(x, deg): + """Vandermonde matrix of given degree. + + Returns the Vandermonde matrix of degree `deg` and sample points + `x`. The Vandermonde matrix is defined by + + .. math:: V[..., i] = x^i, + + where `0 <= i <= deg`. The leading indices of `V` index the elements of + `x` and the last index is the power of `x`. + + If `c` is a 1-D array of coefficients of length `n + 1` and `V` is the + matrix ``V = polyvander(x, n)``, then ``np.dot(V, c)`` and + ``polyval(x, c)`` are the same up to roundoff. This equivalence is + useful both for least squares fitting and for the evaluation of a large + number of polynomials of the same degree and sample points. + + Parameters + ---------- + x : array_like + Array of points. The dtype is converted to float64 or complex128 + depending on whether any of the elements are complex. If `x` is + scalar it is converted to a 1-D array. + deg : int + Degree of the resulting matrix. + + Returns + ------- + vander : ndarray. + The Vandermonde matrix. The shape of the returned matrix is + ``x.shape + (deg + 1,)``, where the last index is the power of `x`. + The dtype will be the same as the converted `x`. + + See Also + -------- + polyvander2d, polyvander3d + + """ + ideg = pu._deprecate_as_int(deg, "deg") + if ideg < 0: + raise ValueError("deg must be non-negative") + + x = np.array(x, copy=False, ndmin=1) + 0.0 + dims = (ideg + 1,) + x.shape + dtyp = x.dtype + v = np.empty(dims, dtype=dtyp) + v[0] = x*0 + 1 + if ideg > 0: + v[1] = x + for i in range(2, ideg + 1): + v[i] = v[i-1]*x + return np.moveaxis(v, 0, -1) + + +def polyvander2d(x, y, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y)`. The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (deg[1] + 1)*i + j] = x^i * y^j, + + where `0 <= i <= deg[0]` and `0 <= j <= deg[1]`. The leading indices of + `V` index the points `(x, y)` and the last index encodes the powers of + `x` and `y`. + + If ``V = polyvander2d(x, y, [xdeg, ydeg])``, then the columns of `V` + correspond to the elements of a 2-D coefficient array `c` of shape + (xdeg + 1, ydeg + 1) in the order + + .. math:: c_{00}, c_{01}, c_{02} ... , c_{10}, c_{11}, c_{12} ... + + and ``np.dot(V, c.flat)`` and ``polyval2d(x, y, c)`` will be the same + up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 2-D polynomials + of the same degrees and sample points. + + Parameters + ---------- + x, y : array_like + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg]. + + Returns + ------- + vander2d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg([1]+1)`. The dtype will be the same + as the converted `x` and `y`. + + See Also + -------- + polyvander, polyvander3d, polyval2d, polyval3d + + """ + return pu._vander_nd_flat((polyvander, polyvander), (x, y), deg) + + +def polyvander3d(x, y, z, deg): + """Pseudo-Vandermonde matrix of given degrees. + + Returns the pseudo-Vandermonde matrix of degrees `deg` and sample + points `(x, y, z)`. If `l, m, n` are the given degrees in `x, y, z`, + then The pseudo-Vandermonde matrix is defined by + + .. math:: V[..., (m+1)(n+1)i + (n+1)j + k] = x^i * y^j * z^k, + + where `0 <= i <= l`, `0 <= j <= m`, and `0 <= j <= n`. The leading + indices of `V` index the points `(x, y, z)` and the last index encodes + the powers of `x`, `y`, and `z`. + + If ``V = polyvander3d(x, y, z, [xdeg, ydeg, zdeg])``, then the columns + of `V` correspond to the elements of a 3-D coefficient array `c` of + shape (xdeg + 1, ydeg + 1, zdeg + 1) in the order + + .. math:: c_{000}, c_{001}, c_{002},... , c_{010}, c_{011}, c_{012},... + + and ``np.dot(V, c.flat)`` and ``polyval3d(x, y, z, c)`` will be the + same up to roundoff. This equivalence is useful both for least squares + fitting and for the evaluation of a large number of 3-D polynomials + of the same degrees and sample points. + + Parameters + ---------- + x, y, z : array_like + Arrays of point coordinates, all of the same shape. The dtypes will + be converted to either float64 or complex128 depending on whether + any of the elements are complex. Scalars are converted to 1-D + arrays. + deg : list of ints + List of maximum degrees of the form [x_deg, y_deg, z_deg]. + + Returns + ------- + vander3d : ndarray + The shape of the returned matrix is ``x.shape + (order,)``, where + :math:`order = (deg[0]+1)*(deg([1]+1)*(deg[2]+1)`. The dtype will + be the same as the converted `x`, `y`, and `z`. + + See Also + -------- + polyvander, polyvander3d, polyval2d, polyval3d + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + return pu._vander_nd_flat((polyvander, polyvander, polyvander), (x, y, z), deg) + + +def polyfit(x, y, deg, rcond=None, full=False, w=None): + """ + Least-squares fit of a polynomial to data. + + Return the coefficients of a polynomial of degree `deg` that is the + least squares fit to the data values `y` given at points `x`. If `y` is + 1-D the returned coefficients will also be 1-D. If `y` is 2-D multiple + fits are done, one for each column of `y`, and the resulting + coefficients are stored in the corresponding columns of a 2-D return. + The fitted polynomial(s) are in the form + + .. math:: p(x) = c_0 + c_1 * x + ... + c_n * x^n, + + where `n` is `deg`. + + Parameters + ---------- + x : array_like, shape (`M`,) + x-coordinates of the `M` sample (data) points ``(x[i], y[i])``. + y : array_like, shape (`M`,) or (`M`, `K`) + y-coordinates of the sample points. Several sets of sample points + sharing the same x-coordinates can be (independently) fit with one + call to `polyfit` by passing in for `y` a 2-D array that contains + one data set per column. + deg : int or 1-D array_like + Degree(s) of the fitting polynomials. If `deg` is a single integer + all terms up to and including the `deg`'th term are included in the + fit. For NumPy versions >= 1.11.0 a list of integers specifying the + degrees of the terms to include may be used instead. + rcond : float, optional + Relative condition number of the fit. Singular values smaller + than `rcond`, relative to the largest singular value, will be + ignored. The default value is ``len(x)*eps``, where `eps` is the + relative precision of the platform's float type, about 2e-16 in + most cases. + full : bool, optional + Switch determining the nature of the return value. When ``False`` + (the default) just the coefficients are returned; when ``True``, + diagnostic information from the singular value decomposition (used + to solve the fit's matrix equation) is also returned. + w : array_like, shape (`M`,), optional + Weights. If not None, the weight ``w[i]`` applies to the unsquared + residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are + chosen so that the errors of the products ``w[i]*y[i]`` all have the + same variance. When using inverse-variance weighting, use + ``w[i] = 1/sigma(y[i])``. The default value is None. + + .. versionadded:: 1.5.0 + + Returns + ------- + coef : ndarray, shape (`deg` + 1,) or (`deg` + 1, `K`) + Polynomial coefficients ordered from low to high. If `y` was 2-D, + the coefficients in column `k` of `coef` represent the polynomial + fit to the data in `y`'s `k`-th column. + + [residuals, rank, singular_values, rcond] : list + These values are only returned if ``full == True`` + + - residuals -- sum of squared residuals of the least squares fit + - rank -- the numerical rank of the scaled Vandermonde matrix + - singular_values -- singular values of the scaled Vandermonde matrix + - rcond -- value of `rcond`. + + For more details, see `numpy.linalg.lstsq`. + + Raises + ------ + RankWarning + Raised if the matrix in the least-squares fit is rank deficient. + The warning is only raised if ``full == False``. The warnings can + be turned off by: + + >>> import warnings + >>> warnings.simplefilter('ignore', np.RankWarning) + + See Also + -------- + numpy.polynomial.chebyshev.chebfit + numpy.polynomial.legendre.legfit + numpy.polynomial.laguerre.lagfit + numpy.polynomial.hermite.hermfit + numpy.polynomial.hermite_e.hermefit + polyval : Evaluates a polynomial. + polyvander : Vandermonde matrix for powers. + numpy.linalg.lstsq : Computes a least-squares fit from the matrix. + scipy.interpolate.UnivariateSpline : Computes spline fits. + + Notes + ----- + The solution is the coefficients of the polynomial `p` that minimizes + the sum of the weighted squared errors + + .. math:: E = \\sum_j w_j^2 * |y_j - p(x_j)|^2, + + where the :math:`w_j` are the weights. This problem is solved by + setting up the (typically) over-determined matrix equation: + + .. math:: V(x) * c = w * y, + + where `V` is the weighted pseudo Vandermonde matrix of `x`, `c` are the + coefficients to be solved for, `w` are the weights, and `y` are the + observed values. This equation is then solved using the singular value + decomposition of `V`. + + If some of the singular values of `V` are so small that they are + neglected (and `full` == ``False``), a `RankWarning` will be raised. + This means that the coefficient values may be poorly determined. + Fitting to a lower order polynomial will usually get rid of the warning + (but may not be what you want, of course; if you have independent + reason(s) for choosing the degree which isn't working, you may have to: + a) reconsider those reasons, and/or b) reconsider the quality of your + data). The `rcond` parameter can also be set to a value smaller than + its default, but the resulting fit may be spurious and have large + contributions from roundoff error. + + Polynomial fits using double precision tend to "fail" at about + (polynomial) degree 20. Fits using Chebyshev or Legendre series are + generally better conditioned, but much can still depend on the + distribution of the sample points and the smoothness of the data. If + the quality of the fit is inadequate, splines may be a good + alternative. + + Examples + -------- + >>> np.random.seed(123) + >>> from numpy.polynomial import polynomial as P + >>> x = np.linspace(-1,1,51) # x "data": [-1, -0.96, ..., 0.96, 1] + >>> y = x**3 - x + np.random.randn(len(x)) # x^3 - x + Gaussian noise + >>> c, stats = P.polyfit(x,y,3,full=True) + >>> np.random.seed(123) + >>> c # c[0], c[2] should be approx. 0, c[1] approx. -1, c[3] approx. 1 + array([ 0.01909725, -1.30598256, -0.00577963, 1.02644286]) # may vary + >>> stats # note the large SSR, explaining the rather poor results + [array([ 38.06116253]), 4, array([ 1.38446749, 1.32119158, 0.50443316, # may vary + 0.28853036]), 1.1324274851176597e-014] + + Same thing without the added noise + + >>> y = x**3 - x + >>> c, stats = P.polyfit(x,y,3,full=True) + >>> c # c[0], c[2] should be "very close to 0", c[1] ~= -1, c[3] ~= 1 + array([-6.36925336e-18, -1.00000000e+00, -4.08053781e-16, 1.00000000e+00]) + >>> stats # note the minuscule SSR + [array([ 7.46346754e-31]), 4, array([ 1.38446749, 1.32119158, # may vary + 0.50443316, 0.28853036]), 1.1324274851176597e-014] + + """ + return pu._fit(polyvander, x, y, deg, rcond, full, w) + + +def polycompanion(c): + """ + Return the companion matrix of c. + + The companion matrix for power series cannot be made symmetric by + scaling the basis, so this function differs from those for the + orthogonal polynomials. + + Parameters + ---------- + c : array_like + 1-D array of polynomial coefficients ordered from low to high + degree. + + Returns + ------- + mat : ndarray + Companion matrix of dimensions (deg, deg). + + Notes + ----- + + .. versionadded:: 1.7.0 + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + raise ValueError('Series must have maximum degree of at least 1.') + if len(c) == 2: + return np.array([[-c[0]/c[1]]]) + + n = len(c) - 1 + mat = np.zeros((n, n), dtype=c.dtype) + bot = mat.reshape(-1)[n::n+1] + bot[...] = 1 + mat[:, -1] -= c[:-1]/c[-1] + return mat + + +def polyroots(c): + """ + Compute the roots of a polynomial. + + Return the roots (a.k.a. "zeros") of the polynomial + + .. math:: p(x) = \\sum_i c[i] * x^i. + + Parameters + ---------- + c : 1-D array_like + 1-D array of polynomial coefficients. + + Returns + ------- + out : ndarray + Array of the roots of the polynomial. If all the roots are real, + then `out` is also real, otherwise it is complex. + + See Also + -------- + numpy.polynomial.chebyshev.chebroots + numpy.polynomial.legendre.legroots + numpy.polynomial.laguerre.lagroots + numpy.polynomial.hermite.hermroots + numpy.polynomial.hermite_e.hermeroots + + Notes + ----- + The root estimates are obtained as the eigenvalues of the companion + matrix, Roots far from the origin of the complex plane may have large + errors due to the numerical instability of the power series for such + values. Roots with multiplicity greater than 1 will also show larger + errors as the value of the series near such points is relatively + insensitive to errors in the roots. Isolated roots near the origin can + be improved by a few iterations of Newton's method. + + Examples + -------- + >>> import numpy.polynomial.polynomial as poly + >>> poly.polyroots(poly.polyfromroots((-1,0,1))) + array([-1., 0., 1.]) + >>> poly.polyroots(poly.polyfromroots((-1,0,1))).dtype + dtype('float64') + >>> j = complex(0,1) + >>> poly.polyroots(poly.polyfromroots((-j,0,j))) + array([ 0.00000000e+00+0.j, 0.00000000e+00+1.j, 2.77555756e-17-1.j]) # may vary + + """ + # c is a trimmed copy + [c] = pu.as_series([c]) + if len(c) < 2: + return np.array([], dtype=c.dtype) + if len(c) == 2: + return np.array([-c[0]/c[1]]) + + # rotated companion matrix reduces error + m = polycompanion(c)[::-1,::-1] + r = la.eigvals(m) + r.sort() + return r + + +# +# polynomial class +# + +class Polynomial(ABCPolyBase): + """A power series class. + + The Polynomial class provides the standard Python numerical methods + '+', '-', '*', '//', '%', 'divmod', '**', and '()' as well as the + attributes and methods listed in the `ABCPolyBase` documentation. + + Parameters + ---------- + coef : array_like + Polynomial coefficients in order of increasing degree, i.e., + ``(1, 2, 3)`` give ``1 + 2*x + 3*x**2``. + domain : (2,) array_like, optional + Domain to use. The interval ``[domain[0], domain[1]]`` is mapped + to the interval ``[window[0], window[1]]`` by shifting and scaling. + The default value is [-1, 1]. + window : (2,) array_like, optional + Window, see `domain` for its use. The default value is [-1, 1]. + + .. versionadded:: 1.6.0 + symbol : str, optional + Symbol used to represent the independent variable in string + representations of the polynomial expression, e.g. for printing. + The symbol must be a valid Python identifier. Default value is 'x'. + + .. versionadded:: 1.24 + + """ + # Virtual Functions + _add = staticmethod(polyadd) + _sub = staticmethod(polysub) + _mul = staticmethod(polymul) + _div = staticmethod(polydiv) + _pow = staticmethod(polypow) + _val = staticmethod(polyval) + _int = staticmethod(polyint) + _der = staticmethod(polyder) + _fit = staticmethod(polyfit) + _line = staticmethod(polyline) + _roots = staticmethod(polyroots) + _fromroots = staticmethod(polyfromroots) + + # Virtual properties + domain = np.array(polydomain) + window = np.array(polydomain) + basis_name = None + + @classmethod + def _str_term_unicode(cls, i, arg_str): + if i == '1': + return f"·{arg_str}" + else: + return f"·{arg_str}{i.translate(cls._superscript_mapping)}" + + @staticmethod + def _str_term_ascii(i, arg_str): + if i == '1': + return f" {arg_str}" + else: + return f" {arg_str}**{i}" + + @staticmethod + def _repr_latex_term(i, arg_str, needs_parens): + if needs_parens: + arg_str = rf"\left({arg_str}\right)" + if i == 0: + return '1' + elif i == 1: + return arg_str + else: + return f"{arg_str}^{{{i}}}" diff --git a/lib/python3.12/site-packages/numpy/polynomial/polynomial.pyi b/lib/python3.12/site-packages/numpy/polynomial/polynomial.pyi new file mode 100644 index 0000000000000000000000000000000000000000..3c87f9d2926615e09bffd03d00306b6f235ec1c2 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/polynomial.pyi @@ -0,0 +1,41 @@ +from typing import Any + +from numpy import ndarray, dtype, int_ +from numpy.polynomial._polybase import ABCPolyBase +from numpy.polynomial.polyutils import trimcoef + +__all__: list[str] + +polytrim = trimcoef + +polydomain: ndarray[Any, dtype[int_]] +polyzero: ndarray[Any, dtype[int_]] +polyone: ndarray[Any, dtype[int_]] +polyx: ndarray[Any, dtype[int_]] + +def polyline(off, scl): ... +def polyfromroots(roots): ... +def polyadd(c1, c2): ... +def polysub(c1, c2): ... +def polymulx(c): ... +def polymul(c1, c2): ... +def polydiv(c1, c2): ... +def polypow(c, pow, maxpower=...): ... +def polyder(c, m=..., scl=..., axis=...): ... +def polyint(c, m=..., k=..., lbnd=..., scl=..., axis=...): ... +def polyval(x, c, tensor=...): ... +def polyvalfromroots(x, r, tensor=...): ... +def polyval2d(x, y, c): ... +def polygrid2d(x, y, c): ... +def polyval3d(x, y, z, c): ... +def polygrid3d(x, y, z, c): ... +def polyvander(x, deg): ... +def polyvander2d(x, y, deg): ... +def polyvander3d(x, y, z, deg): ... +def polyfit(x, y, deg, rcond=..., full=..., w=...): ... +def polyroots(c): ... + +class Polynomial(ABCPolyBase): + domain: Any + window: Any + basis_name: Any diff --git a/lib/python3.12/site-packages/numpy/polynomial/polyutils.py b/lib/python3.12/site-packages/numpy/polynomial/polyutils.py new file mode 100644 index 0000000000000000000000000000000000000000..4829138920169efc5b18b20be4a7d7c9509ba7fb --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/polyutils.py @@ -0,0 +1,789 @@ +""" +Utility classes and functions for the polynomial modules. + +This module provides: error and warning objects; a polynomial base class; +and some routines used in both the `polynomial` and `chebyshev` modules. + +Warning objects +--------------- + +.. autosummary:: + :toctree: generated/ + + RankWarning raised in least-squares fit for rank-deficient matrix. + +Functions +--------- + +.. autosummary:: + :toctree: generated/ + + as_series convert list of array_likes into 1-D arrays of common type. + trimseq remove trailing zeros. + trimcoef remove small trailing coefficients. + getdomain return the domain appropriate for a given set of abscissae. + mapdomain maps points between domains. + mapparms parameters of the linear map between domains. + +""" +import operator +import functools +import warnings + +import numpy as np + +from numpy.core.multiarray import dragon4_positional, dragon4_scientific +from numpy.core.umath import absolute + +__all__ = [ + 'RankWarning', 'as_series', 'trimseq', + 'trimcoef', 'getdomain', 'mapdomain', 'mapparms', + 'format_float'] + +# +# Warnings and Exceptions +# + +class RankWarning(UserWarning): + """Issued by chebfit when the design matrix is rank deficient.""" + pass + +# +# Helper functions to convert inputs to 1-D arrays +# +def trimseq(seq): + """Remove small Poly series coefficients. + + Parameters + ---------- + seq : sequence + Sequence of Poly series coefficients. This routine fails for + empty sequences. + + Returns + ------- + series : sequence + Subsequence with trailing zeros removed. If the resulting sequence + would be empty, return the first element. The returned sequence may + or may not be a view. + + Notes + ----- + Do not lose the type info if the sequence contains unknown objects. + + """ + if len(seq) == 0: + return seq + else: + for i in range(len(seq) - 1, -1, -1): + if seq[i] != 0: + break + return seq[:i+1] + + +def as_series(alist, trim=True): + """ + Return argument as a list of 1-d arrays. + + The returned list contains array(s) of dtype double, complex double, or + object. A 1-d argument of shape ``(N,)`` is parsed into ``N`` arrays of + size one; a 2-d argument of shape ``(M,N)`` is parsed into ``M`` arrays + of size ``N`` (i.e., is "parsed by row"); and a higher dimensional array + raises a Value Error if it is not first reshaped into either a 1-d or 2-d + array. + + Parameters + ---------- + alist : array_like + A 1- or 2-d array_like + trim : boolean, optional + When True, trailing zeros are removed from the inputs. + When False, the inputs are passed through intact. + + Returns + ------- + [a1, a2,...] : list of 1-D arrays + A copy of the input data as a list of 1-d arrays. + + Raises + ------ + ValueError + Raised when `as_series` cannot convert its input to 1-d arrays, or at + least one of the resulting arrays is empty. + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> a = np.arange(4) + >>> pu.as_series(a) + [array([0.]), array([1.]), array([2.]), array([3.])] + >>> b = np.arange(6).reshape((2,3)) + >>> pu.as_series(b) + [array([0., 1., 2.]), array([3., 4., 5.])] + + >>> pu.as_series((1, np.arange(3), np.arange(2, dtype=np.float16))) + [array([1.]), array([0., 1., 2.]), array([0., 1.])] + + >>> pu.as_series([2, [1.1, 0.]]) + [array([2.]), array([1.1])] + + >>> pu.as_series([2, [1.1, 0.]], trim=False) + [array([2.]), array([1.1, 0. ])] + + """ + arrays = [np.array(a, ndmin=1, copy=False) for a in alist] + if min([a.size for a in arrays]) == 0: + raise ValueError("Coefficient array is empty") + if any(a.ndim != 1 for a in arrays): + raise ValueError("Coefficient array is not 1-d") + if trim: + arrays = [trimseq(a) for a in arrays] + + if any(a.dtype == np.dtype(object) for a in arrays): + ret = [] + for a in arrays: + if a.dtype != np.dtype(object): + tmp = np.empty(len(a), dtype=np.dtype(object)) + tmp[:] = a[:] + ret.append(tmp) + else: + ret.append(a.copy()) + else: + try: + dtype = np.common_type(*arrays) + except Exception as e: + raise ValueError("Coefficient arrays have no common type") from e + ret = [np.array(a, copy=True, dtype=dtype) for a in arrays] + return ret + + +def trimcoef(c, tol=0): + """ + Remove "small" "trailing" coefficients from a polynomial. + + "Small" means "small in absolute value" and is controlled by the + parameter `tol`; "trailing" means highest order coefficient(s), e.g., in + ``[0, 1, 1, 0, 0]`` (which represents ``0 + x + x**2 + 0*x**3 + 0*x**4``) + both the 3-rd and 4-th order coefficients would be "trimmed." + + Parameters + ---------- + c : array_like + 1-d array of coefficients, ordered from lowest order to highest. + tol : number, optional + Trailing (i.e., highest order) elements with absolute value less + than or equal to `tol` (default value is zero) are removed. + + Returns + ------- + trimmed : ndarray + 1-d array with trailing zeros removed. If the resulting series + would be empty, a series containing a single zero is returned. + + Raises + ------ + ValueError + If `tol` < 0 + + See Also + -------- + trimseq + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> pu.trimcoef((0,0,3,0,5,0,0)) + array([0., 0., 3., 0., 5.]) + >>> pu.trimcoef((0,0,1e-3,0,1e-5,0,0),1e-3) # item == tol is trimmed + array([0.]) + >>> i = complex(0,1) # works for complex + >>> pu.trimcoef((3e-4,1e-3*(1-i),5e-4,2e-5*(1+i)), 1e-3) + array([0.0003+0.j , 0.001 -0.001j]) + + """ + if tol < 0: + raise ValueError("tol must be non-negative") + + [c] = as_series([c]) + [ind] = np.nonzero(np.abs(c) > tol) + if len(ind) == 0: + return c[:1]*0 + else: + return c[:ind[-1] + 1].copy() + +def getdomain(x): + """ + Return a domain suitable for given abscissae. + + Find a domain suitable for a polynomial or Chebyshev series + defined at the values supplied. + + Parameters + ---------- + x : array_like + 1-d array of abscissae whose domain will be determined. + + Returns + ------- + domain : ndarray + 1-d array containing two values. If the inputs are complex, then + the two returned points are the lower left and upper right corners + of the smallest rectangle (aligned with the axes) in the complex + plane containing the points `x`. If the inputs are real, then the + two points are the ends of the smallest interval containing the + points `x`. + + See Also + -------- + mapparms, mapdomain + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> points = np.arange(4)**2 - 5; points + array([-5, -4, -1, 4]) + >>> pu.getdomain(points) + array([-5., 4.]) + >>> c = np.exp(complex(0,1)*np.pi*np.arange(12)/6) # unit circle + >>> pu.getdomain(c) + array([-1.-1.j, 1.+1.j]) + + """ + [x] = as_series([x], trim=False) + if x.dtype.char in np.typecodes['Complex']: + rmin, rmax = x.real.min(), x.real.max() + imin, imax = x.imag.min(), x.imag.max() + return np.array((complex(rmin, imin), complex(rmax, imax))) + else: + return np.array((x.min(), x.max())) + +def mapparms(old, new): + """ + Linear map parameters between domains. + + Return the parameters of the linear map ``offset + scale*x`` that maps + `old` to `new` such that ``old[i] -> new[i]``, ``i = 0, 1``. + + Parameters + ---------- + old, new : array_like + Domains. Each domain must (successfully) convert to a 1-d array + containing precisely two values. + + Returns + ------- + offset, scale : scalars + The map ``L(x) = offset + scale*x`` maps the first domain to the + second. + + See Also + -------- + getdomain, mapdomain + + Notes + ----- + Also works for complex numbers, and thus can be used to calculate the + parameters required to map any line in the complex plane to any other + line therein. + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> pu.mapparms((-1,1),(-1,1)) + (0.0, 1.0) + >>> pu.mapparms((1,-1),(-1,1)) + (-0.0, -1.0) + >>> i = complex(0,1) + >>> pu.mapparms((-i,-1),(1,i)) + ((1+1j), (1-0j)) + + """ + oldlen = old[1] - old[0] + newlen = new[1] - new[0] + off = (old[1]*new[0] - old[0]*new[1])/oldlen + scl = newlen/oldlen + return off, scl + +def mapdomain(x, old, new): + """ + Apply linear map to input points. + + The linear map ``offset + scale*x`` that maps the domain `old` to + the domain `new` is applied to the points `x`. + + Parameters + ---------- + x : array_like + Points to be mapped. If `x` is a subtype of ndarray the subtype + will be preserved. + old, new : array_like + The two domains that determine the map. Each must (successfully) + convert to 1-d arrays containing precisely two values. + + Returns + ------- + x_out : ndarray + Array of points of the same shape as `x`, after application of the + linear map between the two domains. + + See Also + -------- + getdomain, mapparms + + Notes + ----- + Effectively, this implements: + + .. math:: + x\\_out = new[0] + m(x - old[0]) + + where + + .. math:: + m = \\frac{new[1]-new[0]}{old[1]-old[0]} + + Examples + -------- + >>> from numpy.polynomial import polyutils as pu + >>> old_domain = (-1,1) + >>> new_domain = (0,2*np.pi) + >>> x = np.linspace(-1,1,6); x + array([-1. , -0.6, -0.2, 0.2, 0.6, 1. ]) + >>> x_out = pu.mapdomain(x, old_domain, new_domain); x_out + array([ 0. , 1.25663706, 2.51327412, 3.76991118, 5.02654825, # may vary + 6.28318531]) + >>> x - pu.mapdomain(x_out, new_domain, old_domain) + array([0., 0., 0., 0., 0., 0.]) + + Also works for complex numbers (and thus can be used to map any line in + the complex plane to any other line therein). + + >>> i = complex(0,1) + >>> old = (-1 - i, 1 + i) + >>> new = (-1 + i, 1 - i) + >>> z = np.linspace(old[0], old[1], 6); z + array([-1. -1.j , -0.6-0.6j, -0.2-0.2j, 0.2+0.2j, 0.6+0.6j, 1. +1.j ]) + >>> new_z = pu.mapdomain(z, old, new); new_z + array([-1.0+1.j , -0.6+0.6j, -0.2+0.2j, 0.2-0.2j, 0.6-0.6j, 1.0-1.j ]) # may vary + + """ + x = np.asanyarray(x) + off, scl = mapparms(old, new) + return off + scl*x + + +def _nth_slice(i, ndim): + sl = [np.newaxis] * ndim + sl[i] = slice(None) + return tuple(sl) + + +def _vander_nd(vander_fs, points, degrees): + r""" + A generalization of the Vandermonde matrix for N dimensions + + The result is built by combining the results of 1d Vandermonde matrices, + + .. math:: + W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{V_k(x_k)[i_0, \ldots, i_M, j_k]} + + where + + .. math:: + N &= \texttt{len(points)} = \texttt{len(degrees)} = \texttt{len(vander\_fs)} \\ + M &= \texttt{points[k].ndim} \\ + V_k &= \texttt{vander\_fs[k]} \\ + x_k &= \texttt{points[k]} \\ + 0 \le j_k &\le \texttt{degrees[k]} + + Expanding the one-dimensional :math:`V_k` functions gives: + + .. math:: + W[i_0, \ldots, i_M, j_0, \ldots, j_N] = \prod_{k=0}^N{B_{k, j_k}(x_k[i_0, \ldots, i_M])} + + where :math:`B_{k,m}` is the m'th basis of the polynomial construction used along + dimension :math:`k`. For a regular polynomial, :math:`B_{k, m}(x) = P_m(x) = x^m`. + + Parameters + ---------- + vander_fs : Sequence[function(array_like, int) -> ndarray] + The 1d vander function to use for each axis, such as ``polyvander`` + points : Sequence[array_like] + Arrays of point coordinates, all of the same shape. The dtypes + will be converted to either float64 or complex128 depending on + whether any of the elements are complex. Scalars are converted to + 1-D arrays. + This must be the same length as `vander_fs`. + degrees : Sequence[int] + The maximum degree (inclusive) to use for each axis. + This must be the same length as `vander_fs`. + + Returns + ------- + vander_nd : ndarray + An array of shape ``points[0].shape + tuple(d + 1 for d in degrees)``. + """ + n_dims = len(vander_fs) + if n_dims != len(points): + raise ValueError( + f"Expected {n_dims} dimensions of sample points, got {len(points)}") + if n_dims != len(degrees): + raise ValueError( + f"Expected {n_dims} dimensions of degrees, got {len(degrees)}") + if n_dims == 0: + raise ValueError("Unable to guess a dtype or shape when no points are given") + + # convert to the same shape and type + points = tuple(np.array(tuple(points), copy=False) + 0.0) + + # produce the vandermonde matrix for each dimension, placing the last + # axis of each in an independent trailing axis of the output + vander_arrays = ( + vander_fs[i](points[i], degrees[i])[(...,) + _nth_slice(i, n_dims)] + for i in range(n_dims) + ) + + # we checked this wasn't empty already, so no `initial` needed + return functools.reduce(operator.mul, vander_arrays) + + +def _vander_nd_flat(vander_fs, points, degrees): + """ + Like `_vander_nd`, but flattens the last ``len(degrees)`` axes into a single axis + + Used to implement the public ``vanderd`` functions. + """ + v = _vander_nd(vander_fs, points, degrees) + return v.reshape(v.shape[:-len(degrees)] + (-1,)) + + +def _fromroots(line_f, mul_f, roots): + """ + Helper function used to implement the ``fromroots`` functions. + + Parameters + ---------- + line_f : function(float, float) -> ndarray + The ``line`` function, such as ``polyline`` + mul_f : function(array_like, array_like) -> ndarray + The ``mul`` function, such as ``polymul`` + roots + See the ``fromroots`` functions for more detail + """ + if len(roots) == 0: + return np.ones(1) + else: + [roots] = as_series([roots], trim=False) + roots.sort() + p = [line_f(-r, 1) for r in roots] + n = len(p) + while n > 1: + m, r = divmod(n, 2) + tmp = [mul_f(p[i], p[i+m]) for i in range(m)] + if r: + tmp[0] = mul_f(tmp[0], p[-1]) + p = tmp + n = m + return p[0] + + +def _valnd(val_f, c, *args): + """ + Helper function used to implement the ``vald`` functions. + + Parameters + ---------- + val_f : function(array_like, array_like, tensor: bool) -> array_like + The ``val`` function, such as ``polyval`` + c, args + See the ``vald`` functions for more detail + """ + args = [np.asanyarray(a) for a in args] + shape0 = args[0].shape + if not all((a.shape == shape0 for a in args[1:])): + if len(args) == 3: + raise ValueError('x, y, z are incompatible') + elif len(args) == 2: + raise ValueError('x, y are incompatible') + else: + raise ValueError('ordinates are incompatible') + it = iter(args) + x0 = next(it) + + # use tensor on only the first + c = val_f(x0, c) + for xi in it: + c = val_f(xi, c, tensor=False) + return c + + +def _gridnd(val_f, c, *args): + """ + Helper function used to implement the ``gridd`` functions. + + Parameters + ---------- + val_f : function(array_like, array_like, tensor: bool) -> array_like + The ``val`` function, such as ``polyval`` + c, args + See the ``gridd`` functions for more detail + """ + for xi in args: + c = val_f(xi, c) + return c + + +def _div(mul_f, c1, c2): + """ + Helper function used to implement the ``div`` functions. + + Implementation uses repeated subtraction of c2 multiplied by the nth basis. + For some polynomial types, a more efficient approach may be possible. + + Parameters + ---------- + mul_f : function(array_like, array_like) -> array_like + The ``mul`` function, such as ``polymul`` + c1, c2 + See the ``div`` functions for more detail + """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if c2[-1] == 0: + raise ZeroDivisionError() + + lc1 = len(c1) + lc2 = len(c2) + if lc1 < lc2: + return c1[:1]*0, c1 + elif lc2 == 1: + return c1/c2[-1], c1[:1]*0 + else: + quo = np.empty(lc1 - lc2 + 1, dtype=c1.dtype) + rem = c1 + for i in range(lc1 - lc2, - 1, -1): + p = mul_f([0]*i + [1], c2) + q = rem[-1]/p[-1] + rem = rem[:-1] - q*p[:-1] + quo[i] = q + return quo, trimseq(rem) + + +def _add(c1, c2): + """ Helper function used to implement the ``add`` functions. """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if len(c1) > len(c2): + c1[:c2.size] += c2 + ret = c1 + else: + c2[:c1.size] += c1 + ret = c2 + return trimseq(ret) + + +def _sub(c1, c2): + """ Helper function used to implement the ``sub`` functions. """ + # c1, c2 are trimmed copies + [c1, c2] = as_series([c1, c2]) + if len(c1) > len(c2): + c1[:c2.size] -= c2 + ret = c1 + else: + c2 = -c2 + c2[:c1.size] += c1 + ret = c2 + return trimseq(ret) + + +def _fit(vander_f, x, y, deg, rcond=None, full=False, w=None): + """ + Helper function used to implement the ``fit`` functions. + + Parameters + ---------- + vander_f : function(array_like, int) -> ndarray + The 1d vander function, such as ``polyvander`` + c1, c2 + See the ``fit`` functions for more detail + """ + x = np.asarray(x) + 0.0 + y = np.asarray(y) + 0.0 + deg = np.asarray(deg) + + # check arguments. + if deg.ndim > 1 or deg.dtype.kind not in 'iu' or deg.size == 0: + raise TypeError("deg must be an int or non-empty 1-D array of int") + if deg.min() < 0: + raise ValueError("expected deg >= 0") + if x.ndim != 1: + raise TypeError("expected 1D vector for x") + if x.size == 0: + raise TypeError("expected non-empty vector for x") + if y.ndim < 1 or y.ndim > 2: + raise TypeError("expected 1D or 2D array for y") + if len(x) != len(y): + raise TypeError("expected x and y to have same length") + + if deg.ndim == 0: + lmax = deg + order = lmax + 1 + van = vander_f(x, lmax) + else: + deg = np.sort(deg) + lmax = deg[-1] + order = len(deg) + van = vander_f(x, lmax)[:, deg] + + # set up the least squares matrices in transposed form + lhs = van.T + rhs = y.T + if w is not None: + w = np.asarray(w) + 0.0 + if w.ndim != 1: + raise TypeError("expected 1D vector for w") + if len(x) != len(w): + raise TypeError("expected x and w to have same length") + # apply weights. Don't use inplace operations as they + # can cause problems with NA. + lhs = lhs * w + rhs = rhs * w + + # set rcond + if rcond is None: + rcond = len(x)*np.finfo(x.dtype).eps + + # Determine the norms of the design matrix columns. + if issubclass(lhs.dtype.type, np.complexfloating): + scl = np.sqrt((np.square(lhs.real) + np.square(lhs.imag)).sum(1)) + else: + scl = np.sqrt(np.square(lhs).sum(1)) + scl[scl == 0] = 1 + + # Solve the least squares problem. + c, resids, rank, s = np.linalg.lstsq(lhs.T/scl, rhs.T, rcond) + c = (c.T/scl).T + + # Expand c to include non-fitted coefficients which are set to zero + if deg.ndim > 0: + if c.ndim == 2: + cc = np.zeros((lmax+1, c.shape[1]), dtype=c.dtype) + else: + cc = np.zeros(lmax+1, dtype=c.dtype) + cc[deg] = c + c = cc + + # warn on rank reduction + if rank != order and not full: + msg = "The fit may be poorly conditioned" + warnings.warn(msg, RankWarning, stacklevel=2) + + if full: + return c, [resids, rank, s, rcond] + else: + return c + + +def _pow(mul_f, c, pow, maxpower): + """ + Helper function used to implement the ``pow`` functions. + + Parameters + ---------- + mul_f : function(array_like, array_like) -> ndarray + The ``mul`` function, such as ``polymul`` + c : array_like + 1-D array of array of series coefficients + pow, maxpower + See the ``pow`` functions for more detail + """ + # c is a trimmed copy + [c] = as_series([c]) + power = int(pow) + if power != pow or power < 0: + raise ValueError("Power must be a non-negative integer.") + elif maxpower is not None and power > maxpower: + raise ValueError("Power is too large") + elif power == 0: + return np.array([1], dtype=c.dtype) + elif power == 1: + return c + else: + # This can be made more efficient by using powers of two + # in the usual way. + prd = c + for i in range(2, power + 1): + prd = mul_f(prd, c) + return prd + + +def _deprecate_as_int(x, desc): + """ + Like `operator.index`, but emits a deprecation warning when passed a float + + Parameters + ---------- + x : int-like, or float with integral value + Value to interpret as an integer + desc : str + description to include in any error message + + Raises + ------ + TypeError : if x is a non-integral float or non-numeric + DeprecationWarning : if x is an integral float + """ + try: + return operator.index(x) + except TypeError as e: + # Numpy 1.17.0, 2019-03-11 + try: + ix = int(x) + except TypeError: + pass + else: + if ix == x: + warnings.warn( + f"In future, this will raise TypeError, as {desc} will " + "need to be an integer not just an integral float.", + DeprecationWarning, + stacklevel=3 + ) + return ix + + raise TypeError(f"{desc} must be an integer") from e + + +def format_float(x, parens=False): + if not np.issubdtype(type(x), np.floating): + return str(x) + + opts = np.get_printoptions() + + if np.isnan(x): + return opts['nanstr'] + elif np.isinf(x): + return opts['infstr'] + + exp_format = False + if x != 0: + a = absolute(x) + if a >= 1.e8 or a < 10**min(0, -(opts['precision']-1)//2): + exp_format = True + + trim, unique = '0', True + if opts['floatmode'] == 'fixed': + trim, unique = 'k', False + + if exp_format: + s = dragon4_scientific(x, precision=opts['precision'], + unique=unique, trim=trim, + sign=opts['sign'] == '+') + if parens: + s = '(' + s + ')' + else: + s = dragon4_positional(x, precision=opts['precision'], + fractional=True, + unique=unique, trim=trim, + sign=opts['sign'] == '+') + return s diff --git a/lib/python3.12/site-packages/numpy/polynomial/polyutils.pyi b/lib/python3.12/site-packages/numpy/polynomial/polyutils.pyi new file mode 100644 index 0000000000000000000000000000000000000000..c0bcc67847f6b466c8d4fcf6f9b323df736c1c5f --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/polyutils.pyi @@ -0,0 +1,11 @@ +__all__: list[str] + +class RankWarning(UserWarning): ... + +def trimseq(seq): ... +def as_series(alist, trim=...): ... +def trimcoef(c, tol=...): ... +def getdomain(x): ... +def mapparms(old, new): ... +def mapdomain(x, old, new): ... +def format_float(x, parens=...): ... diff --git a/lib/python3.12/site-packages/numpy/polynomial/setup.py b/lib/python3.12/site-packages/numpy/polynomial/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..b58e867a133f804fbaf0d31099258a11e29058aa --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/setup.py @@ -0,0 +1,10 @@ +def configuration(parent_package='',top_path=None): + from numpy.distutils.misc_util import Configuration + config = Configuration('polynomial', parent_package, top_path) + config.add_subpackage('tests') + config.add_data_files('*.pyi') + return config + +if __name__ == '__main__': + from numpy.distutils.core import setup + setup(configuration=configuration) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__init__.py b/lib/python3.12/site-packages/numpy/polynomial/tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c5cbe356aff4342459b19e07858d86713fb4c310 Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_classes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9151233caaec684242220631dca2f9efa35424a7 Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_hermite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6aab650000f1c67ac5ee1264465213aad9da5084 Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_laguerre.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d64df9a94bb85290d0e8fba85d9ffb58dc66a402 Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_legendre.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7d315be5eadd979c7796c89956fe1c762f264ba Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polynomial.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94f0fe587b80446a73b9cea7ab65942193a0a3d1 Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_polyutils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0bf6db78bf6d5bfe8c81d8ca3c1f9fc73a598fb Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_printing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-312.pyc b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a0bac3e12e3368faabce0e08eaea26cebec878d Binary files /dev/null and b/lib/python3.12/site-packages/numpy/polynomial/tests/__pycache__/test_symbol.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_chebyshev.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_chebyshev.py new file mode 100644 index 0000000000000000000000000000000000000000..2f54bebfdb27d54f436378e4ab6d6c8f2426dd90 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_chebyshev.py @@ -0,0 +1,619 @@ +"""Tests for chebyshev module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.chebyshev as cheb +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + + +def trim(x): + return cheb.chebtrim(x, tol=1e-6) + +T0 = [1] +T1 = [0, 1] +T2 = [-1, 0, 2] +T3 = [0, -3, 0, 4] +T4 = [1, 0, -8, 0, 8] +T5 = [0, 5, 0, -20, 0, 16] +T6 = [-1, 0, 18, 0, -48, 0, 32] +T7 = [0, -7, 0, 56, 0, -112, 0, 64] +T8 = [1, 0, -32, 0, 160, 0, -256, 0, 128] +T9 = [0, 9, 0, -120, 0, 432, 0, -576, 0, 256] + +Tlist = [T0, T1, T2, T3, T4, T5, T6, T7, T8, T9] + + +class TestPrivate: + + def test__cseries_to_zseries(self): + for i in range(5): + inp = np.array([2] + [1]*i, np.double) + tgt = np.array([.5]*i + [2] + [.5]*i, np.double) + res = cheb._cseries_to_zseries(inp) + assert_equal(res, tgt) + + def test__zseries_to_cseries(self): + for i in range(5): + inp = np.array([.5]*i + [2] + [.5]*i, np.double) + tgt = np.array([2] + [1]*i, np.double) + res = cheb._zseries_to_cseries(inp) + assert_equal(res, tgt) + + +class TestConstants: + + def test_chebdomain(self): + assert_equal(cheb.chebdomain, [-1, 1]) + + def test_chebzero(self): + assert_equal(cheb.chebzero, [0]) + + def test_chebone(self): + assert_equal(cheb.chebone, [1]) + + def test_chebx(self): + assert_equal(cheb.chebx, [0, 1]) + + +class TestArithmetic: + + def test_chebadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = cheb.chebadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = cheb.chebsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebmulx(self): + assert_equal(cheb.chebmulx([0]), [0]) + assert_equal(cheb.chebmulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [.5, 0, .5] + assert_equal(cheb.chebmulx(ser), tgt) + + def test_chebmul(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(i + j + 1) + tgt[i + j] += .5 + tgt[abs(i - j)] += .5 + res = cheb.chebmul([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = cheb.chebadd(ci, cj) + quo, rem = cheb.chebdiv(tgt, ci) + res = cheb.chebadd(cheb.chebmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_chebpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(cheb.chebmul, [c]*j, np.array([1])) + res = cheb.chebpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2.5, 2., 1.5]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_chebval(self): + #check empty input + assert_equal(cheb.chebval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Tlist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = cheb.chebval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(cheb.chebval(x, [1]).shape, dims) + assert_equal(cheb.chebval(x, [1, 0]).shape, dims) + assert_equal(cheb.chebval(x, [1, 0, 0]).shape, dims) + + def test_chebval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, cheb.chebval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = cheb.chebval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_chebval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, cheb.chebval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = cheb.chebval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_chebgrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = cheb.chebgrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebgrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_chebgrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = cheb.chebgrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = cheb.chebgrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_chebint(self): + # check exceptions + assert_raises(TypeError, cheb.chebint, [0], .5) + assert_raises(ValueError, cheb.chebint, [0], -1) + assert_raises(ValueError, cheb.chebint, [0], 1, [0, 0]) + assert_raises(ValueError, cheb.chebint, [0], lbnd=[0]) + assert_raises(ValueError, cheb.chebint, [0], scl=[0]) + assert_raises(TypeError, cheb.chebint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = cheb.chebint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i]) + res = cheb.cheb2poly(chebint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(cheb.chebval(-1, chebint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + chebpol = cheb.poly2cheb(pol) + chebint = cheb.chebint(chebpol, m=1, k=[i], scl=2) + res = cheb.cheb2poly(chebint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1) + res = cheb.chebint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k]) + res = cheb.chebint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k], lbnd=-1) + res = cheb.chebint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = cheb.chebint(tgt, m=1, k=[k], scl=2) + res = cheb.chebint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([cheb.chebint(c) for c in c2d.T]).T + res = cheb.chebint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebint(c) for c in c2d]) + res = cheb.chebint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebint(c, k=3) for c in c2d]) + res = cheb.chebint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_chebder(self): + # check exceptions + assert_raises(TypeError, cheb.chebder, [0], .5) + assert_raises(ValueError, cheb.chebder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = cheb.chebder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = cheb.chebder(cheb.chebint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = cheb.chebder(cheb.chebint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([cheb.chebder(c) for c in c2d.T]).T + res = cheb.chebder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([cheb.chebder(c) for c in c2d]) + res = cheb.chebder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_chebvander(self): + # check for 1d x + x = np.arange(3) + v = cheb.chebvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], cheb.chebval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = cheb.chebvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], cheb.chebval(x, coef)) + + def test_chebvander2d(self): + # also tests chebval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = cheb.chebvander2d(x1, x2, [1, 2]) + tgt = cheb.chebval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = cheb.chebvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_chebvander3d(self): + # also tests chebval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = cheb.chebvander3d(x1, x2, x3, [1, 2, 3]) + tgt = cheb.chebval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = cheb.chebvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_chebfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, cheb.chebfit, [1], [1], -1) + assert_raises(TypeError, cheb.chebfit, [[1]], [1], 0) + assert_raises(TypeError, cheb.chebfit, [], [1], 0) + assert_raises(TypeError, cheb.chebfit, [1], [[[1]]], 0) + assert_raises(TypeError, cheb.chebfit, [1, 2], [1], 0) + assert_raises(TypeError, cheb.chebfit, [1], [1, 2], 0) + assert_raises(TypeError, cheb.chebfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, cheb.chebfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, cheb.chebfit, [1], [1], [-1,]) + assert_raises(ValueError, cheb.chebfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, cheb.chebfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = cheb.chebfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(cheb.chebval(x, coef3), y) + coef3 = cheb.chebfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(cheb.chebval(x, coef3), y) + # + coef4 = cheb.chebfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + coef4 = cheb.chebfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = cheb.chebfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(cheb.chebval(x, coef4), y) + # + coef2d = cheb.chebfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = cheb.chebfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = cheb.chebfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = cheb.chebfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = cheb.chebfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = cheb.chebfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(cheb.chebfit(x, x, 1), [0, 1]) + assert_almost_equal(cheb.chebfit(x, x, [0, 1]), [0, 1]) + # test fitting only even polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = cheb.chebfit(x, y, 4) + assert_almost_equal(cheb.chebval(x, coef1), y) + coef2 = cheb.chebfit(x, y, [0, 2, 4]) + assert_almost_equal(cheb.chebval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestInterpolate: + + def f(self, x): + return x * (x - 1) * (x - 2) + + def test_raises(self): + assert_raises(ValueError, cheb.chebinterpolate, self.f, -1) + assert_raises(TypeError, cheb.chebinterpolate, self.f, 10.) + + def test_dimensions(self): + for deg in range(1, 5): + assert_(cheb.chebinterpolate(self.f, deg).shape == (deg + 1,)) + + def test_approximation(self): + + def powx(x, p): + return x**p + + x = np.linspace(-1, 1, 10) + for deg in range(0, 10): + for p in range(0, deg + 1): + c = cheb.chebinterpolate(powx, deg, (p,)) + assert_almost_equal(cheb.chebval(x, c), powx(x, p), decimal=12) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, cheb.chebcompanion, []) + assert_raises(ValueError, cheb.chebcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(cheb.chebcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(cheb.chebcompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = cheb.chebgauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = cheb.chebvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.pi + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_chebfromroots(self): + res = cheb.chebfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + tgt = [0]*i + [1] + res = cheb.chebfromroots(roots)*2**(i-1) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebroots(self): + assert_almost_equal(cheb.chebroots([1]), []) + assert_almost_equal(cheb.chebroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = cheb.chebroots(cheb.chebfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_chebtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, cheb.chebtrim, coef, -1) + + # Test results + assert_equal(cheb.chebtrim(coef), coef[:-1]) + assert_equal(cheb.chebtrim(coef, 1), coef[:-3]) + assert_equal(cheb.chebtrim(coef, 2), [0]) + + def test_chebline(self): + assert_equal(cheb.chebline(3, 4), [3, 4]) + + def test_cheb2poly(self): + for i in range(10): + assert_almost_equal(cheb.cheb2poly([0]*i + [1]), Tlist[i]) + + def test_poly2cheb(self): + for i in range(10): + assert_almost_equal(cheb.poly2cheb(Tlist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-1, 1, 11)[1:-1] + tgt = 1./(np.sqrt(1 + x) * np.sqrt(1 - x)) + res = cheb.chebweight(x) + assert_almost_equal(res, tgt) + + def test_chebpts1(self): + #test exceptions + assert_raises(ValueError, cheb.chebpts1, 1.5) + assert_raises(ValueError, cheb.chebpts1, 0) + + #test points + tgt = [0] + assert_almost_equal(cheb.chebpts1(1), tgt) + tgt = [-0.70710678118654746, 0.70710678118654746] + assert_almost_equal(cheb.chebpts1(2), tgt) + tgt = [-0.86602540378443871, 0, 0.86602540378443871] + assert_almost_equal(cheb.chebpts1(3), tgt) + tgt = [-0.9238795325, -0.3826834323, 0.3826834323, 0.9238795325] + assert_almost_equal(cheb.chebpts1(4), tgt) + + def test_chebpts2(self): + #test exceptions + assert_raises(ValueError, cheb.chebpts2, 1.5) + assert_raises(ValueError, cheb.chebpts2, 1) + + #test points + tgt = [-1, 1] + assert_almost_equal(cheb.chebpts2(2), tgt) + tgt = [-1, 0, 1] + assert_almost_equal(cheb.chebpts2(3), tgt) + tgt = [-1, -0.5, .5, 1] + assert_almost_equal(cheb.chebpts2(4), tgt) + tgt = [-1.0, -0.707106781187, 0, 0.707106781187, 1.0] + assert_almost_equal(cheb.chebpts2(5), tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_classes.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_classes.py new file mode 100644 index 0000000000000000000000000000000000000000..6322062f29ece2f52754ac7aedf2591b3a983709 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_classes.py @@ -0,0 +1,600 @@ +"""Test inter-conversion of different polynomial classes. + +This tests the convert and cast methods of all the polynomial classes. + +""" +import operator as op +from numbers import Number + +import pytest +import numpy as np +from numpy.polynomial import ( + Polynomial, Legendre, Chebyshev, Laguerre, Hermite, HermiteE) +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) +from numpy.polynomial.polyutils import RankWarning + +# +# fixtures +# + +classes = ( + Polynomial, Legendre, Chebyshev, Laguerre, + Hermite, HermiteE + ) +classids = tuple(cls.__name__ for cls in classes) + +@pytest.fixture(params=classes, ids=classids) +def Poly(request): + return request.param + +# +# helper functions +# +random = np.random.random + + +def assert_poly_almost_equal(p1, p2, msg=""): + try: + assert_(np.all(p1.domain == p2.domain)) + assert_(np.all(p1.window == p2.window)) + assert_almost_equal(p1.coef, p2.coef) + except AssertionError: + msg = f"Result: {p1}\nTarget: {p2}" + raise AssertionError(msg) + + +# +# Test conversion methods that depend on combinations of two classes. +# + +Poly1 = Poly +Poly2 = Poly + + +def test_conversion(Poly1, Poly2): + x = np.linspace(0, 1, 10) + coef = random((3,)) + + d1 = Poly1.domain + random((2,))*.25 + w1 = Poly1.window + random((2,))*.25 + p1 = Poly1(coef, domain=d1, window=w1) + + d2 = Poly2.domain + random((2,))*.25 + w2 = Poly2.window + random((2,))*.25 + p2 = p1.convert(kind=Poly2, domain=d2, window=w2) + + assert_almost_equal(p2.domain, d2) + assert_almost_equal(p2.window, w2) + assert_almost_equal(p2(x), p1(x)) + + +def test_cast(Poly1, Poly2): + x = np.linspace(0, 1, 10) + coef = random((3,)) + + d1 = Poly1.domain + random((2,))*.25 + w1 = Poly1.window + random((2,))*.25 + p1 = Poly1(coef, domain=d1, window=w1) + + d2 = Poly2.domain + random((2,))*.25 + w2 = Poly2.window + random((2,))*.25 + p2 = Poly2.cast(p1, domain=d2, window=w2) + + assert_almost_equal(p2.domain, d2) + assert_almost_equal(p2.window, w2) + assert_almost_equal(p2(x), p1(x)) + + +# +# test methods that depend on one class +# + + +def test_identity(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + x = np.linspace(d[0], d[1], 11) + p = Poly.identity(domain=d, window=w) + assert_equal(p.domain, d) + assert_equal(p.window, w) + assert_almost_equal(p(x), x) + + +def test_basis(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly.basis(5, domain=d, window=w) + assert_equal(p.domain, d) + assert_equal(p.window, w) + assert_equal(p.coef, [0]*5 + [1]) + + +def test_fromroots(Poly): + # check that requested roots are zeros of a polynomial + # of correct degree, domain, and window. + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + r = random((5,)) + p1 = Poly.fromroots(r, domain=d, window=w) + assert_equal(p1.degree(), len(r)) + assert_equal(p1.domain, d) + assert_equal(p1.window, w) + assert_almost_equal(p1(r), 0) + + # check that polynomial is monic + pdom = Polynomial.domain + pwin = Polynomial.window + p2 = Polynomial.cast(p1, domain=pdom, window=pwin) + assert_almost_equal(p2.coef[-1], 1) + + +def test_bad_conditioned_fit(Poly): + + x = [0., 0., 1.] + y = [1., 2., 3.] + + # check RankWarning is raised + with pytest.warns(RankWarning) as record: + Poly.fit(x, y, 2) + assert record[0].message.args[0] == "The fit may be poorly conditioned" + + +def test_fit(Poly): + + def f(x): + return x*(x - 1)*(x - 2) + x = np.linspace(0, 3) + y = f(x) + + # check default value of domain and window + p = Poly.fit(x, y, 3) + assert_almost_equal(p.domain, [0, 3]) + assert_almost_equal(p(x), y) + assert_equal(p.degree(), 3) + + # check with given domains and window + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly.fit(x, y, 3, domain=d, window=w) + assert_almost_equal(p(x), y) + assert_almost_equal(p.domain, d) + assert_almost_equal(p.window, w) + p = Poly.fit(x, y, [0, 1, 2, 3], domain=d, window=w) + assert_almost_equal(p(x), y) + assert_almost_equal(p.domain, d) + assert_almost_equal(p.window, w) + + # check with class domain default + p = Poly.fit(x, y, 3, []) + assert_equal(p.domain, Poly.domain) + assert_equal(p.window, Poly.window) + p = Poly.fit(x, y, [0, 1, 2, 3], []) + assert_equal(p.domain, Poly.domain) + assert_equal(p.window, Poly.window) + + # check that fit accepts weights. + w = np.zeros_like(x) + z = y + random(y.shape)*.25 + w[::2] = 1 + p1 = Poly.fit(x[::2], z[::2], 3) + p2 = Poly.fit(x, z, 3, w=w) + p3 = Poly.fit(x, z, [0, 1, 2, 3], w=w) + assert_almost_equal(p1(x), p2(x)) + assert_almost_equal(p2(x), p3(x)) + + +def test_equal(Poly): + p1 = Poly([1, 2, 3], domain=[0, 1], window=[2, 3]) + p2 = Poly([1, 1, 1], domain=[0, 1], window=[2, 3]) + p3 = Poly([1, 2, 3], domain=[1, 2], window=[2, 3]) + p4 = Poly([1, 2, 3], domain=[0, 1], window=[1, 2]) + assert_(p1 == p1) + assert_(not p1 == p2) + assert_(not p1 == p3) + assert_(not p1 == p4) + + +def test_not_equal(Poly): + p1 = Poly([1, 2, 3], domain=[0, 1], window=[2, 3]) + p2 = Poly([1, 1, 1], domain=[0, 1], window=[2, 3]) + p3 = Poly([1, 2, 3], domain=[1, 2], window=[2, 3]) + p4 = Poly([1, 2, 3], domain=[0, 1], window=[1, 2]) + assert_(not p1 != p1) + assert_(p1 != p2) + assert_(p1 != p3) + assert_(p1 != p4) + + +def test_add(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 + p2 + assert_poly_almost_equal(p2 + p1, p3) + assert_poly_almost_equal(p1 + c2, p3) + assert_poly_almost_equal(c2 + p1, p3) + assert_poly_almost_equal(p1 + tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) + p1, p3) + assert_poly_almost_equal(p1 + np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) + p1, p3) + assert_raises(TypeError, op.add, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.add, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.add, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.add, p1, Polynomial([0])) + + +def test_sub(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 - p2 + assert_poly_almost_equal(p2 - p1, -p3) + assert_poly_almost_equal(p1 - c2, p3) + assert_poly_almost_equal(c2 - p1, -p3) + assert_poly_almost_equal(p1 - tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) - p1, -p3) + assert_poly_almost_equal(p1 - np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) - p1, -p3) + assert_raises(TypeError, op.sub, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.sub, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.sub, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.sub, p1, Polynomial([0])) + + +def test_mul(Poly): + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = p1 * p2 + assert_poly_almost_equal(p2 * p1, p3) + assert_poly_almost_equal(p1 * c2, p3) + assert_poly_almost_equal(c2 * p1, p3) + assert_poly_almost_equal(p1 * tuple(c2), p3) + assert_poly_almost_equal(tuple(c2) * p1, p3) + assert_poly_almost_equal(p1 * np.array(c2), p3) + assert_poly_almost_equal(np.array(c2) * p1, p3) + assert_poly_almost_equal(p1 * 2, p1 * Poly([2])) + assert_poly_almost_equal(2 * p1, p1 * Poly([2])) + assert_raises(TypeError, op.mul, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.mul, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.mul, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.mul, p1, Polynomial([0])) + + +def test_floordiv(Poly): + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + assert_poly_almost_equal(p4 // p2, p1) + assert_poly_almost_equal(p4 // c2, p1) + assert_poly_almost_equal(c4 // p2, p1) + assert_poly_almost_equal(p4 // tuple(c2), p1) + assert_poly_almost_equal(tuple(c4) // p2, p1) + assert_poly_almost_equal(p4 // np.array(c2), p1) + assert_poly_almost_equal(np.array(c4) // p2, p1) + assert_poly_almost_equal(2 // p2, Poly([0])) + assert_poly_almost_equal(p2 // 2, 0.5*p2) + assert_raises( + TypeError, op.floordiv, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises( + TypeError, op.floordiv, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.floordiv, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.floordiv, p1, Polynomial([0])) + + +def test_truediv(Poly): + # true division is valid only if the denominator is a Number and + # not a python bool. + p1 = Poly([1,2,3]) + p2 = p1 * 5 + + for stype in np.ScalarType: + if not issubclass(stype, Number) or issubclass(stype, bool): + continue + s = stype(5) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for stype in (int, float): + s = stype(5) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for stype in [complex]: + s = stype(5, 0) + assert_poly_almost_equal(op.truediv(p2, s), p1) + assert_raises(TypeError, op.truediv, s, p2) + for s in [tuple(), list(), dict(), bool(), np.array([1])]: + assert_raises(TypeError, op.truediv, p2, s) + assert_raises(TypeError, op.truediv, s, p2) + for ptype in classes: + assert_raises(TypeError, op.truediv, p2, ptype(1)) + + +def test_mod(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + assert_poly_almost_equal(p4 % p2, p3) + assert_poly_almost_equal(p4 % c2, p3) + assert_poly_almost_equal(c4 % p2, p3) + assert_poly_almost_equal(p4 % tuple(c2), p3) + assert_poly_almost_equal(tuple(c4) % p2, p3) + assert_poly_almost_equal(p4 % np.array(c2), p3) + assert_poly_almost_equal(np.array(c4) % p2, p3) + assert_poly_almost_equal(2 % p2, Poly([2])) + assert_poly_almost_equal(p2 % 2, Poly([0])) + assert_raises(TypeError, op.mod, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, op.mod, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, op.mod, p1, Chebyshev([0])) + else: + assert_raises(TypeError, op.mod, p1, Polynomial([0])) + + +def test_divmod(Poly): + # This checks commutation, not numerical correctness + c1 = list(random((4,)) + .5) + c2 = list(random((3,)) + .5) + c3 = list(random((2,)) + .5) + p1 = Poly(c1) + p2 = Poly(c2) + p3 = Poly(c3) + p4 = p1 * p2 + p3 + c4 = list(p4.coef) + quo, rem = divmod(p4, p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, c2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(c4, p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, tuple(c2)) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(tuple(c4), p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p4, np.array(c2)) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(np.array(c4), p2) + assert_poly_almost_equal(quo, p1) + assert_poly_almost_equal(rem, p3) + quo, rem = divmod(p2, 2) + assert_poly_almost_equal(quo, 0.5*p2) + assert_poly_almost_equal(rem, Poly([0])) + quo, rem = divmod(2, p2) + assert_poly_almost_equal(quo, Poly([0])) + assert_poly_almost_equal(rem, Poly([2])) + assert_raises(TypeError, divmod, p1, Poly([0], domain=Poly.domain + 1)) + assert_raises(TypeError, divmod, p1, Poly([0], window=Poly.window + 1)) + if Poly is Polynomial: + assert_raises(TypeError, divmod, p1, Chebyshev([0])) + else: + assert_raises(TypeError, divmod, p1, Polynomial([0])) + + +def test_roots(Poly): + d = Poly.domain * 1.25 + .25 + w = Poly.window + tgt = np.linspace(d[0], d[1], 5) + res = np.sort(Poly.fromroots(tgt, domain=d, window=w).roots()) + assert_almost_equal(res, tgt) + # default domain and window + res = np.sort(Poly.fromroots(tgt).roots()) + assert_almost_equal(res, tgt) + + +def test_degree(Poly): + p = Poly.basis(5) + assert_equal(p.degree(), 5) + + +def test_copy(Poly): + p1 = Poly.basis(5) + p2 = p1.copy() + assert_(p1 == p2) + assert_(p1 is not p2) + assert_(p1.coef is not p2.coef) + assert_(p1.domain is not p2.domain) + assert_(p1.window is not p2.window) + + +def test_integ(Poly): + P = Polynomial + # Check defaults + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ()) + p2 = P.cast(p0.integ(2)) + assert_poly_almost_equal(p1, P([0, 2, 3, 4])) + assert_poly_almost_equal(p2, P([0, 0, 1, 1, 1])) + # Check with k + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ(k=1)) + p2 = P.cast(p0.integ(2, k=[1, 1])) + assert_poly_almost_equal(p1, P([1, 2, 3, 4])) + assert_poly_almost_equal(p2, P([1, 1, 1, 1, 1])) + # Check with lbnd + p0 = Poly.cast(P([1*2, 2*3, 3*4])) + p1 = P.cast(p0.integ(lbnd=1)) + p2 = P.cast(p0.integ(2, lbnd=1)) + assert_poly_almost_equal(p1, P([-9, 2, 3, 4])) + assert_poly_almost_equal(p2, P([6, -9, 1, 1, 1])) + # Check scaling + d = 2*Poly.domain + p0 = Poly.cast(P([1*2, 2*3, 3*4]), domain=d) + p1 = P.cast(p0.integ()) + p2 = P.cast(p0.integ(2)) + assert_poly_almost_equal(p1, P([0, 2, 3, 4])) + assert_poly_almost_equal(p2, P([0, 0, 1, 1, 1])) + + +def test_deriv(Poly): + # Check that the derivative is the inverse of integration. It is + # assumes that the integration has been checked elsewhere. + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p1 = Poly([1, 2, 3], domain=d, window=w) + p2 = p1.integ(2, k=[1, 2]) + p3 = p1.integ(1, k=[1]) + assert_almost_equal(p2.deriv(1).coef, p3.coef) + assert_almost_equal(p2.deriv(2).coef, p1.coef) + # default domain and window + p1 = Poly([1, 2, 3]) + p2 = p1.integ(2, k=[1, 2]) + p3 = p1.integ(1, k=[1]) + assert_almost_equal(p2.deriv(1).coef, p3.coef) + assert_almost_equal(p2.deriv(2).coef, p1.coef) + + +def test_linspace(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + p = Poly([1, 2, 3], domain=d, window=w) + # check default domain + xtgt = np.linspace(d[0], d[1], 20) + ytgt = p(xtgt) + xres, yres = p.linspace(20) + assert_almost_equal(xres, xtgt) + assert_almost_equal(yres, ytgt) + # check specified domain + xtgt = np.linspace(0, 2, 20) + ytgt = p(xtgt) + xres, yres = p.linspace(20, domain=[0, 2]) + assert_almost_equal(xres, xtgt) + assert_almost_equal(yres, ytgt) + + +def test_pow(Poly): + d = Poly.domain + random((2,))*.25 + w = Poly.window + random((2,))*.25 + tgt = Poly([1], domain=d, window=w) + tst = Poly([1, 2, 3], domain=d, window=w) + for i in range(5): + assert_poly_almost_equal(tst**i, tgt) + tgt = tgt * tst + # default domain and window + tgt = Poly([1]) + tst = Poly([1, 2, 3]) + for i in range(5): + assert_poly_almost_equal(tst**i, tgt) + tgt = tgt * tst + # check error for invalid powers + assert_raises(ValueError, op.pow, tgt, 1.5) + assert_raises(ValueError, op.pow, tgt, -1) + + +def test_call(Poly): + P = Polynomial + d = Poly.domain + x = np.linspace(d[0], d[1], 11) + + # Check defaults + p = Poly.cast(P([1, 2, 3])) + tgt = 1 + x*(2 + 3*x) + res = p(x) + assert_almost_equal(res, tgt) + + +def test_cutdeg(Poly): + p = Poly([1, 2, 3]) + assert_raises(ValueError, p.cutdeg, .5) + assert_raises(ValueError, p.cutdeg, -1) + assert_equal(len(p.cutdeg(3)), 3) + assert_equal(len(p.cutdeg(2)), 3) + assert_equal(len(p.cutdeg(1)), 2) + assert_equal(len(p.cutdeg(0)), 1) + + +def test_truncate(Poly): + p = Poly([1, 2, 3]) + assert_raises(ValueError, p.truncate, .5) + assert_raises(ValueError, p.truncate, 0) + assert_equal(len(p.truncate(4)), 3) + assert_equal(len(p.truncate(3)), 3) + assert_equal(len(p.truncate(2)), 2) + assert_equal(len(p.truncate(1)), 1) + + +def test_trim(Poly): + c = [1, 1e-6, 1e-12, 0] + p = Poly(c) + assert_equal(p.trim().coef, c[:3]) + assert_equal(p.trim(1e-10).coef, c[:2]) + assert_equal(p.trim(1e-5).coef, c[:1]) + + +def test_mapparms(Poly): + # check with defaults. Should be identity. + d = Poly.domain + w = Poly.window + p = Poly([1], domain=d, window=w) + assert_almost_equal([0, 1], p.mapparms()) + # + w = 2*d + 1 + p = Poly([1], domain=d, window=w) + assert_almost_equal([1, 2], p.mapparms()) + + +def test_ufunc_override(Poly): + p = Poly([1, 2, 3]) + x = np.ones(3) + assert_raises(TypeError, np.add, p, x) + assert_raises(TypeError, np.add, x, p) + + +# +# Test class method that only exists for some classes +# + + +class TestInterpolate: + + def f(self, x): + return x * (x - 1) * (x - 2) + + def test_raises(self): + assert_raises(ValueError, Chebyshev.interpolate, self.f, -1) + assert_raises(TypeError, Chebyshev.interpolate, self.f, 10.) + + def test_dimensions(self): + for deg in range(1, 5): + assert_(Chebyshev.interpolate(self.f, deg).degree() == deg) + + def test_approximation(self): + + def powx(x, p): + return x**p + + x = np.linspace(0, 2, 10) + for deg in range(0, 10): + for t in range(0, deg + 1): + p = Chebyshev.interpolate(powx, deg, domain=[0, 2], args=(t,)) + assert_almost_equal(p(x), powx(x, t), decimal=11) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite.py new file mode 100644 index 0000000000000000000000000000000000000000..53ee0844e3c58456807bfd7828bdb9cf58f8ed76 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite.py @@ -0,0 +1,555 @@ +"""Tests for hermite module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.hermite as herm +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +H0 = np.array([1]) +H1 = np.array([0, 2]) +H2 = np.array([-2, 0, 4]) +H3 = np.array([0, -12, 0, 8]) +H4 = np.array([12, 0, -48, 0, 16]) +H5 = np.array([0, 120, 0, -160, 0, 32]) +H6 = np.array([-120, 0, 720, 0, -480, 0, 64]) +H7 = np.array([0, -1680, 0, 3360, 0, -1344, 0, 128]) +H8 = np.array([1680, 0, -13440, 0, 13440, 0, -3584, 0, 256]) +H9 = np.array([0, 30240, 0, -80640, 0, 48384, 0, -9216, 0, 512]) + +Hlist = [H0, H1, H2, H3, H4, H5, H6, H7, H8, H9] + + +def trim(x): + return herm.hermtrim(x, tol=1e-6) + + +class TestConstants: + + def test_hermdomain(self): + assert_equal(herm.hermdomain, [-1, 1]) + + def test_hermzero(self): + assert_equal(herm.hermzero, [0]) + + def test_hermone(self): + assert_equal(herm.hermone, [1]) + + def test_hermx(self): + assert_equal(herm.hermx, [0, .5]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_hermadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = herm.hermadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = herm.hermsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermmulx(self): + assert_equal(herm.hermmulx([0]), [0]) + assert_equal(herm.hermmulx([1]), [0, .5]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i, 0, .5] + assert_equal(herm.hermmulx(ser), tgt) + + def test_hermmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = herm.hermval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = herm.hermval(self.x, pol2) + pol3 = herm.hermmul(pol1, pol2) + val3 = herm.hermval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_hermdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = herm.hermadd(ci, cj) + quo, rem = herm.hermdiv(tgt, ci) + res = herm.hermadd(herm.hermmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(herm.hermmul, [c]*j, np.array([1])) + res = herm.hermpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2.5, 1., .75]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_hermval(self): + #check empty input + assert_equal(herm.hermval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Hlist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = herm.hermval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(herm.hermval(x, [1]).shape, dims) + assert_equal(herm.hermval(x, [1, 0]).shape, dims) + assert_equal(herm.hermval(x, [1, 0, 0]).shape, dims) + + def test_hermval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herm.hermval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = herm.hermval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_hermval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herm.hermval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = herm.hermval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_hermgrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = herm.hermgrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermgrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_hermgrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = herm.hermgrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herm.hermgrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_hermint(self): + # check exceptions + assert_raises(TypeError, herm.hermint, [0], .5) + assert_raises(ValueError, herm.hermint, [0], -1) + assert_raises(ValueError, herm.hermint, [0], 1, [0, 0]) + assert_raises(ValueError, herm.hermint, [0], lbnd=[0]) + assert_raises(ValueError, herm.hermint, [0], scl=[0]) + assert_raises(TypeError, herm.hermint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = herm.hermint([0], m=i, k=k) + assert_almost_equal(res, [0, .5]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i]) + res = herm.herm2poly(hermint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(herm.hermval(-1, hermint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + hermpol = herm.poly2herm(pol) + hermint = herm.hermint(hermpol, m=1, k=[i], scl=2) + res = herm.herm2poly(hermint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1) + res = herm.hermint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k]) + res = herm.hermint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k], lbnd=-1) + res = herm.hermint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herm.hermint(tgt, m=1, k=[k], scl=2) + res = herm.hermint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herm.hermint(c) for c in c2d.T]).T + res = herm.hermint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermint(c) for c in c2d]) + res = herm.hermint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermint(c, k=3) for c in c2d]) + res = herm.hermint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_hermder(self): + # check exceptions + assert_raises(TypeError, herm.hermder, [0], .5) + assert_raises(ValueError, herm.hermder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = herm.hermder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herm.hermder(herm.hermint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herm.hermder(herm.hermint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herm.hermder(c) for c in c2d.T]).T + res = herm.hermder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herm.hermder(c) for c in c2d]) + res = herm.hermder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_hermvander(self): + # check for 1d x + x = np.arange(3) + v = herm.hermvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herm.hermval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = herm.hermvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herm.hermval(x, coef)) + + def test_hermvander2d(self): + # also tests hermval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = herm.hermvander2d(x1, x2, [1, 2]) + tgt = herm.hermval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herm.hermvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_hermvander3d(self): + # also tests hermval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = herm.hermvander3d(x1, x2, x3, [1, 2, 3]) + tgt = herm.hermval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herm.hermvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_hermfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, herm.hermfit, [1], [1], -1) + assert_raises(TypeError, herm.hermfit, [[1]], [1], 0) + assert_raises(TypeError, herm.hermfit, [], [1], 0) + assert_raises(TypeError, herm.hermfit, [1], [[[1]]], 0) + assert_raises(TypeError, herm.hermfit, [1, 2], [1], 0) + assert_raises(TypeError, herm.hermfit, [1], [1, 2], 0) + assert_raises(TypeError, herm.hermfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, herm.hermfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, herm.hermfit, [1], [1], [-1,]) + assert_raises(ValueError, herm.hermfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, herm.hermfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = herm.hermfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(herm.hermval(x, coef3), y) + coef3 = herm.hermfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(herm.hermval(x, coef3), y) + # + coef4 = herm.hermfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + coef4 = herm.hermfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = herm.hermfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(herm.hermval(x, coef4), y) + # + coef2d = herm.hermfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = herm.hermfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = herm.hermfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = herm.hermfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = herm.hermfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = herm.hermfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(herm.hermfit(x, x, 1), [0, .5]) + assert_almost_equal(herm.hermfit(x, x, [0, 1]), [0, .5]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = herm.hermfit(x, y, 4) + assert_almost_equal(herm.hermval(x, coef1), y) + coef2 = herm.hermfit(x, y, [0, 2, 4]) + assert_almost_equal(herm.hermval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, herm.hermcompanion, []) + assert_raises(ValueError, herm.hermcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(herm.hermcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(herm.hermcompanion([1, 2])[0, 0] == -.25) + + +class TestGauss: + + def test_100(self): + x, w = herm.hermgauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = herm.hermvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.sqrt(np.pi) + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_hermfromroots(self): + res = herm.hermfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = herm.hermfromroots(roots) + res = herm.hermval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(herm.herm2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_hermroots(self): + assert_almost_equal(herm.hermroots([1]), []) + assert_almost_equal(herm.hermroots([1, 1]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = herm.hermroots(herm.hermfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, herm.hermtrim, coef, -1) + + # Test results + assert_equal(herm.hermtrim(coef), coef[:-1]) + assert_equal(herm.hermtrim(coef, 1), coef[:-3]) + assert_equal(herm.hermtrim(coef, 2), [0]) + + def test_hermline(self): + assert_equal(herm.hermline(3, 4), [3, 2]) + + def test_herm2poly(self): + for i in range(10): + assert_almost_equal(herm.herm2poly([0]*i + [1]), Hlist[i]) + + def test_poly2herm(self): + for i in range(10): + assert_almost_equal(herm.poly2herm(Hlist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-5, 5, 11) + tgt = np.exp(-x**2) + res = herm.hermweight(x) + assert_almost_equal(res, tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite_e.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite_e.py new file mode 100644 index 0000000000000000000000000000000000000000..2d262a3306222bd79f682b09763b0bd2b90ba8fe --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_hermite_e.py @@ -0,0 +1,556 @@ +"""Tests for hermite_e module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.hermite_e as herme +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +He0 = np.array([1]) +He1 = np.array([0, 1]) +He2 = np.array([-1, 0, 1]) +He3 = np.array([0, -3, 0, 1]) +He4 = np.array([3, 0, -6, 0, 1]) +He5 = np.array([0, 15, 0, -10, 0, 1]) +He6 = np.array([-15, 0, 45, 0, -15, 0, 1]) +He7 = np.array([0, -105, 0, 105, 0, -21, 0, 1]) +He8 = np.array([105, 0, -420, 0, 210, 0, -28, 0, 1]) +He9 = np.array([0, 945, 0, -1260, 0, 378, 0, -36, 0, 1]) + +Helist = [He0, He1, He2, He3, He4, He5, He6, He7, He8, He9] + + +def trim(x): + return herme.hermetrim(x, tol=1e-6) + + +class TestConstants: + + def test_hermedomain(self): + assert_equal(herme.hermedomain, [-1, 1]) + + def test_hermezero(self): + assert_equal(herme.hermezero, [0]) + + def test_hermeone(self): + assert_equal(herme.hermeone, [1]) + + def test_hermex(self): + assert_equal(herme.hermex, [0, 1]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_hermeadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = herme.hermeadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermesub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = herme.hermesub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermemulx(self): + assert_equal(herme.hermemulx([0]), [0]) + assert_equal(herme.hermemulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i, 0, 1] + assert_equal(herme.hermemulx(ser), tgt) + + def test_hermemul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = herme.hermeval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = herme.hermeval(self.x, pol2) + pol3 = herme.hermemul(pol1, pol2) + val3 = herme.hermeval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_hermediv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = herme.hermeadd(ci, cj) + quo, rem = herme.hermediv(tgt, ci) + res = herme.hermeadd(herme.hermemul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_hermepow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(herme.hermemul, [c]*j, np.array([1])) + res = herme.hermepow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([4., 2., 3.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_hermeval(self): + #check empty input + assert_equal(herme.hermeval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Helist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = herme.hermeval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(herme.hermeval(x, [1]).shape, dims) + assert_equal(herme.hermeval(x, [1, 0]).shape, dims) + assert_equal(herme.hermeval(x, [1, 0, 0]).shape, dims) + + def test_hermeval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herme.hermeval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = herme.hermeval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermeval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_hermeval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, herme.hermeval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = herme.hermeval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermeval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_hermegrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = herme.hermegrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermegrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_hermegrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = herme.hermegrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = herme.hermegrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_hermeint(self): + # check exceptions + assert_raises(TypeError, herme.hermeint, [0], .5) + assert_raises(ValueError, herme.hermeint, [0], -1) + assert_raises(ValueError, herme.hermeint, [0], 1, [0, 0]) + assert_raises(ValueError, herme.hermeint, [0], lbnd=[0]) + assert_raises(ValueError, herme.hermeint, [0], scl=[0]) + assert_raises(TypeError, herme.hermeint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = herme.hermeint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i]) + res = herme.herme2poly(hermeint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i], lbnd=-1) + assert_almost_equal(herme.hermeval(-1, hermeint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + hermepol = herme.poly2herme(pol) + hermeint = herme.hermeint(hermepol, m=1, k=[i], scl=2) + res = herme.herme2poly(hermeint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1) + res = herme.hermeint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k]) + res = herme.hermeint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k], lbnd=-1) + res = herme.hermeint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = herme.hermeint(tgt, m=1, k=[k], scl=2) + res = herme.hermeint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermeint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herme.hermeint(c) for c in c2d.T]).T + res = herme.hermeint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeint(c) for c in c2d]) + res = herme.hermeint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeint(c, k=3) for c in c2d]) + res = herme.hermeint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_hermeder(self): + # check exceptions + assert_raises(TypeError, herme.hermeder, [0], .5) + assert_raises(ValueError, herme.hermeder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = herme.hermeder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herme.hermeder(herme.hermeint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = herme.hermeder( + herme.hermeint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermeder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([herme.hermeder(c) for c in c2d.T]).T + res = herme.hermeder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([herme.hermeder(c) for c in c2d]) + res = herme.hermeder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_hermevander(self): + # check for 1d x + x = np.arange(3) + v = herme.hermevander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herme.hermeval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = herme.hermevander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], herme.hermeval(x, coef)) + + def test_hermevander2d(self): + # also tests hermeval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = herme.hermevander2d(x1, x2, [1, 2]) + tgt = herme.hermeval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herme.hermevander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_hermevander3d(self): + # also tests hermeval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = herme.hermevander3d(x1, x2, x3, [1, 2, 3]) + tgt = herme.hermeval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = herme.hermevander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_hermefit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, herme.hermefit, [1], [1], -1) + assert_raises(TypeError, herme.hermefit, [[1]], [1], 0) + assert_raises(TypeError, herme.hermefit, [], [1], 0) + assert_raises(TypeError, herme.hermefit, [1], [[[1]]], 0) + assert_raises(TypeError, herme.hermefit, [1, 2], [1], 0) + assert_raises(TypeError, herme.hermefit, [1], [1, 2], 0) + assert_raises(TypeError, herme.hermefit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, herme.hermefit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, herme.hermefit, [1], [1], [-1,]) + assert_raises(ValueError, herme.hermefit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, herme.hermefit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = herme.hermefit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(herme.hermeval(x, coef3), y) + coef3 = herme.hermefit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(herme.hermeval(x, coef3), y) + # + coef4 = herme.hermefit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + coef4 = herme.hermefit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = herme.hermefit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(herme.hermeval(x, coef4), y) + # + coef2d = herme.hermefit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = herme.hermefit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = herme.hermefit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = herme.hermefit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = herme.hermefit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = herme.hermefit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(herme.hermefit(x, x, 1), [0, 1]) + assert_almost_equal(herme.hermefit(x, x, [0, 1]), [0, 1]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = herme.hermefit(x, y, 4) + assert_almost_equal(herme.hermeval(x, coef1), y) + coef2 = herme.hermefit(x, y, [0, 2, 4]) + assert_almost_equal(herme.hermeval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, herme.hermecompanion, []) + assert_raises(ValueError, herme.hermecompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(herme.hermecompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(herme.hermecompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = herme.hermegauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = herme.hermevander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = np.sqrt(2*np.pi) + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_hermefromroots(self): + res = herme.hermefromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = herme.hermefromroots(roots) + res = herme.hermeval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(herme.herme2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_hermeroots(self): + assert_almost_equal(herme.hermeroots([1]), []) + assert_almost_equal(herme.hermeroots([1, 1]), [-1]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = herme.hermeroots(herme.hermefromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_hermetrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, herme.hermetrim, coef, -1) + + # Test results + assert_equal(herme.hermetrim(coef), coef[:-1]) + assert_equal(herme.hermetrim(coef, 1), coef[:-3]) + assert_equal(herme.hermetrim(coef, 2), [0]) + + def test_hermeline(self): + assert_equal(herme.hermeline(3, 4), [3, 4]) + + def test_herme2poly(self): + for i in range(10): + assert_almost_equal(herme.herme2poly([0]*i + [1]), Helist[i]) + + def test_poly2herme(self): + for i in range(10): + assert_almost_equal(herme.poly2herme(Helist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-5, 5, 11) + tgt = np.exp(-.5*x**2) + res = herme.hermeweight(x) + assert_almost_equal(res, tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_laguerre.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_laguerre.py new file mode 100644 index 0000000000000000000000000000000000000000..227ef3c5576dd666e2eb76576eb260d5ba48cb0e --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_laguerre.py @@ -0,0 +1,537 @@ +"""Tests for laguerre module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.laguerre as lag +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +L0 = np.array([1])/1 +L1 = np.array([1, -1])/1 +L2 = np.array([2, -4, 1])/2 +L3 = np.array([6, -18, 9, -1])/6 +L4 = np.array([24, -96, 72, -16, 1])/24 +L5 = np.array([120, -600, 600, -200, 25, -1])/120 +L6 = np.array([720, -4320, 5400, -2400, 450, -36, 1])/720 + +Llist = [L0, L1, L2, L3, L4, L5, L6] + + +def trim(x): + return lag.lagtrim(x, tol=1e-6) + + +class TestConstants: + + def test_lagdomain(self): + assert_equal(lag.lagdomain, [0, 1]) + + def test_lagzero(self): + assert_equal(lag.lagzero, [0]) + + def test_lagone(self): + assert_equal(lag.lagone, [1]) + + def test_lagx(self): + assert_equal(lag.lagx, [1, -1]) + + +class TestArithmetic: + x = np.linspace(-3, 3, 100) + + def test_lagadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = lag.lagadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = lag.lagsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagmulx(self): + assert_equal(lag.lagmulx([0]), [0]) + assert_equal(lag.lagmulx([1]), [1, -1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [-i, 2*i + 1, -(i + 1)] + assert_almost_equal(lag.lagmulx(ser), tgt) + + def test_lagmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = lag.lagval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = lag.lagval(self.x, pol2) + pol3 = lag.lagmul(pol1, pol2) + val3 = lag.lagval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_lagdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = lag.lagadd(ci, cj) + quo, rem = lag.lagdiv(tgt, ci) + res = lag.lagadd(lag.lagmul(quo, ci), rem) + assert_almost_equal(trim(res), trim(tgt), err_msg=msg) + + def test_lagpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(lag.lagmul, [c]*j, np.array([1])) + res = lag.lagpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([9., -14., 6.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_lagval(self): + #check empty input + assert_equal(lag.lagval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Llist] + for i in range(7): + msg = f"At i={i}" + tgt = y[i] + res = lag.lagval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(lag.lagval(x, [1]).shape, dims) + assert_equal(lag.lagval(x, [1, 0]).shape, dims) + assert_equal(lag.lagval(x, [1, 0, 0]).shape, dims) + + def test_lagval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, lag.lagval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = lag.lagval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.lagval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_lagval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, lag.lagval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = lag.lagval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.lagval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_laggrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = lag.laggrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.laggrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_laggrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = lag.laggrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = lag.laggrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_lagint(self): + # check exceptions + assert_raises(TypeError, lag.lagint, [0], .5) + assert_raises(ValueError, lag.lagint, [0], -1) + assert_raises(ValueError, lag.lagint, [0], 1, [0, 0]) + assert_raises(ValueError, lag.lagint, [0], lbnd=[0]) + assert_raises(ValueError, lag.lagint, [0], scl=[0]) + assert_raises(TypeError, lag.lagint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = lag.lagint([0], m=i, k=k) + assert_almost_equal(res, [1, -1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i]) + res = lag.lag2poly(lagint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(lag.lagval(-1, lagint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + lagpol = lag.poly2lag(pol) + lagint = lag.lagint(lagpol, m=1, k=[i], scl=2) + res = lag.lag2poly(lagint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1) + res = lag.lagint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k]) + res = lag.lagint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k], lbnd=-1) + res = lag.lagint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = lag.lagint(tgt, m=1, k=[k], scl=2) + res = lag.lagint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([lag.lagint(c) for c in c2d.T]).T + res = lag.lagint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagint(c) for c in c2d]) + res = lag.lagint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagint(c, k=3) for c in c2d]) + res = lag.lagint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_lagder(self): + # check exceptions + assert_raises(TypeError, lag.lagder, [0], .5) + assert_raises(ValueError, lag.lagder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = lag.lagder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = lag.lagder(lag.lagint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = lag.lagder(lag.lagint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([lag.lagder(c) for c in c2d.T]).T + res = lag.lagder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([lag.lagder(c) for c in c2d]) + res = lag.lagder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_lagvander(self): + # check for 1d x + x = np.arange(3) + v = lag.lagvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], lag.lagval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = lag.lagvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], lag.lagval(x, coef)) + + def test_lagvander2d(self): + # also tests lagval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = lag.lagvander2d(x1, x2, [1, 2]) + tgt = lag.lagval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = lag.lagvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_lagvander3d(self): + # also tests lagval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = lag.lagvander3d(x1, x2, x3, [1, 2, 3]) + tgt = lag.lagval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = lag.lagvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + +class TestFitting: + + def test_lagfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + # Test exceptions + assert_raises(ValueError, lag.lagfit, [1], [1], -1) + assert_raises(TypeError, lag.lagfit, [[1]], [1], 0) + assert_raises(TypeError, lag.lagfit, [], [1], 0) + assert_raises(TypeError, lag.lagfit, [1], [[[1]]], 0) + assert_raises(TypeError, lag.lagfit, [1, 2], [1], 0) + assert_raises(TypeError, lag.lagfit, [1], [1, 2], 0) + assert_raises(TypeError, lag.lagfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, lag.lagfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, lag.lagfit, [1], [1], [-1,]) + assert_raises(ValueError, lag.lagfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, lag.lagfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = lag.lagfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(lag.lagval(x, coef3), y) + coef3 = lag.lagfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(lag.lagval(x, coef3), y) + # + coef4 = lag.lagfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(lag.lagval(x, coef4), y) + coef4 = lag.lagfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(lag.lagval(x, coef4), y) + # + coef2d = lag.lagfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = lag.lagfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = lag.lagfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = lag.lagfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = lag.lagfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = lag.lagfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(lag.lagfit(x, x, 1), [1, -1]) + assert_almost_equal(lag.lagfit(x, x, [0, 1]), [1, -1]) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, lag.lagcompanion, []) + assert_raises(ValueError, lag.lagcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(lag.lagcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(lag.lagcompanion([1, 2])[0, 0] == 1.5) + + +class TestGauss: + + def test_100(self): + x, w = lag.laggauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = lag.lagvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = 1.0 + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_lagfromroots(self): + res = lag.lagfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = lag.lagfromroots(roots) + res = lag.lagval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(lag.lag2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_lagroots(self): + assert_almost_equal(lag.lagroots([1]), []) + assert_almost_equal(lag.lagroots([0, 1]), [1]) + for i in range(2, 5): + tgt = np.linspace(0, 3, i) + res = lag.lagroots(lag.lagfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_lagtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, lag.lagtrim, coef, -1) + + # Test results + assert_equal(lag.lagtrim(coef), coef[:-1]) + assert_equal(lag.lagtrim(coef, 1), coef[:-3]) + assert_equal(lag.lagtrim(coef, 2), [0]) + + def test_lagline(self): + assert_equal(lag.lagline(3, 4), [7, -4]) + + def test_lag2poly(self): + for i in range(7): + assert_almost_equal(lag.lag2poly([0]*i + [1]), Llist[i]) + + def test_poly2lag(self): + for i in range(7): + assert_almost_equal(lag.poly2lag(Llist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(0, 10, 11) + tgt = np.exp(-x) + res = lag.lagweight(x) + assert_almost_equal(res, tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_legendre.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_legendre.py new file mode 100644 index 0000000000000000000000000000000000000000..92399c160ecb75fbb1f9a5a7f2bba0fe90d84a54 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_legendre.py @@ -0,0 +1,568 @@ +"""Tests for legendre module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.legendre as leg +from numpy.polynomial.polynomial import polyval +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + +L0 = np.array([1]) +L1 = np.array([0, 1]) +L2 = np.array([-1, 0, 3])/2 +L3 = np.array([0, -3, 0, 5])/2 +L4 = np.array([3, 0, -30, 0, 35])/8 +L5 = np.array([0, 15, 0, -70, 0, 63])/8 +L6 = np.array([-5, 0, 105, 0, -315, 0, 231])/16 +L7 = np.array([0, -35, 0, 315, 0, -693, 0, 429])/16 +L8 = np.array([35, 0, -1260, 0, 6930, 0, -12012, 0, 6435])/128 +L9 = np.array([0, 315, 0, -4620, 0, 18018, 0, -25740, 0, 12155])/128 + +Llist = [L0, L1, L2, L3, L4, L5, L6, L7, L8, L9] + + +def trim(x): + return leg.legtrim(x, tol=1e-6) + + +class TestConstants: + + def test_legdomain(self): + assert_equal(leg.legdomain, [-1, 1]) + + def test_legzero(self): + assert_equal(leg.legzero, [0]) + + def test_legone(self): + assert_equal(leg.legone, [1]) + + def test_legx(self): + assert_equal(leg.legx, [0, 1]) + + +class TestArithmetic: + x = np.linspace(-1, 1, 100) + + def test_legadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = leg.legadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legsub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = leg.legsub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legmulx(self): + assert_equal(leg.legmulx([0]), [0]) + assert_equal(leg.legmulx([1]), [0, 1]) + for i in range(1, 5): + tmp = 2*i + 1 + ser = [0]*i + [1] + tgt = [0]*(i - 1) + [i/tmp, 0, (i + 1)/tmp] + assert_equal(leg.legmulx(ser), tgt) + + def test_legmul(self): + # check values of result + for i in range(5): + pol1 = [0]*i + [1] + val1 = leg.legval(self.x, pol1) + for j in range(5): + msg = f"At i={i}, j={j}" + pol2 = [0]*j + [1] + val2 = leg.legval(self.x, pol2) + pol3 = leg.legmul(pol1, pol2) + val3 = leg.legval(self.x, pol3) + assert_(len(pol3) == i + j + 1, msg) + assert_almost_equal(val3, val1*val2, err_msg=msg) + + def test_legdiv(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1] + cj = [0]*j + [1] + tgt = leg.legadd(ci, cj) + quo, rem = leg.legdiv(tgt, ci) + res = leg.legadd(leg.legmul(quo, ci), rem) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_legpow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(leg.legmul, [c]*j, np.array([1])) + res = leg.legpow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([2., 2., 2.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = polyval(x, [1., 2., 3.]) + + def test_legval(self): + #check empty input + assert_equal(leg.legval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [polyval(x, c) for c in Llist] + for i in range(10): + msg = f"At i={i}" + tgt = y[i] + res = leg.legval(x, [0]*i + [1]) + assert_almost_equal(res, tgt, err_msg=msg) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(leg.legval(x, [1]).shape, dims) + assert_equal(leg.legval(x, [1, 0]).shape, dims) + assert_equal(leg.legval(x, [1, 0, 0]).shape, dims) + + def test_legval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, leg.legval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = leg.legval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.legval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_legval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises(ValueError, leg.legval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = leg.legval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.legval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_leggrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = leg.leggrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.leggrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_leggrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = leg.leggrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = leg.leggrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_legint(self): + # check exceptions + assert_raises(TypeError, leg.legint, [0], .5) + assert_raises(ValueError, leg.legint, [0], -1) + assert_raises(ValueError, leg.legint, [0], 1, [0, 0]) + assert_raises(ValueError, leg.legint, [0], lbnd=[0]) + assert_raises(ValueError, leg.legint, [0], scl=[0]) + assert_raises(TypeError, leg.legint, [0], axis=.5) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = leg.legint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i]) + res = leg.leg2poly(legint) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i], lbnd=-1) + assert_almost_equal(leg.legval(-1, legint), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + legpol = leg.poly2leg(pol) + legint = leg.legint(legpol, m=1, k=[i], scl=2) + res = leg.leg2poly(legint) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1) + res = leg.legint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k]) + res = leg.legint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k], lbnd=-1) + res = leg.legint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = leg.legint(tgt, m=1, k=[k], scl=2) + res = leg.legint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([leg.legint(c) for c in c2d.T]).T + res = leg.legint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legint(c) for c in c2d]) + res = leg.legint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legint(c, k=3) for c in c2d]) + res = leg.legint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + def test_legint_zerointord(self): + assert_equal(leg.legint((1, 2, 3), 0), (1, 2, 3)) + + +class TestDerivative: + + def test_legder(self): + # check exceptions + assert_raises(TypeError, leg.legder, [0], .5) + assert_raises(ValueError, leg.legder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = leg.legder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = leg.legder(leg.legint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = leg.legder(leg.legint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([leg.legder(c) for c in c2d.T]).T + res = leg.legder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([leg.legder(c) for c in c2d]) + res = leg.legder(c2d, axis=1) + assert_almost_equal(res, tgt) + + def test_legder_orderhigherthancoeff(self): + c = (1, 2, 3, 4) + assert_equal(leg.legder(c, 4), [0]) + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_legvander(self): + # check for 1d x + x = np.arange(3) + v = leg.legvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], leg.legval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = leg.legvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], leg.legval(x, coef)) + + def test_legvander2d(self): + # also tests polyval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = leg.legvander2d(x1, x2, [1, 2]) + tgt = leg.legval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = leg.legvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_legvander3d(self): + # also tests polyval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = leg.legvander3d(x1, x2, x3, [1, 2, 3]) + tgt = leg.legval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = leg.legvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + def test_legvander_negdeg(self): + assert_raises(ValueError, leg.legvander, (1, 2, 3), -1) + + +class TestFitting: + + def test_legfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, leg.legfit, [1], [1], -1) + assert_raises(TypeError, leg.legfit, [[1]], [1], 0) + assert_raises(TypeError, leg.legfit, [], [1], 0) + assert_raises(TypeError, leg.legfit, [1], [[[1]]], 0) + assert_raises(TypeError, leg.legfit, [1, 2], [1], 0) + assert_raises(TypeError, leg.legfit, [1], [1, 2], 0) + assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, leg.legfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, leg.legfit, [1], [1], [-1,]) + assert_raises(ValueError, leg.legfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, leg.legfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = leg.legfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(leg.legval(x, coef3), y) + coef3 = leg.legfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(leg.legval(x, coef3), y) + # + coef4 = leg.legfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + coef4 = leg.legfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + # check things still work if deg is not in strict increasing + coef4 = leg.legfit(x, y, [2, 3, 4, 1, 0]) + assert_equal(len(coef4), 5) + assert_almost_equal(leg.legval(x, coef4), y) + # + coef2d = leg.legfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = leg.legfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + y[0::2] = 0 + wcoef3 = leg.legfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = leg.legfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = leg.legfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = leg.legfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(leg.legfit(x, x, 1), [0, 1]) + assert_almost_equal(leg.legfit(x, x, [0, 1]), [0, 1]) + # test fitting only even Legendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = leg.legfit(x, y, 4) + assert_almost_equal(leg.legval(x, coef1), y) + coef2 = leg.legfit(x, y, [0, 2, 4]) + assert_almost_equal(leg.legval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, leg.legcompanion, []) + assert_raises(ValueError, leg.legcompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(leg.legcompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(leg.legcompanion([1, 2])[0, 0] == -.5) + + +class TestGauss: + + def test_100(self): + x, w = leg.leggauss(100) + + # test orthogonality. Note that the results need to be normalized, + # otherwise the huge values that can arise from fast growing + # functions like Laguerre can be very confusing. + v = leg.legvander(x, 99) + vv = np.dot(v.T * w, v) + vd = 1/np.sqrt(vv.diagonal()) + vv = vd[:, None] * vv * vd + assert_almost_equal(vv, np.eye(100)) + + # check that the integral of 1 is correct + tgt = 2.0 + assert_almost_equal(w.sum(), tgt) + + +class TestMisc: + + def test_legfromroots(self): + res = leg.legfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + pol = leg.legfromroots(roots) + res = leg.legval(roots, pol) + tgt = 0 + assert_(len(pol) == i + 1) + assert_almost_equal(leg.leg2poly(pol)[-1], 1) + assert_almost_equal(res, tgt) + + def test_legroots(self): + assert_almost_equal(leg.legroots([1]), []) + assert_almost_equal(leg.legroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = leg.legroots(leg.legfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_legtrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, leg.legtrim, coef, -1) + + # Test results + assert_equal(leg.legtrim(coef), coef[:-1]) + assert_equal(leg.legtrim(coef, 1), coef[:-3]) + assert_equal(leg.legtrim(coef, 2), [0]) + + def test_legline(self): + assert_equal(leg.legline(3, 4), [3, 4]) + + def test_legline_zeroscl(self): + assert_equal(leg.legline(3, 0), [3]) + + def test_leg2poly(self): + for i in range(10): + assert_almost_equal(leg.leg2poly([0]*i + [1]), Llist[i]) + + def test_poly2leg(self): + for i in range(10): + assert_almost_equal(leg.poly2leg(Llist[i]), [0]*i + [1]) + + def test_weight(self): + x = np.linspace(-1, 1, 11) + tgt = 1. + res = leg.legweight(x) + assert_almost_equal(res, tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_polynomial.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_polynomial.py new file mode 100644 index 0000000000000000000000000000000000000000..6b3ef2388f630f0233c79f31a9a1f4039f4e4f4a --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_polynomial.py @@ -0,0 +1,611 @@ +"""Tests for polynomial module. + +""" +from functools import reduce + +import numpy as np +import numpy.polynomial.polynomial as poly +import pickle +from copy import deepcopy +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + assert_warns, assert_array_equal, assert_raises_regex) + + +def trim(x): + return poly.polytrim(x, tol=1e-6) + +T0 = [1] +T1 = [0, 1] +T2 = [-1, 0, 2] +T3 = [0, -3, 0, 4] +T4 = [1, 0, -8, 0, 8] +T5 = [0, 5, 0, -20, 0, 16] +T6 = [-1, 0, 18, 0, -48, 0, 32] +T7 = [0, -7, 0, 56, 0, -112, 0, 64] +T8 = [1, 0, -32, 0, 160, 0, -256, 0, 128] +T9 = [0, 9, 0, -120, 0, 432, 0, -576, 0, 256] + +Tlist = [T0, T1, T2, T3, T4, T5, T6, T7, T8, T9] + + +class TestConstants: + + def test_polydomain(self): + assert_equal(poly.polydomain, [-1, 1]) + + def test_polyzero(self): + assert_equal(poly.polyzero, [0]) + + def test_polyone(self): + assert_equal(poly.polyone, [1]) + + def test_polyx(self): + assert_equal(poly.polyx, [0, 1]) + + def test_copy(self): + x = poly.Polynomial([1, 2, 3]) + y = deepcopy(x) + assert_equal(x, y) + + def test_pickle(self): + x = poly.Polynomial([1, 2, 3]) + y = pickle.loads(pickle.dumps(x)) + assert_equal(x, y) + +class TestArithmetic: + + def test_polyadd(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] += 1 + res = poly.polyadd([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polysub(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(max(i, j) + 1) + tgt[i] += 1 + tgt[j] -= 1 + res = poly.polysub([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polymulx(self): + assert_equal(poly.polymulx([0]), [0]) + assert_equal(poly.polymulx([1]), [0, 1]) + for i in range(1, 5): + ser = [0]*i + [1] + tgt = [0]*(i + 1) + [1] + assert_equal(poly.polymulx(ser), tgt) + + def test_polymul(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + tgt = np.zeros(i + j + 1) + tgt[i + j] += 1 + res = poly.polymul([0]*i + [1], [0]*j + [1]) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + def test_polydiv(self): + # check zero division + assert_raises(ZeroDivisionError, poly.polydiv, [1], [0]) + + # check scalar division + quo, rem = poly.polydiv([2], [2]) + assert_equal((quo, rem), (1, 0)) + quo, rem = poly.polydiv([2, 2], [2]) + assert_equal((quo, rem), ((1, 1), 0)) + + # check rest. + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + ci = [0]*i + [1, 2] + cj = [0]*j + [1, 2] + tgt = poly.polyadd(ci, cj) + quo, rem = poly.polydiv(tgt, ci) + res = poly.polyadd(poly.polymul(quo, ci), rem) + assert_equal(res, tgt, err_msg=msg) + + def test_polypow(self): + for i in range(5): + for j in range(5): + msg = f"At i={i}, j={j}" + c = np.arange(i + 1) + tgt = reduce(poly.polymul, [c]*j, np.array([1])) + res = poly.polypow(c, j) + assert_equal(trim(res), trim(tgt), err_msg=msg) + + +class TestEvaluation: + # coefficients of 1 + 2*x + 3*x**2 + c1d = np.array([1., 2., 3.]) + c2d = np.einsum('i,j->ij', c1d, c1d) + c3d = np.einsum('i,j,k->ijk', c1d, c1d, c1d) + + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + y = poly.polyval(x, [1., 2., 3.]) + + def test_polyval(self): + #check empty input + assert_equal(poly.polyval([], [1]).size, 0) + + #check normal input) + x = np.linspace(-1, 1) + y = [x**i for i in range(5)] + for i in range(5): + tgt = y[i] + res = poly.polyval(x, [0]*i + [1]) + assert_almost_equal(res, tgt) + tgt = x*(x**2 - 1) + res = poly.polyval(x, [0, -1, 0, 1]) + assert_almost_equal(res, tgt) + + #check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(poly.polyval(x, [1]).shape, dims) + assert_equal(poly.polyval(x, [1, 0]).shape, dims) + assert_equal(poly.polyval(x, [1, 0, 0]).shape, dims) + + #check masked arrays are processed correctly + mask = [False, True, False] + mx = np.ma.array([1, 2, 3], mask=mask) + res = np.polyval([7, 5, 3], mx) + assert_array_equal(res.mask, mask) + + #check subtypes of ndarray are preserved + class C(np.ndarray): + pass + + cx = np.array([1, 2, 3]).view(C) + assert_equal(type(np.polyval([2, 3, 4], cx)), C) + + def test_polyvalfromroots(self): + # check exception for broadcasting x values over root array with + # too few dimensions + assert_raises(ValueError, poly.polyvalfromroots, + [1], [1], tensor=False) + + # check empty input + assert_equal(poly.polyvalfromroots([], [1]).size, 0) + assert_(poly.polyvalfromroots([], [1]).shape == (0,)) + + # check empty input + multidimensional roots + assert_equal(poly.polyvalfromroots([], [[1] * 5]).size, 0) + assert_(poly.polyvalfromroots([], [[1] * 5]).shape == (5, 0)) + + # check scalar input + assert_equal(poly.polyvalfromroots(1, 1), 0) + assert_(poly.polyvalfromroots(1, np.ones((3, 3))).shape == (3,)) + + # check normal input) + x = np.linspace(-1, 1) + y = [x**i for i in range(5)] + for i in range(1, 5): + tgt = y[i] + res = poly.polyvalfromroots(x, [0]*i) + assert_almost_equal(res, tgt) + tgt = x*(x - 1)*(x + 1) + res = poly.polyvalfromroots(x, [-1, 0, 1]) + assert_almost_equal(res, tgt) + + # check that shape is preserved + for i in range(3): + dims = [2]*i + x = np.zeros(dims) + assert_equal(poly.polyvalfromroots(x, [1]).shape, dims) + assert_equal(poly.polyvalfromroots(x, [1, 0]).shape, dims) + assert_equal(poly.polyvalfromroots(x, [1, 0, 0]).shape, dims) + + # check compatibility with factorization + ptest = [15, 2, -16, -2, 1] + r = poly.polyroots(ptest) + x = np.linspace(-1, 1) + assert_almost_equal(poly.polyval(x, ptest), + poly.polyvalfromroots(x, r)) + + # check multidimensional arrays of roots and values + # check tensor=False + rshape = (3, 5) + x = np.arange(-3, 2) + r = np.random.randint(-5, 5, size=rshape) + res = poly.polyvalfromroots(x, r, tensor=False) + tgt = np.empty(r.shape[1:]) + for ii in range(tgt.size): + tgt[ii] = poly.polyvalfromroots(x[ii], r[:, ii]) + assert_equal(res, tgt) + + # check tensor=True + x = np.vstack([x, 2*x]) + res = poly.polyvalfromroots(x, r, tensor=True) + tgt = np.empty(r.shape[1:] + x.shape) + for ii in range(r.shape[1]): + for jj in range(x.shape[0]): + tgt[ii, jj, :] = poly.polyvalfromroots(x[jj], r[:, ii]) + assert_equal(res, tgt) + + def test_polyval2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises_regex(ValueError, 'incompatible', + poly.polyval2d, x1, x2[:2], self.c2d) + + #test values + tgt = y1*y2 + res = poly.polyval2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polyval2d(z, z, self.c2d) + assert_(res.shape == (2, 3)) + + def test_polyval3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test exceptions + assert_raises_regex(ValueError, 'incompatible', + poly.polyval3d, x1, x2, x3[:2], self.c3d) + + #test values + tgt = y1*y2*y3 + res = poly.polyval3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polyval3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)) + + def test_polygrid2d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j->ij', y1, y2) + res = poly.polygrid2d(x1, x2, self.c2d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polygrid2d(z, z, self.c2d) + assert_(res.shape == (2, 3)*2) + + def test_polygrid3d(self): + x1, x2, x3 = self.x + y1, y2, y3 = self.y + + #test values + tgt = np.einsum('i,j,k->ijk', y1, y2, y3) + res = poly.polygrid3d(x1, x2, x3, self.c3d) + assert_almost_equal(res, tgt) + + #test shape + z = np.ones((2, 3)) + res = poly.polygrid3d(z, z, z, self.c3d) + assert_(res.shape == (2, 3)*3) + + +class TestIntegral: + + def test_polyint(self): + # check exceptions + assert_raises(TypeError, poly.polyint, [0], .5) + assert_raises(ValueError, poly.polyint, [0], -1) + assert_raises(ValueError, poly.polyint, [0], 1, [0, 0]) + assert_raises(ValueError, poly.polyint, [0], lbnd=[0]) + assert_raises(ValueError, poly.polyint, [0], scl=[0]) + assert_raises(TypeError, poly.polyint, [0], axis=.5) + with assert_warns(DeprecationWarning): + poly.polyint([1, 1], 1.) + + # test integration of zero polynomial + for i in range(2, 5): + k = [0]*(i - 2) + [1] + res = poly.polyint([0], m=i, k=k) + assert_almost_equal(res, [0, 1]) + + # check single integration with integration constant + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [1/scl] + res = poly.polyint(pol, m=1, k=[i]) + assert_almost_equal(trim(res), trim(tgt)) + + # check single integration with integration constant and lbnd + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + res = poly.polyint(pol, m=1, k=[i], lbnd=-1) + assert_almost_equal(poly.polyval(-1, res), i) + + # check single integration with integration constant and scaling + for i in range(5): + scl = i + 1 + pol = [0]*i + [1] + tgt = [i] + [0]*i + [2/scl] + res = poly.polyint(pol, m=1, k=[i], scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with default k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1) + res = poly.polyint(pol, m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with defined k + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k]) + res = poly.polyint(pol, m=j, k=list(range(j))) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with lbnd + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k], lbnd=-1) + res = poly.polyint(pol, m=j, k=list(range(j)), lbnd=-1) + assert_almost_equal(trim(res), trim(tgt)) + + # check multiple integrations with scaling + for i in range(5): + for j in range(2, 5): + pol = [0]*i + [1] + tgt = pol[:] + for k in range(j): + tgt = poly.polyint(tgt, m=1, k=[k], scl=2) + res = poly.polyint(pol, m=j, k=list(range(j)), scl=2) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyint_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([poly.polyint(c) for c in c2d.T]).T + res = poly.polyint(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyint(c) for c in c2d]) + res = poly.polyint(c2d, axis=1) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyint(c, k=3) for c in c2d]) + res = poly.polyint(c2d, k=3, axis=1) + assert_almost_equal(res, tgt) + + +class TestDerivative: + + def test_polyder(self): + # check exceptions + assert_raises(TypeError, poly.polyder, [0], .5) + assert_raises(ValueError, poly.polyder, [0], -1) + + # check that zeroth derivative does nothing + for i in range(5): + tgt = [0]*i + [1] + res = poly.polyder(tgt, m=0) + assert_equal(trim(res), trim(tgt)) + + # check that derivation is the inverse of integration + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = poly.polyder(poly.polyint(tgt, m=j), m=j) + assert_almost_equal(trim(res), trim(tgt)) + + # check derivation with scaling + for i in range(5): + for j in range(2, 5): + tgt = [0]*i + [1] + res = poly.polyder(poly.polyint(tgt, m=j, scl=2), m=j, scl=.5) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyder_axis(self): + # check that axis keyword works + c2d = np.random.random((3, 4)) + + tgt = np.vstack([poly.polyder(c) for c in c2d.T]).T + res = poly.polyder(c2d, axis=0) + assert_almost_equal(res, tgt) + + tgt = np.vstack([poly.polyder(c) for c in c2d]) + res = poly.polyder(c2d, axis=1) + assert_almost_equal(res, tgt) + + +class TestVander: + # some random values in [-1, 1) + x = np.random.random((3, 5))*2 - 1 + + def test_polyvander(self): + # check for 1d x + x = np.arange(3) + v = poly.polyvander(x, 3) + assert_(v.shape == (3, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], poly.polyval(x, coef)) + + # check for 2d x + x = np.array([[1, 2], [3, 4], [5, 6]]) + v = poly.polyvander(x, 3) + assert_(v.shape == (3, 2, 4)) + for i in range(4): + coef = [0]*i + [1] + assert_almost_equal(v[..., i], poly.polyval(x, coef)) + + def test_polyvander2d(self): + # also tests polyval2d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3)) + van = poly.polyvander2d(x1, x2, [1, 2]) + tgt = poly.polyval2d(x1, x2, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = poly.polyvander2d([x1], [x2], [1, 2]) + assert_(van.shape == (1, 5, 6)) + + def test_polyvander3d(self): + # also tests polyval3d for non-square coefficient array + x1, x2, x3 = self.x + c = np.random.random((2, 3, 4)) + van = poly.polyvander3d(x1, x2, x3, [1, 2, 3]) + tgt = poly.polyval3d(x1, x2, x3, c) + res = np.dot(van, c.flat) + assert_almost_equal(res, tgt) + + # check shape + van = poly.polyvander3d([x1], [x2], [x3], [1, 2, 3]) + assert_(van.shape == (1, 5, 24)) + + def test_polyvandernegdeg(self): + x = np.arange(3) + assert_raises(ValueError, poly.polyvander, x, -1) + + +class TestCompanion: + + def test_raises(self): + assert_raises(ValueError, poly.polycompanion, []) + assert_raises(ValueError, poly.polycompanion, [1]) + + def test_dimensions(self): + for i in range(1, 5): + coef = [0]*i + [1] + assert_(poly.polycompanion(coef).shape == (i, i)) + + def test_linear_root(self): + assert_(poly.polycompanion([1, 2])[0, 0] == -.5) + + +class TestMisc: + + def test_polyfromroots(self): + res = poly.polyfromroots([]) + assert_almost_equal(trim(res), [1]) + for i in range(1, 5): + roots = np.cos(np.linspace(-np.pi, 0, 2*i + 1)[1::2]) + tgt = Tlist[i] + res = poly.polyfromroots(roots)*2**(i-1) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyroots(self): + assert_almost_equal(poly.polyroots([1]), []) + assert_almost_equal(poly.polyroots([1, 2]), [-.5]) + for i in range(2, 5): + tgt = np.linspace(-1, 1, i) + res = poly.polyroots(poly.polyfromroots(tgt)) + assert_almost_equal(trim(res), trim(tgt)) + + def test_polyfit(self): + def f(x): + return x*(x - 1)*(x - 2) + + def f2(x): + return x**4 + x**2 + 1 + + # Test exceptions + assert_raises(ValueError, poly.polyfit, [1], [1], -1) + assert_raises(TypeError, poly.polyfit, [[1]], [1], 0) + assert_raises(TypeError, poly.polyfit, [], [1], 0) + assert_raises(TypeError, poly.polyfit, [1], [[[1]]], 0) + assert_raises(TypeError, poly.polyfit, [1, 2], [1], 0) + assert_raises(TypeError, poly.polyfit, [1], [1, 2], 0) + assert_raises(TypeError, poly.polyfit, [1], [1], 0, w=[[1]]) + assert_raises(TypeError, poly.polyfit, [1], [1], 0, w=[1, 1]) + assert_raises(ValueError, poly.polyfit, [1], [1], [-1,]) + assert_raises(ValueError, poly.polyfit, [1], [1], [2, -1, 6]) + assert_raises(TypeError, poly.polyfit, [1], [1], []) + + # Test fit + x = np.linspace(0, 2) + y = f(x) + # + coef3 = poly.polyfit(x, y, 3) + assert_equal(len(coef3), 4) + assert_almost_equal(poly.polyval(x, coef3), y) + coef3 = poly.polyfit(x, y, [0, 1, 2, 3]) + assert_equal(len(coef3), 4) + assert_almost_equal(poly.polyval(x, coef3), y) + # + coef4 = poly.polyfit(x, y, 4) + assert_equal(len(coef4), 5) + assert_almost_equal(poly.polyval(x, coef4), y) + coef4 = poly.polyfit(x, y, [0, 1, 2, 3, 4]) + assert_equal(len(coef4), 5) + assert_almost_equal(poly.polyval(x, coef4), y) + # + coef2d = poly.polyfit(x, np.array([y, y]).T, 3) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + coef2d = poly.polyfit(x, np.array([y, y]).T, [0, 1, 2, 3]) + assert_almost_equal(coef2d, np.array([coef3, coef3]).T) + # test weighting + w = np.zeros_like(x) + yw = y.copy() + w[1::2] = 1 + yw[0::2] = 0 + wcoef3 = poly.polyfit(x, yw, 3, w=w) + assert_almost_equal(wcoef3, coef3) + wcoef3 = poly.polyfit(x, yw, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef3, coef3) + # + wcoef2d = poly.polyfit(x, np.array([yw, yw]).T, 3, w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + wcoef2d = poly.polyfit(x, np.array([yw, yw]).T, [0, 1, 2, 3], w=w) + assert_almost_equal(wcoef2d, np.array([coef3, coef3]).T) + # test scaling with complex values x points whose square + # is zero when summed. + x = [1, 1j, -1, -1j] + assert_almost_equal(poly.polyfit(x, x, 1), [0, 1]) + assert_almost_equal(poly.polyfit(x, x, [0, 1]), [0, 1]) + # test fitting only even Polyendre polynomials + x = np.linspace(-1, 1) + y = f2(x) + coef1 = poly.polyfit(x, y, 4) + assert_almost_equal(poly.polyval(x, coef1), y) + coef2 = poly.polyfit(x, y, [0, 2, 4]) + assert_almost_equal(poly.polyval(x, coef2), y) + assert_almost_equal(coef1, coef2) + + def test_polytrim(self): + coef = [2, -1, 1, 0] + + # Test exceptions + assert_raises(ValueError, poly.polytrim, coef, -1) + + # Test results + assert_equal(poly.polytrim(coef), coef[:-1]) + assert_equal(poly.polytrim(coef, 1), coef[:-3]) + assert_equal(poly.polytrim(coef, 2), [0]) + + def test_polyline(self): + assert_equal(poly.polyline(3, 4), [3, 4]) + + def test_polyline_zero(self): + assert_equal(poly.polyline(3, 0), [3]) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_polyutils.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_polyutils.py new file mode 100644 index 0000000000000000000000000000000000000000..cc630790da1ce8fd1ca413cd530ae5636cce5aa8 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_polyutils.py @@ -0,0 +1,121 @@ +"""Tests for polyutils module. + +""" +import numpy as np +import numpy.polynomial.polyutils as pu +from numpy.testing import ( + assert_almost_equal, assert_raises, assert_equal, assert_, + ) + + +class TestMisc: + + def test_trimseq(self): + for i in range(5): + tgt = [1] + res = pu.trimseq([1] + [0]*5) + assert_equal(res, tgt) + + def test_as_series(self): + # check exceptions + assert_raises(ValueError, pu.as_series, [[]]) + assert_raises(ValueError, pu.as_series, [[[1, 2]]]) + assert_raises(ValueError, pu.as_series, [[1], ['a']]) + # check common types + types = ['i', 'd', 'O'] + for i in range(len(types)): + for j in range(i): + ci = np.ones(1, types[i]) + cj = np.ones(1, types[j]) + [resi, resj] = pu.as_series([ci, cj]) + assert_(resi.dtype.char == resj.dtype.char) + assert_(resj.dtype.char == types[i]) + + def test_trimcoef(self): + coef = [2, -1, 1, 0] + # Test exceptions + assert_raises(ValueError, pu.trimcoef, coef, -1) + # Test results + assert_equal(pu.trimcoef(coef), coef[:-1]) + assert_equal(pu.trimcoef(coef, 1), coef[:-3]) + assert_equal(pu.trimcoef(coef, 2), [0]) + + def test_vander_nd_exception(self): + # n_dims != len(points) + assert_raises(ValueError, pu._vander_nd, (), (1, 2, 3), [90]) + # n_dims != len(degrees) + assert_raises(ValueError, pu._vander_nd, (), (), [90.65]) + # n_dims == 0 + assert_raises(ValueError, pu._vander_nd, (), (), []) + + def test_div_zerodiv(self): + # c2[-1] == 0 + assert_raises(ZeroDivisionError, pu._div, pu._div, (1, 2, 3), [0]) + + def test_pow_too_large(self): + # power > maxpower + assert_raises(ValueError, pu._pow, (), [1, 2, 3], 5, 4) + +class TestDomain: + + def test_getdomain(self): + # test for real values + x = [1, 10, 3, -1] + tgt = [-1, 10] + res = pu.getdomain(x) + assert_almost_equal(res, tgt) + + # test for complex values + x = [1 + 1j, 1 - 1j, 0, 2] + tgt = [-1j, 2 + 1j] + res = pu.getdomain(x) + assert_almost_equal(res, tgt) + + def test_mapdomain(self): + # test for real values + dom1 = [0, 4] + dom2 = [1, 3] + tgt = dom2 + res = pu.mapdomain(dom1, dom1, dom2) + assert_almost_equal(res, tgt) + + # test for complex values + dom1 = [0 - 1j, 2 + 1j] + dom2 = [-2, 2] + tgt = dom2 + x = dom1 + res = pu.mapdomain(x, dom1, dom2) + assert_almost_equal(res, tgt) + + # test for multidimensional arrays + dom1 = [0, 4] + dom2 = [1, 3] + tgt = np.array([dom2, dom2]) + x = np.array([dom1, dom1]) + res = pu.mapdomain(x, dom1, dom2) + assert_almost_equal(res, tgt) + + # test that subtypes are preserved. + class MyNDArray(np.ndarray): + pass + + dom1 = [0, 4] + dom2 = [1, 3] + x = np.array([dom1, dom1]).view(MyNDArray) + res = pu.mapdomain(x, dom1, dom2) + assert_(isinstance(res, MyNDArray)) + + def test_mapparms(self): + # test for real values + dom1 = [0, 4] + dom2 = [1, 3] + tgt = [1, .5] + res = pu. mapparms(dom1, dom2) + assert_almost_equal(res, tgt) + + # test for complex values + dom1 = [0 - 1j, 2 + 1j] + dom2 = [-2, 2] + tgt = [-1 + 1j, 1 - 1j] + res = pu.mapparms(dom1, dom2) + assert_almost_equal(res, tgt) diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_printing.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_printing.py new file mode 100644 index 0000000000000000000000000000000000000000..6f2a5092d7225c797b60fd8f2602f2f9276cdd74 --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_printing.py @@ -0,0 +1,530 @@ +from math import nan, inf +import pytest +from numpy.core import array, arange, printoptions +import numpy.polynomial as poly +from numpy.testing import assert_equal, assert_ + +# For testing polynomial printing with object arrays +from fractions import Fraction +from decimal import Decimal + + +class TestStrUnicodeSuperSubscripts: + + @pytest.fixture(scope='class', autouse=True) + def use_unicode(self): + poly.set_default_printstyle('unicode') + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·x + 3.0·x²"), + ([-1, 0, 3, -1], "-1.0 + 0.0·x + 3.0·x² - 1.0·x³"), + (arange(12), ("0.0 + 1.0·x + 2.0·x² + 3.0·x³ + 4.0·x⁴ + 5.0·x⁵ + " + "6.0·x⁶ + 7.0·x⁷ +\n8.0·x⁸ + 9.0·x⁹ + 10.0·x¹⁰ + " + "11.0·x¹¹")), + )) + def test_polynomial_str(self, inp, tgt): + res = str(poly.Polynomial(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·T₁(x) + 3.0·T₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·T₁(x) + 3.0·T₂(x) - 1.0·T₃(x)"), + (arange(12), ("0.0 + 1.0·T₁(x) + 2.0·T₂(x) + 3.0·T₃(x) + 4.0·T₄(x) + " + "5.0·T₅(x) +\n6.0·T₆(x) + 7.0·T₇(x) + 8.0·T₈(x) + " + "9.0·T₉(x) + 10.0·T₁₀(x) + 11.0·T₁₁(x)")), + )) + def test_chebyshev_str(self, inp, tgt): + res = str(poly.Chebyshev(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·P₁(x) + 3.0·P₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·P₁(x) + 3.0·P₂(x) - 1.0·P₃(x)"), + (arange(12), ("0.0 + 1.0·P₁(x) + 2.0·P₂(x) + 3.0·P₃(x) + 4.0·P₄(x) + " + "5.0·P₅(x) +\n6.0·P₆(x) + 7.0·P₇(x) + 8.0·P₈(x) + " + "9.0·P₉(x) + 10.0·P₁₀(x) + 11.0·P₁₁(x)")), + )) + def test_legendre_str(self, inp, tgt): + res = str(poly.Legendre(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·H₁(x) + 3.0·H₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·H₁(x) + 3.0·H₂(x) - 1.0·H₃(x)"), + (arange(12), ("0.0 + 1.0·H₁(x) + 2.0·H₂(x) + 3.0·H₃(x) + 4.0·H₄(x) + " + "5.0·H₅(x) +\n6.0·H₆(x) + 7.0·H₇(x) + 8.0·H₈(x) + " + "9.0·H₉(x) + 10.0·H₁₀(x) + 11.0·H₁₁(x)")), + )) + def test_hermite_str(self, inp, tgt): + res = str(poly.Hermite(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·He₁(x) + 3.0·He₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·He₁(x) + 3.0·He₂(x) - 1.0·He₃(x)"), + (arange(12), ("0.0 + 1.0·He₁(x) + 2.0·He₂(x) + 3.0·He₃(x) + " + "4.0·He₄(x) + 5.0·He₅(x) +\n6.0·He₆(x) + 7.0·He₇(x) + " + "8.0·He₈(x) + 9.0·He₉(x) + 10.0·He₁₀(x) +\n" + "11.0·He₁₁(x)")), + )) + def test_hermiteE_str(self, inp, tgt): + res = str(poly.HermiteE(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0·L₁(x) + 3.0·L₂(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0·L₁(x) + 3.0·L₂(x) - 1.0·L₃(x)"), + (arange(12), ("0.0 + 1.0·L₁(x) + 2.0·L₂(x) + 3.0·L₃(x) + 4.0·L₄(x) + " + "5.0·L₅(x) +\n6.0·L₆(x) + 7.0·L₇(x) + 8.0·L₈(x) + " + "9.0·L₉(x) + 10.0·L₁₀(x) + 11.0·L₁₁(x)")), + )) + def test_laguerre_str(self, inp, tgt): + res = str(poly.Laguerre(inp)) + assert_equal(res, tgt) + + +class TestStrAscii: + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 x + 3.0 x**2"), + ([-1, 0, 3, -1], "-1.0 + 0.0 x + 3.0 x**2 - 1.0 x**3"), + (arange(12), ("0.0 + 1.0 x + 2.0 x**2 + 3.0 x**3 + 4.0 x**4 + " + "5.0 x**5 + 6.0 x**6 +\n7.0 x**7 + 8.0 x**8 + " + "9.0 x**9 + 10.0 x**10 + 11.0 x**11")), + )) + def test_polynomial_str(self, inp, tgt): + res = str(poly.Polynomial(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 T_1(x) + 3.0 T_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 T_1(x) + 3.0 T_2(x) - 1.0 T_3(x)"), + (arange(12), ("0.0 + 1.0 T_1(x) + 2.0 T_2(x) + 3.0 T_3(x) + " + "4.0 T_4(x) + 5.0 T_5(x) +\n6.0 T_6(x) + 7.0 T_7(x) + " + "8.0 T_8(x) + 9.0 T_9(x) + 10.0 T_10(x) +\n" + "11.0 T_11(x)")), + )) + def test_chebyshev_str(self, inp, tgt): + res = str(poly.Chebyshev(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 P_1(x) + 3.0 P_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 P_1(x) + 3.0 P_2(x) - 1.0 P_3(x)"), + (arange(12), ("0.0 + 1.0 P_1(x) + 2.0 P_2(x) + 3.0 P_3(x) + " + "4.0 P_4(x) + 5.0 P_5(x) +\n6.0 P_6(x) + 7.0 P_7(x) + " + "8.0 P_8(x) + 9.0 P_9(x) + 10.0 P_10(x) +\n" + "11.0 P_11(x)")), + )) + def test_legendre_str(self, inp, tgt): + res = str(poly.Legendre(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 H_1(x) + 3.0 H_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 H_1(x) + 3.0 H_2(x) - 1.0 H_3(x)"), + (arange(12), ("0.0 + 1.0 H_1(x) + 2.0 H_2(x) + 3.0 H_3(x) + " + "4.0 H_4(x) + 5.0 H_5(x) +\n6.0 H_6(x) + 7.0 H_7(x) + " + "8.0 H_8(x) + 9.0 H_9(x) + 10.0 H_10(x) +\n" + "11.0 H_11(x)")), + )) + def test_hermite_str(self, inp, tgt): + res = str(poly.Hermite(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 He_1(x) + 3.0 He_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 He_1(x) + 3.0 He_2(x) - 1.0 He_3(x)"), + (arange(12), ("0.0 + 1.0 He_1(x) + 2.0 He_2(x) + 3.0 He_3(x) + " + "4.0 He_4(x) +\n5.0 He_5(x) + 6.0 He_6(x) + " + "7.0 He_7(x) + 8.0 He_8(x) + 9.0 He_9(x) +\n" + "10.0 He_10(x) + 11.0 He_11(x)")), + )) + def test_hermiteE_str(self, inp, tgt): + res = str(poly.HermiteE(inp)) + assert_equal(res, tgt) + + @pytest.mark.parametrize(('inp', 'tgt'), ( + ([1, 2, 3], "1.0 + 2.0 L_1(x) + 3.0 L_2(x)"), + ([-1, 0, 3, -1], "-1.0 + 0.0 L_1(x) + 3.0 L_2(x) - 1.0 L_3(x)"), + (arange(12), ("0.0 + 1.0 L_1(x) + 2.0 L_2(x) + 3.0 L_3(x) + " + "4.0 L_4(x) + 5.0 L_5(x) +\n6.0 L_6(x) + 7.0 L_7(x) + " + "8.0 L_8(x) + 9.0 L_9(x) + 10.0 L_10(x) +\n" + "11.0 L_11(x)")), + )) + def test_laguerre_str(self, inp, tgt): + res = str(poly.Laguerre(inp)) + assert_equal(res, tgt) + + +class TestLinebreaking: + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + def test_single_line_one_less(self): + # With 'ascii' style, len(str(p)) is default linewidth - 1 (i.e. 74) + p = poly.Polynomial([12345678, 12345678, 12345678, 12345678, 123]) + assert_equal(len(str(p)), 74) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 + 123.0 x**4' + )) + + def test_num_chars_is_linewidth(self): + # len(str(p)) == default linewidth == 75 + p = poly.Polynomial([12345678, 12345678, 12345678, 12345678, 1234]) + assert_equal(len(str(p)), 75) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 +\n1234.0 x**4' + )) + + def test_first_linebreak_multiline_one_less_than_linewidth(self): + # Multiline str where len(first_line) + len(next_term) == lw - 1 == 74 + p = poly.Polynomial( + [12345678, 12345678, 12345678, 12345678, 1, 12345678] + ) + assert_equal(len(str(p).split('\n')[0]), 74) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.0 x**3 + 1.0 x**4 +\n12345678.0 x**5' + )) + + def test_first_linebreak_multiline_on_linewidth(self): + # First line is one character longer than previous test + p = poly.Polynomial( + [12345678, 12345678, 12345678, 12345678.12, 1, 12345678] + ) + assert_equal(str(p), ( + '12345678.0 + 12345678.0 x + 12345678.0 x**2 + ' + '12345678.12 x**3 +\n1.0 x**4 + 12345678.0 x**5' + )) + + @pytest.mark.parametrize(('lw', 'tgt'), ( + (75, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 + 40000.0 x**4 + ' + '500000.0 x**5 +\n600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 + ' + '900.0 x**9')), + (45, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 +\n40000.0 x**4 + ' + '500000.0 x**5 +\n600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 +\n' + '900.0 x**9')), + (132, ('0.0 + 10.0 x + 200.0 x**2 + 3000.0 x**3 + 40000.0 x**4 + ' + '500000.0 x**5 + 600000.0 x**6 + 70000.0 x**7 + 8000.0 x**8 + ' + '900.0 x**9')), + )) + def test_linewidth_printoption(self, lw, tgt): + p = poly.Polynomial( + [0, 10, 200, 3000, 40000, 500000, 600000, 70000, 8000, 900] + ) + with printoptions(linewidth=lw): + assert_equal(str(p), tgt) + for line in str(p).split('\n'): + assert_(len(line) < lw) + + +def test_set_default_printoptions(): + p = poly.Polynomial([1, 2, 3]) + c = poly.Chebyshev([1, 2, 3]) + poly.set_default_printstyle('ascii') + assert_equal(str(p), "1.0 + 2.0 x + 3.0 x**2") + assert_equal(str(c), "1.0 + 2.0 T_1(x) + 3.0 T_2(x)") + poly.set_default_printstyle('unicode') + assert_equal(str(p), "1.0 + 2.0·x + 3.0·x²") + assert_equal(str(c), "1.0 + 2.0·T₁(x) + 3.0·T₂(x)") + with pytest.raises(ValueError): + poly.set_default_printstyle('invalid_input') + + +def test_complex_coefficients(): + """Test both numpy and built-in complex.""" + coefs = [0+1j, 1+1j, -2+2j, 3+0j] + # numpy complex + p1 = poly.Polynomial(coefs) + # Python complex + p2 = poly.Polynomial(array(coefs, dtype=object)) + poly.set_default_printstyle('unicode') + assert_equal(str(p1), "1j + (1+1j)·x - (2-2j)·x² + (3+0j)·x³") + assert_equal(str(p2), "1j + (1+1j)·x + (-2+2j)·x² + (3+0j)·x³") + poly.set_default_printstyle('ascii') + assert_equal(str(p1), "1j + (1+1j) x - (2-2j) x**2 + (3+0j) x**3") + assert_equal(str(p2), "1j + (1+1j) x + (-2+2j) x**2 + (3+0j) x**3") + + +@pytest.mark.parametrize(('coefs', 'tgt'), ( + (array([Fraction(1, 2), Fraction(3, 4)], dtype=object), ( + "1/2 + 3/4·x" + )), + (array([1, 2, Fraction(5, 7)], dtype=object), ( + "1 + 2·x + 5/7·x²" + )), + (array([Decimal('1.00'), Decimal('2.2'), 3], dtype=object), ( + "1.00 + 2.2·x + 3·x²" + )), +)) +def test_numeric_object_coefficients(coefs, tgt): + p = poly.Polynomial(coefs) + poly.set_default_printstyle('unicode') + assert_equal(str(p), tgt) + + +@pytest.mark.parametrize(('coefs', 'tgt'), ( + (array([1, 2, 'f'], dtype=object), '1 + 2·x + f·x²'), + (array([1, 2, [3, 4]], dtype=object), '1 + 2·x + [3, 4]·x²'), +)) +def test_nonnumeric_object_coefficients(coefs, tgt): + """ + Test coef fallback for object arrays of non-numeric coefficients. + """ + p = poly.Polynomial(coefs) + poly.set_default_printstyle('unicode') + assert_equal(str(p), tgt) + + +class TestFormat: + def test_format_unicode(self): + poly.set_default_printstyle('ascii') + p = poly.Polynomial([1, 2, 0, -1]) + assert_equal(format(p, 'unicode'), "1.0 + 2.0·x + 0.0·x² - 1.0·x³") + + def test_format_ascii(self): + poly.set_default_printstyle('unicode') + p = poly.Polynomial([1, 2, 0, -1]) + assert_equal( + format(p, 'ascii'), "1.0 + 2.0 x + 0.0 x**2 - 1.0 x**3" + ) + + def test_empty_formatstr(self): + poly.set_default_printstyle('ascii') + p = poly.Polynomial([1, 2, 3]) + assert_equal(format(p), "1.0 + 2.0 x + 3.0 x**2") + assert_equal(f"{p}", "1.0 + 2.0 x + 3.0 x**2") + + def test_bad_formatstr(self): + p = poly.Polynomial([1, 2, 0, -1]) + with pytest.raises(ValueError): + format(p, '.2f') + + +@pytest.mark.parametrize(('poly', 'tgt'), ( + (poly.Polynomial, '1.0 + 2.0·z + 3.0·z²'), + (poly.Chebyshev, '1.0 + 2.0·T₁(z) + 3.0·T₂(z)'), + (poly.Hermite, '1.0 + 2.0·H₁(z) + 3.0·H₂(z)'), + (poly.HermiteE, '1.0 + 2.0·He₁(z) + 3.0·He₂(z)'), + (poly.Laguerre, '1.0 + 2.0·L₁(z) + 3.0·L₂(z)'), + (poly.Legendre, '1.0 + 2.0·P₁(z) + 3.0·P₂(z)'), +)) +def test_symbol(poly, tgt): + p = poly([1, 2, 3], symbol='z') + assert_equal(f"{p:unicode}", tgt) + + +class TestRepr: + def test_polynomial_str(self): + res = repr(poly.Polynomial([0, 1])) + tgt = ( + "Polynomial([0., 1.], domain=[-1, 1], window=[-1, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_chebyshev_str(self): + res = repr(poly.Chebyshev([0, 1])) + tgt = ( + "Chebyshev([0., 1.], domain=[-1, 1], window=[-1, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_legendre_repr(self): + res = repr(poly.Legendre([0, 1])) + tgt = ( + "Legendre([0., 1.], domain=[-1, 1], window=[-1, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_hermite_repr(self): + res = repr(poly.Hermite([0, 1])) + tgt = ( + "Hermite([0., 1.], domain=[-1, 1], window=[-1, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_hermiteE_repr(self): + res = repr(poly.HermiteE([0, 1])) + tgt = ( + "HermiteE([0., 1.], domain=[-1, 1], window=[-1, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + def test_laguerre_repr(self): + res = repr(poly.Laguerre([0, 1])) + tgt = ( + "Laguerre([0., 1.], domain=[0, 1], window=[0, 1], " + "symbol='x')" + ) + assert_equal(res, tgt) + + +class TestLatexRepr: + """Test the latex repr used by Jupyter""" + + def as_latex(self, obj): + # right now we ignore the formatting of scalars in our tests, since + # it makes them too verbose. Ideally, the formatting of scalars will + # be fixed such that tests below continue to pass + obj._repr_latex_scalar = lambda x, parens=False: str(x) + try: + return obj._repr_latex_() + finally: + del obj._repr_latex_scalar + + def test_simple_polynomial(self): + # default input + p = poly.Polynomial([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,x + 3.0\,x^{2}$') + + # translated input + p = poly.Polynomial([1, 2, 3], domain=[-2, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(1.0 + x\right) + 3.0\,\left(1.0 + x\right)^{2}$') + + # scaled input + p = poly.Polynomial([1, 2, 3], domain=[-0.5, 0.5]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(2.0x\right) + 3.0\,\left(2.0x\right)^{2}$') + + # affine input + p = poly.Polynomial([1, 2, 3], domain=[-1, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0 + 2.0\,\left(1.0 + 2.0x\right) + 3.0\,\left(1.0 + 2.0x\right)^{2}$') + + def test_basis_func(self): + p = poly.Chebyshev([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{T}_{0}(x) + 2.0\,{T}_{1}(x) + 3.0\,{T}_{2}(x)$') + # affine input - check no surplus parens are added + p = poly.Chebyshev([1, 2, 3], domain=[-1, 0]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{T}_{0}(1.0 + 2.0x) + 2.0\,{T}_{1}(1.0 + 2.0x) + 3.0\,{T}_{2}(1.0 + 2.0x)$') + + def test_multichar_basis_func(self): + p = poly.HermiteE([1, 2, 3]) + assert_equal(self.as_latex(p), + r'$x \mapsto 1.0\,{He}_{0}(x) + 2.0\,{He}_{1}(x) + 3.0\,{He}_{2}(x)$') + + def test_symbol_basic(self): + # default input + p = poly.Polynomial([1, 2, 3], symbol='z') + assert_equal(self.as_latex(p), + r'$z \mapsto 1.0 + 2.0\,z + 3.0\,z^{2}$') + + # translated input + p = poly.Polynomial([1, 2, 3], domain=[-2, 0], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(1.0 + z\right) + 3.0\,' + r'\left(1.0 + z\right)^{2}$' + ), + ) + + # scaled input + p = poly.Polynomial([1, 2, 3], domain=[-0.5, 0.5], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(2.0z\right) + 3.0\,' + r'\left(2.0z\right)^{2}$' + ), + ) + + # affine input + p = poly.Polynomial([1, 2, 3], domain=[-1, 0], symbol='z') + assert_equal( + self.as_latex(p), + ( + r'$z \mapsto 1.0 + 2.0\,\left(1.0 + 2.0z\right) + 3.0\,' + r'\left(1.0 + 2.0z\right)^{2}$' + ), + ) + + +SWITCH_TO_EXP = ( + '1.0 + (1.0e-01) x + (1.0e-02) x**2', + '1.2 + (1.2e-01) x + (1.2e-02) x**2', + '1.23 + 0.12 x + (1.23e-02) x**2 + (1.23e-03) x**3', + '1.235 + 0.123 x + (1.235e-02) x**2 + (1.235e-03) x**3', + '1.2346 + 0.1235 x + 0.0123 x**2 + (1.2346e-03) x**3 + (1.2346e-04) x**4', + '1.23457 + 0.12346 x + 0.01235 x**2 + (1.23457e-03) x**3 + ' + '(1.23457e-04) x**4', + '1.234568 + 0.123457 x + 0.012346 x**2 + 0.001235 x**3 + ' + '(1.234568e-04) x**4 + (1.234568e-05) x**5', + '1.2345679 + 0.1234568 x + 0.0123457 x**2 + 0.0012346 x**3 + ' + '(1.2345679e-04) x**4 + (1.2345679e-05) x**5') + +class TestPrintOptions: + """ + Test the output is properly configured via printoptions. + The exponential notation is enabled automatically when the values + are too small or too large. + """ + + @pytest.fixture(scope='class', autouse=True) + def use_ascii(self): + poly.set_default_printstyle('ascii') + + def test_str(self): + p = poly.Polynomial([1/2, 1/7, 1/7*10**8, 1/7*10**9]) + assert_equal(str(p), '0.5 + 0.14285714 x + 14285714.28571429 x**2 ' + '+ (1.42857143e+08) x**3') + + with printoptions(precision=3): + assert_equal(str(p), '0.5 + 0.143 x + 14285714.286 x**2 ' + '+ (1.429e+08) x**3') + + def test_latex(self): + p = poly.Polynomial([1/2, 1/7, 1/7*10**8, 1/7*10**9]) + assert_equal(p._repr_latex_(), + r'$x \mapsto \text{0.5} + \text{0.14285714}\,x + ' + r'\text{14285714.28571429}\,x^{2} + ' + r'\text{(1.42857143e+08)}\,x^{3}$') + + with printoptions(precision=3): + assert_equal(p._repr_latex_(), + r'$x \mapsto \text{0.5} + \text{0.143}\,x + ' + r'\text{14285714.286}\,x^{2} + \text{(1.429e+08)}\,x^{3}$') + + def test_fixed(self): + p = poly.Polynomial([1/2]) + assert_equal(str(p), '0.5') + + with printoptions(floatmode='fixed'): + assert_equal(str(p), '0.50000000') + + with printoptions(floatmode='fixed', precision=4): + assert_equal(str(p), '0.5000') + + def test_switch_to_exp(self): + for i, s in enumerate(SWITCH_TO_EXP): + with printoptions(precision=i): + p = poly.Polynomial([1.23456789*10**-i + for i in range(i//2+3)]) + assert str(p).replace('\n', ' ') == s + + def test_non_finite(self): + p = poly.Polynomial([nan, inf]) + assert str(p) == 'nan + inf x' + assert p._repr_latex_() == r'$x \mapsto \text{nan} + \text{inf}\,x$' + with printoptions(nanstr='NAN', infstr='INF'): + assert str(p) == 'NAN + INF x' + assert p._repr_latex_() == \ + r'$x \mapsto \text{NAN} + \text{INF}\,x$' diff --git a/lib/python3.12/site-packages/numpy/polynomial/tests/test_symbol.py b/lib/python3.12/site-packages/numpy/polynomial/tests/test_symbol.py new file mode 100644 index 0000000000000000000000000000000000000000..4ea6035ef7a75e6807634ba894e42015c83edb7d --- /dev/null +++ b/lib/python3.12/site-packages/numpy/polynomial/tests/test_symbol.py @@ -0,0 +1,216 @@ +""" +Tests related to the ``symbol`` attribute of the ABCPolyBase class. +""" + +import pytest +import numpy.polynomial as poly +from numpy.core import array +from numpy.testing import assert_equal, assert_raises, assert_ + + +class TestInit: + """ + Test polynomial creation with symbol kwarg. + """ + c = [1, 2, 3] + + def test_default_symbol(self): + p = poly.Polynomial(self.c) + assert_equal(p.symbol, 'x') + + @pytest.mark.parametrize(('bad_input', 'exception'), ( + ('', ValueError), + ('3', ValueError), + (None, TypeError), + (1, TypeError), + )) + def test_symbol_bad_input(self, bad_input, exception): + with pytest.raises(exception): + p = poly.Polynomial(self.c, symbol=bad_input) + + @pytest.mark.parametrize('symbol', ( + 'x', + 'x_1', + 'A', + 'xyz', + 'β', + )) + def test_valid_symbols(self, symbol): + """ + Values for symbol that should pass input validation. + """ + p = poly.Polynomial(self.c, symbol=symbol) + assert_equal(p.symbol, symbol) + + def test_property(self): + """ + 'symbol' attribute is read only. + """ + p = poly.Polynomial(self.c, symbol='x') + with pytest.raises(AttributeError): + p.symbol = 'z' + + def test_change_symbol(self): + p = poly.Polynomial(self.c, symbol='y') + # Create new polynomial from p with different symbol + pt = poly.Polynomial(p.coef, symbol='t') + assert_equal(pt.symbol, 't') + + +class TestUnaryOperators: + p = poly.Polynomial([1, 2, 3], symbol='z') + + def test_neg(self): + n = -self.p + assert_equal(n.symbol, 'z') + + def test_scalarmul(self): + out = self.p * 10 + assert_equal(out.symbol, 'z') + + def test_rscalarmul(self): + out = 10 * self.p + assert_equal(out.symbol, 'z') + + def test_pow(self): + out = self.p ** 3 + assert_equal(out.symbol, 'z') + + +@pytest.mark.parametrize( + 'rhs', + ( + poly.Polynomial([4, 5, 6], symbol='z'), + array([4, 5, 6]), + ), +) +class TestBinaryOperatorsSameSymbol: + """ + Ensure symbol is preserved for numeric operations on polynomials with + the same symbol + """ + p = poly.Polynomial([1, 2, 3], symbol='z') + + def test_add(self, rhs): + out = self.p + rhs + assert_equal(out.symbol, 'z') + + def test_sub(self, rhs): + out = self.p - rhs + assert_equal(out.symbol, 'z') + + def test_polymul(self, rhs): + out = self.p * rhs + assert_equal(out.symbol, 'z') + + def test_divmod(self, rhs): + for out in divmod(self.p, rhs): + assert_equal(out.symbol, 'z') + + def test_radd(self, rhs): + out = rhs + self.p + assert_equal(out.symbol, 'z') + + def test_rsub(self, rhs): + out = rhs - self.p + assert_equal(out.symbol, 'z') + + def test_rmul(self, rhs): + out = rhs * self.p + assert_equal(out.symbol, 'z') + + def test_rdivmod(self, rhs): + for out in divmod(rhs, self.p): + assert_equal(out.symbol, 'z') + + +class TestBinaryOperatorsDifferentSymbol: + p = poly.Polynomial([1, 2, 3], symbol='x') + other = poly.Polynomial([4, 5, 6], symbol='y') + ops = (p.__add__, p.__sub__, p.__mul__, p.__floordiv__, p.__mod__) + + @pytest.mark.parametrize('f', ops) + def test_binops_fails(self, f): + assert_raises(ValueError, f, self.other) + + +class TestEquality: + p = poly.Polynomial([1, 2, 3], symbol='x') + + def test_eq(self): + other = poly.Polynomial([1, 2, 3], symbol='x') + assert_(self.p == other) + + def test_neq(self): + other = poly.Polynomial([1, 2, 3], symbol='y') + assert_(not self.p == other) + + +class TestExtraMethods: + """ + Test other methods for manipulating/creating polynomial objects. + """ + p = poly.Polynomial([1, 2, 3, 0], symbol='z') + + def test_copy(self): + other = self.p.copy() + assert_equal(other.symbol, 'z') + + def test_trim(self): + other = self.p.trim() + assert_equal(other.symbol, 'z') + + def test_truncate(self): + other = self.p.truncate(2) + assert_equal(other.symbol, 'z') + + @pytest.mark.parametrize('kwarg', ( + {'domain': [-10, 10]}, + {'window': [-10, 10]}, + {'kind': poly.Chebyshev}, + )) + def test_convert(self, kwarg): + other = self.p.convert(**kwarg) + assert_equal(other.symbol, 'z') + + def test_integ(self): + other = self.p.integ() + assert_equal(other.symbol, 'z') + + def test_deriv(self): + other = self.p.deriv() + assert_equal(other.symbol, 'z') + + +def test_composition(): + p = poly.Polynomial([3, 2, 1], symbol="t") + q = poly.Polynomial([5, 1, 0, -1], symbol="λ_1") + r = p(q) + assert r.symbol == "λ_1" + + +# +# Class methods that result in new polynomial class instances +# + + +def test_fit(): + x, y = (range(10),)*2 + p = poly.Polynomial.fit(x, y, deg=1, symbol='z') + assert_equal(p.symbol, 'z') + + +def test_froomroots(): + roots = [-2, 2] + p = poly.Polynomial.fromroots(roots, symbol='z') + assert_equal(p.symbol, 'z') + + +def test_identity(): + p = poly.Polynomial.identity(domain=[-1, 1], window=[5, 20], symbol='z') + assert_equal(p.symbol, 'z') + + +def test_basis(): + p = poly.Polynomial.basis(3, symbol='z') + assert_equal(p.symbol, 'z') diff --git a/lib/python3.12/site-packages/numpy/py.typed b/lib/python3.12/site-packages/numpy/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/numpy/version.py b/lib/python3.12/site-packages/numpy/version.py new file mode 100644 index 0000000000000000000000000000000000000000..fb108fcb4742144d1255ea5f3364fb6cb1c1752f --- /dev/null +++ b/lib/python3.12/site-packages/numpy/version.py @@ -0,0 +1,8 @@ + +version = "1.26.4" +__version__ = version +full_version = version + +git_revision = "9815c16f449e12915ef35a8255329ba26dacd5c0" +release = 'dev' not in version and '+' not in version +short_version = version.split("+")[0] diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/INSTALLER b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/LICENSE b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..754db5afcb8260c5539b1e0ede3990b8ddbc3e29 --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/LICENSE @@ -0,0 +1,20 @@ +ISC LICENSE + + This license is approved by the OSI and FSF as GPL-compatible. + http://opensource.org/licenses/isc-license.txt + + Copyright (c) 2013-2014, Pexpect development team + Copyright (c) 2012, Noah Spurrier + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/METADATA b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..c5ea4a1a59ae2f7d7d1bdb4efcc4c3d40caeb85e --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/METADATA @@ -0,0 +1,52 @@ +Metadata-Version: 2.1 +Name: pexpect +Version: 4.9.0 +Summary: Pexpect allows easy control of interactive console applications. +Home-page: https://pexpect.readthedocs.io/ +Author: Noah Spurrier; Thomas Kluyver; Jeff Quast +Author-email: noah@noah.org, thomas@kluyver.me.uk, contact@jeffquast.com +License: ISC license +Project-URL: Bug Tracker, https://github.com/pexpect/pexpect/issues +Project-URL: Documentation, https://pexpect.readthedocs.io/ +Project-URL: Source Code, https://github.com/pexpect/pexpect +Project-URL: History, https://pexpect.readthedocs.io/en/stable/history.html +Platform: UNIX +Classifier: Development Status :: 5 - Production/Stable +Classifier: Environment :: Console +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: System Administrators +Classifier: License :: OSI Approved :: ISC License (ISCL) +Classifier: Operating System :: POSIX +Classifier: Operating System :: MacOS :: MacOS X +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 2.7 +Classifier: Programming Language :: Python :: 3 +Classifier: Topic :: Software Development +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Topic :: Software Development :: Quality Assurance +Classifier: Topic :: Software Development :: Testing +Classifier: Topic :: System +Classifier: Topic :: System :: Archiving :: Packaging +Classifier: Topic :: System :: Installation/Setup +Classifier: Topic :: System :: Shells +Classifier: Topic :: System :: Software Distribution +Classifier: Topic :: Terminals +License-File: LICENSE +Requires-Dist: ptyprocess (>=0.5) + + +Pexpect is a pure Python module for spawning child applications; controlling +them; and responding to expected patterns in their output. Pexpect works like +Don Libes' Expect. Pexpect allows your script to spawn a child application and +control it as if a human were typing commands. + +Pexpect can be used for automating interactive applications such as ssh, ftp, +passwd, telnet, etc. It can be used to automate setup scripts for duplicating +software package installations on different servers. It can be used for +automated software testing. Pexpect is in the spirit of Don Libes' Expect, but +Pexpect is pure Python. + +The main features of Pexpect require the pty module in the Python standard +library, which is only available on Unix-like systems. Some features—waiting +for patterns from file descriptors or subprocesses—are also available on +Windows. diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/RECORD b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..1a346cf153a3e766a7c3eaf160ea9bf1e48a25e1 --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/RECORD @@ -0,0 +1,43 @@ +pexpect-4.9.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +pexpect-4.9.0.dist-info/LICENSE,sha256=Skg64cTcc4psi3P-tJB04YNdoCq1qmhvJnUCmQb6Nk0,987 +pexpect-4.9.0.dist-info/METADATA,sha256=LXxaxf3v3koBlOxSL7zAynDpnhoqlW7qREFo-FkWUKo,2463 +pexpect-4.9.0.dist-info/RECORD,, +pexpect-4.9.0.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110 +pexpect-4.9.0.dist-info/top_level.txt,sha256=O-b3UY9VQZkW3yDAeFNatUOKO4GojVWO4TTHoI9-E7k,8 +pexpect/ANSI.py,sha256=aA-3tdXz_FZ4G7PAqFZi5g1KBGQ6PzJzS0gm3ALZKZw,12177 +pexpect/FSM.py,sha256=tluiyUGMyIH3q_wLG6Ak1NZVuXUAGNDjq6k6BK1q8RY,13419 +pexpect/__init__.py,sha256=SuQYzpVxpzqLwZ1f68ov5Tvcy8Qv_eD1_NyuSIDibaU,4089 +pexpect/__pycache__/ANSI.cpython-312.pyc,, +pexpect/__pycache__/FSM.cpython-312.pyc,, +pexpect/__pycache__/__init__.cpython-312.pyc,, +pexpect/__pycache__/_async.cpython-312.pyc,, +pexpect/__pycache__/_async_pre_await.cpython-312.pyc,, +pexpect/__pycache__/_async_w_await.cpython-312.pyc,, +pexpect/__pycache__/exceptions.cpython-312.pyc,, +pexpect/__pycache__/expect.cpython-312.pyc,, +pexpect/__pycache__/fdpexpect.cpython-312.pyc,, +pexpect/__pycache__/popen_spawn.cpython-312.pyc,, +pexpect/__pycache__/pty_spawn.cpython-312.pyc,, +pexpect/__pycache__/pxssh.cpython-312.pyc,, +pexpect/__pycache__/replwrap.cpython-312.pyc,, +pexpect/__pycache__/run.cpython-312.pyc,, +pexpect/__pycache__/screen.cpython-312.pyc,, +pexpect/__pycache__/socket_pexpect.cpython-312.pyc,, +pexpect/__pycache__/spawnbase.cpython-312.pyc,, +pexpect/__pycache__/utils.cpython-312.pyc,, +pexpect/_async.py,sha256=kNQ073RymFLQgpmNzseyIC50jhBsLtoxsEW-7TRrxys,907 +pexpect/_async_pre_await.py,sha256=vkv3IAhRhVKzxvXydnl7Q7m9MjVC24E0OuveVsKoYRs,3465 +pexpect/_async_w_await.py,sha256=3QtNwnRSxz5L6pCEA8-aYTzSx647cmALPYwk8A_fTt0,3802 +pexpect/bashrc.sh,sha256=uPuvSvtNmi1ZYW1--6L5deebMjfUEq_LE3-A25zwDyo,419 +pexpect/exceptions.py,sha256=A9C1PWbBc2j9AKvnv7UkPCawhFTEGYmeULW0vwbMvXQ,1068 +pexpect/expect.py,sha256=KKtBmx2MYa-yDE715XlHUcloKe5ndBD359a4OYVXD84,13827 +pexpect/fdpexpect.py,sha256=zdSiPvZlBSuqMk_BNBcpXWb3kGJk6MIVI-1NWZv2lU4,5991 +pexpect/popen_spawn.py,sha256=bxLlZLG8BBbRFmv3YLeR38g2ZxoLRIhE_-RJt8dorgE,6159 +pexpect/pty_spawn.py,sha256=ZygSYsdnVJ5acxiNM9gLvLrT2AVqgwJvbDcPaTxxv9E,37382 +pexpect/pxssh.py,sha256=808mJbdWlDocew12LrvNHz9-ba_a1ZfwB8XFRrcest4,24445 +pexpect/replwrap.py,sha256=_I5PDohJ5qWWiXzO9pS32kfcVllJyungubrn5_ATcVY,5951 +pexpect/run.py,sha256=PkTA_IGOHzLVK89tVIzpLJ7I2jANN5o7GDcH5k1Ly8c,6629 +pexpect/screen.py,sha256=-twD4sIEp83nzuYH9lRDzwHfesoTgVGWglsBYWOK7Ks,13704 +pexpect/socket_pexpect.py,sha256=SV1strRUbcplsL-VbMXgL3sn_q4TB_wRFzGyMzxjzRY,4814 +pexpect/spawnbase.py,sha256=SThkQQ25wiSA-7uqveL3hbkS8FnLFAfk-iXgX2OtOY8,21685 +pexpect/utils.py,sha256=1jIhzU7eBvY3pbW3LZoJhCOU2KWqgty5HgQ6VBYIp5U,6019 diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/WHEEL b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..9d8f872bbf2275e6d1785238e90b0321f4b6f323 --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.38.4) +Root-Is-Purelib: true +Tag: py2-none-any +Tag: py3-none-any + diff --git a/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/top_level.txt b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..808fb07afdc86c2930ec8c3be7eb2a49d669a046 --- /dev/null +++ b/lib/python3.12/site-packages/pexpect-4.9.0.dist-info/top_level.txt @@ -0,0 +1 @@ +pexpect diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e42f33b0ca9773db51ef9b0fc8d5cc62da23a6c Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/_compute_docstrings.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/_compute_docstrings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..888af159df1f03729e91fffb27b072514ebda49e Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/_compute_docstrings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/_generated_version.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/_generated_version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d507c17fef24d6278ff0bcce720b442b3f83f49d Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/_generated_version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/acero.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/acero.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0d0e6c54f562346c1a13648f233e980f2b309677 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/acero.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/benchmark.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/benchmark.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e29003a26970365b808313104911148aa91268f Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/benchmark.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/cffi.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/cffi.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a9cf798510c7bc196e2183518aa0f2e6836c7e3 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/cffi.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/compute.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/compute.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c7891306d184a609c6ad3d4e9809d560a4dda2a Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/compute.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/conftest.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/conftest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e08098d0157bb7f7e78a58ee267f1993cf855068 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/conftest.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/csv.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/csv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2fb3f5bebf1d922a285528fb36eb4ee22dd327b Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/csv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/cuda.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/cuda.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98dfbc0dc4d477c4dcfe0fcfe295e780083d1604 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/cuda.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/dataset.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/dataset.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6e8911db507e550e78066d456cddc4c37426b98 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/dataset.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/feather.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/feather.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3d65af2364cb44ec05f4a8b3eb1fc65434d8555 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/feather.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/flight.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/flight.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb36b1fd0682f5224cc4a9cf9c08c65e8c082237 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/flight.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/fs.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/fs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbf99ef3f59beded39aeddd96402bd39b167dcdf Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/fs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/ipc.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/ipc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9111a8713d143bc6ecd147f2dec30638655176df Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/ipc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/json.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/json.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6fbff15d3d2b3773bbbe6639f18a9338610a211 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/json.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/jvm.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/jvm.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15718fe9d6fa1bc3374107a48e9c8bc314a5dd5b Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/jvm.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/orc.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/orc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a915d4c17ec50c5efc98c744f0bdadf2664912f Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/orc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/pandas_compat.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/pandas_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..337e656f2c37a8babc1f23eb081972f60756ce22 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/pandas_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/substrait.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/substrait.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..327965c18b9c6d8c0f40d4bd1b2b88118fd118cb Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/substrait.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c5384202c6317d7a32f5361576108d182d6e3c1 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..889e6111f37ff9129cc84d107ab1738f44be81a0 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/api/io.h b/lib/python3.12/site-packages/pyarrow/include/parquet/api/io.h new file mode 100644 index 0000000000000000000000000000000000000000..28a00f12a7a616136beb328d20120d6458294eab --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/api/io.h @@ -0,0 +1,20 @@ +// 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. + +#pragma once + +#include "parquet/exception.h" diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/api/reader.h b/lib/python3.12/site-packages/pyarrow/include/parquet/api/reader.h new file mode 100644 index 0000000000000000000000000000000000000000..7e746e8c5bbf551e84431552f688a493e2d62bc4 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/api/reader.h @@ -0,0 +1,35 @@ +// 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. + +#pragma once + +// Column reader API +#include "parquet/column_reader.h" +#include "parquet/column_scanner.h" +#include "parquet/exception.h" +#include "parquet/file_reader.h" +#include "parquet/metadata.h" +#include "parquet/platform.h" +#include "parquet/printer.h" +#include "parquet/properties.h" +#include "parquet/statistics.h" + +// Schemas +#include "parquet/api/schema.h" + +// IO +#include "parquet/api/io.h" diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/api/schema.h b/lib/python3.12/site-packages/pyarrow/include/parquet/api/schema.h new file mode 100644 index 0000000000000000000000000000000000000000..7ca714f47b5448974c460e424ab3821d10f7a384 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/api/schema.h @@ -0,0 +1,21 @@ +// 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. + +#pragma once + +// Schemas +#include "parquet/schema.h" diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/api/writer.h b/lib/python3.12/site-packages/pyarrow/include/parquet/api/writer.h new file mode 100644 index 0000000000000000000000000000000000000000..b072dcf74dea7233723ae55599d95be47c674716 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/api/writer.h @@ -0,0 +1,25 @@ +// 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. + +#pragma once + +#include "parquet/api/io.h" +#include "parquet/api/schema.h" +#include "parquet/column_writer.h" +#include "parquet/exception.h" +#include "parquet/file_writer.h" +#include "parquet/statistics.h" diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/reader.h b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/reader.h new file mode 100644 index 0000000000000000000000000000000000000000..54620b3d0f564bb3e6680193c61a8fa843f38f63 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/reader.h @@ -0,0 +1,392 @@ +// 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. + +#pragma once + +#include +// N.B. we don't include async_generator.h as it's relatively heavy +#include +#include +#include + +#include "parquet/file_reader.h" +#include "parquet/platform.h" +#include "parquet/properties.h" + +namespace arrow { + +class ChunkedArray; +class KeyValueMetadata; +class RecordBatchReader; +struct Scalar; +class Schema; +class Table; +class RecordBatch; + +} // namespace arrow + +namespace parquet { + +class FileMetaData; +class SchemaDescriptor; + +namespace arrow { + +class ColumnChunkReader; +class ColumnReader; +struct SchemaManifest; +class RowGroupReader; + +/// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches. +/// +/// This interfaces caters for different use cases and thus provides different +/// interfaces. In its most simplistic form, we cater for a user that wants to +/// read the whole Parquet at once with the `FileReader::ReadTable` method. +/// +/// More advanced users that also want to implement parallelism on top of each +/// single Parquet files should do this on the RowGroup level. For this, they can +/// call `FileReader::RowGroup(i)->ReadTable` to receive only the specified +/// RowGroup as a table. +/// +/// In the most advanced situation, where a consumer wants to independently read +/// RowGroups in parallel and consume each column individually, they can call +/// `FileReader::RowGroup(i)->Column(j)->Read` and receive an `arrow::Column` +/// instance. +/// +/// Finally, one can also get a stream of record batches using +/// `FileReader::GetRecordBatchReader()`. This can internally decode columns +/// in parallel if use_threads was enabled in the ArrowReaderProperties. +/// +/// The parquet format supports an optional integer field_id which can be assigned +/// to a field. Arrow will convert these field IDs to a metadata key named +/// PARQUET:field_id on the appropriate field. +// TODO(wesm): nested data does not always make sense with this user +// interface unless you are only reading a single leaf node from a branch of +// a table. For example: +// +// repeated group data { +// optional group record { +// optional int32 val1; +// optional byte_array val2; +// optional bool val3; +// } +// optional int32 val4; +// } +// +// In the Parquet file, there are 4 leaf nodes: +// +// * data.record.val1 +// * data.record.val2 +// * data.record.val3 +// * data.val4 +// +// When materializing this data in an Arrow array, we would have: +// +// data: list), +// val3: bool, +// >, +// val4: int32 +// >> +// +// However, in the Parquet format, each leaf node has its own repetition and +// definition levels describing the structure of the intermediate nodes in +// this array structure. Thus, we will need to scan the leaf data for a group +// of leaf nodes part of the same type tree to create a single result Arrow +// nested array structure. +// +// This is additionally complicated "chunky" repeated fields or very large byte +// arrays +class PARQUET_EXPORT FileReader { + public: + /// Factory function to create a FileReader from a ParquetFileReader and properties + /// \deprecated Deprecated in 23.0.0. Use arrow::Result version instead. + ARROW_DEPRECATED("Deprecated in 23.0.0. Use arrow::Result version instead.") + static ::arrow::Status Make(::arrow::MemoryPool* pool, + std::unique_ptr reader, + const ArrowReaderProperties& properties, + std::unique_ptr* out); + + /// Factory function to create a FileReader from a ParquetFileReader + /// \deprecated Deprecated in 23.0.0. Use arrow::Result version instead. + ARROW_DEPRECATED("Deprecated in 23.0.0. Use arrow::Result version instead.") + static ::arrow::Status Make(::arrow::MemoryPool* pool, + std::unique_ptr reader, + std::unique_ptr* out); + + /// Factory function to create a FileReader from a ParquetFileReader and properties + static ::arrow::Result> Make( + ::arrow::MemoryPool* pool, std::unique_ptr reader, + const ArrowReaderProperties& properties); + + /// Factory function to create a FileReader from a ParquetFileReader + static ::arrow::Result> Make( + ::arrow::MemoryPool* pool, std::unique_ptr reader); + + // Since the distribution of columns amongst a Parquet file's row groups may + // be uneven (the number of values in each column chunk can be different), we + // provide a column-oriented read interface. The ColumnReader hides the + // details of paging through the file's row groups and yielding + // fully-materialized arrow::Array instances + // + // Returns error status if the column of interest is not flat. + // The indicated column index is relative to the schema + virtual ::arrow::Status GetColumn(int i, std::unique_ptr* out) = 0; + + /// \brief Return arrow schema for all the columns. + virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0; + + /// \brief Read column as a whole into a chunked array. + /// + /// The index i refers the index of the top level schema field, which may + /// be nested or flat - e.g. + /// + /// 0 foo.bar + /// foo.bar.baz + /// foo.qux + /// 1 foo2 + /// 2 foo3 + /// + /// i=0 will read the entire foo struct, i=1 the foo2 primitive column etc + virtual ::arrow::Status ReadColumn(int i, + std::shared_ptr<::arrow::ChunkedArray>* out) = 0; + + /// \brief Return a RecordBatchReader of all row groups and columns. + virtual ::arrow::Result> + GetRecordBatchReader() = 0; + + /// \brief Return a RecordBatchReader of row groups selected from row_group_indices. + /// + /// Note that the ordering in row_group_indices matters. FileReaders must outlive + /// their RecordBatchReaders. + /// + /// \returns error Result if row_group_indices contains an invalid index + virtual ::arrow::Result> + GetRecordBatchReader(const std::vector& row_group_indices) = 0; + + /// \brief Return a RecordBatchReader of row groups selected from + /// row_group_indices, whose columns are selected by column_indices. + /// + /// Note that the ordering in row_group_indices and column_indices + /// matter. FileReaders must outlive their RecordBatchReaders. + /// + /// \returns error Result if either row_group_indices or column_indices + /// contains an invalid index + virtual ::arrow::Result> + GetRecordBatchReader(const std::vector& row_group_indices, + const std::vector& column_indices) = 0; + + /// \brief Return a RecordBatchReader of row groups selected from + /// row_group_indices, whose columns are selected by column_indices. + /// + /// Note that the ordering in row_group_indices and column_indices + /// matter. FileReaders must outlive their RecordBatchReaders. + /// + /// \param row_group_indices which row groups to read (order determines read order). + /// \param column_indices which columns to read (order determines output schema). + /// \param[out] out record batch stream from parquet data. + /// + /// \returns error Status if either row_group_indices or column_indices + /// contains an invalid index + /// \deprecated Deprecated in 21.0.0. Use arrow::Result version instead. + ARROW_DEPRECATED("Deprecated in 21.0.0. Use arrow::Result version instead.") + ::arrow::Status GetRecordBatchReader(const std::vector& row_group_indices, + const std::vector& column_indices, + std::shared_ptr<::arrow::RecordBatchReader>* out); + + /// \deprecated Deprecated in 21.0.0. Use arrow::Result version instead. + ARROW_DEPRECATED("Deprecated in 21.0.0. Use arrow::Result version instead.") + ::arrow::Status GetRecordBatchReader(const std::vector& row_group_indices, + std::shared_ptr<::arrow::RecordBatchReader>* out); + + /// \deprecated Deprecated in 21.0.0. Use arrow::Result version instead. + ARROW_DEPRECATED("Deprecated in 21.0.0. Use arrow::Result version instead.") + ::arrow::Status GetRecordBatchReader(std::shared_ptr<::arrow::RecordBatchReader>* out); + + /// \brief Return a generator of record batches. + /// + /// The FileReader must outlive the generator, so this requires that you pass in a + /// shared_ptr. + /// + /// \returns error Result if either row_group_indices or column_indices contains an + /// invalid index + virtual ::arrow::Result< + std::function<::arrow::Future>()>> + GetRecordBatchGenerator(std::shared_ptr reader, + const std::vector row_group_indices, + const std::vector column_indices, + ::arrow::internal::Executor* cpu_executor = NULLPTR, + int64_t rows_to_readahead = 0) = 0; + + /// Read all columns into a Table + virtual ::arrow::Status ReadTable(std::shared_ptr<::arrow::Table>* out) = 0; + + /// \brief Read the given columns into a Table + /// + /// The indicated column indices are relative to the internal representation + /// of the parquet table. For instance : + /// 0 foo.bar + /// foo.bar.baz 0 + /// foo.bar.baz2 1 + /// foo.qux 2 + /// 1 foo2 3 + /// 2 foo3 4 + /// + /// i=0 will read foo.bar.baz, i=1 will read only foo.bar.baz2 and so on. + /// Only leaf fields have indices; foo itself doesn't have an index. + /// To get the index for a particular leaf field, one can use + /// manifest().schema_fields to get the top level fields, and then walk the + /// tree to identify the relevant leaf fields and access its column_index. + /// To get the total number of leaf fields, use FileMetadata.num_columns(). + virtual ::arrow::Status ReadTable(const std::vector& column_indices, + std::shared_ptr<::arrow::Table>* out) = 0; + + virtual ::arrow::Status ReadRowGroup(int i, const std::vector& column_indices, + std::shared_ptr<::arrow::Table>* out) = 0; + + virtual ::arrow::Status ReadRowGroup(int i, std::shared_ptr<::arrow::Table>* out) = 0; + + virtual ::arrow::Status ReadRowGroups(const std::vector& row_groups, + const std::vector& column_indices, + std::shared_ptr<::arrow::Table>* out) = 0; + + virtual ::arrow::Status ReadRowGroups(const std::vector& row_groups, + std::shared_ptr<::arrow::Table>* out) = 0; + + /// \brief Scan file contents with one thread, return number of rows + virtual ::arrow::Status ScanContents(std::vector columns, + const int32_t column_batch_size, + int64_t* num_rows) = 0; + + /// \brief Return a reader for the RowGroup, this object must not outlive the + /// FileReader. + virtual std::shared_ptr RowGroup(int row_group_index) = 0; + + /// \brief The number of row groups in the file + virtual int num_row_groups() const = 0; + + virtual ParquetFileReader* parquet_reader() const = 0; + + /// Set whether to use multiple threads during reads of multiple columns. + /// By default only one thread is used. + virtual void set_use_threads(bool use_threads) = 0; + + /// Set number of records to read per batch for the RecordBatchReader. + virtual void set_batch_size(int64_t batch_size) = 0; + + virtual const ArrowReaderProperties& properties() const = 0; + + virtual const SchemaManifest& manifest() const = 0; + + virtual ~FileReader() = default; +}; + +class RowGroupReader { + public: + virtual ~RowGroupReader() = default; + virtual std::shared_ptr Column(int column_index) = 0; + virtual ::arrow::Status ReadTable(const std::vector& column_indices, + std::shared_ptr<::arrow::Table>* out) = 0; + virtual ::arrow::Status ReadTable(std::shared_ptr<::arrow::Table>* out) = 0; + + private: + struct Iterator; +}; + +class ColumnChunkReader { + public: + virtual ~ColumnChunkReader() = default; + virtual ::arrow::Status Read(std::shared_ptr<::arrow::ChunkedArray>* out) = 0; +}; + +// At this point, the column reader is a stream iterator. It only knows how to +// read the next batch of values for a particular column from the file until it +// runs out. +// +// We also do not expose any internal Parquet details, such as row groups. This +// might change in the future. +class PARQUET_EXPORT ColumnReader { + public: + virtual ~ColumnReader() = default; + + // Scan the next array of the indicated size. The actual size of the + // returned array may be less than the passed size depending how much data is + // available in the file. + // + // When all the data in the file has been exhausted, the result is set to + // nullptr. + // + // Returns Status::OK on a successful read, including if you have exhausted + // the data available in the file. + virtual ::arrow::Status NextBatch(int64_t batch_size, + std::shared_ptr<::arrow::ChunkedArray>* out) = 0; +}; + +/// \brief Experimental helper class for bindings (like Python) that struggle +/// either with std::move or C++ exceptions +class PARQUET_EXPORT FileReaderBuilder { + public: + FileReaderBuilder(); + + /// Create FileReaderBuilder from Arrow file and optional properties / metadata + ::arrow::Status Open(std::shared_ptr<::arrow::io::RandomAccessFile> file, + const ReaderProperties& properties = default_reader_properties(), + std::shared_ptr metadata = NULLPTR); + + /// Create FileReaderBuilder from file path and optional properties / metadata + ::arrow::Status OpenFile(const std::string& path, bool memory_map = false, + const ReaderProperties& props = default_reader_properties(), + std::shared_ptr metadata = NULLPTR); + + ParquetFileReader* raw_reader() { return raw_reader_.get(); } + + /// Set Arrow MemoryPool for memory allocation + FileReaderBuilder* memory_pool(::arrow::MemoryPool* pool); + /// Set Arrow reader properties + FileReaderBuilder* properties(const ArrowReaderProperties& arg_properties); + /// Build FileReader instance + ::arrow::Status Build(std::unique_ptr* out); + ::arrow::Result> Build(); + + private: + ::arrow::MemoryPool* pool_; + ArrowReaderProperties properties_; + std::unique_ptr raw_reader_; +}; + +/// \defgroup parquet-arrow-reader-factories Factory functions for Parquet Arrow readers +/// +/// @{ + +/// \brief Build FileReader from Arrow file and MemoryPool +/// +/// Advanced settings are supported through the FileReaderBuilder class. +PARQUET_EXPORT +::arrow::Result> OpenFile( + std::shared_ptr<::arrow::io::RandomAccessFile>, ::arrow::MemoryPool* allocator); + +/// @} + +PARQUET_EXPORT +::arrow::Status StatisticsAsScalars(const Statistics& Statistics, + std::shared_ptr<::arrow::Scalar>* min, + std::shared_ptr<::arrow::Scalar>* max); + +} // namespace arrow +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/schema.h b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/schema.h new file mode 100644 index 0000000000000000000000000000000000000000..dd60fde43422889c53ebd7cf86fbac99c8c6f282 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/schema.h @@ -0,0 +1,184 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/result.h" +#include "arrow/status.h" +#include "arrow/type.h" +#include "arrow/type_fwd.h" + +#include "parquet/level_conversion.h" +#include "parquet/platform.h" +#include "parquet/schema.h" + +namespace parquet { + +class ArrowReaderProperties; +class ArrowWriterProperties; +class WriterProperties; + +namespace arrow { + +/// \defgroup arrow-to-parquet-schema-conversion Functions to convert an Arrow +/// schema into a Parquet schema. +/// +/// @{ + +PARQUET_EXPORT +::arrow::Status FieldToNode(const std::shared_ptr<::arrow::Field>& field, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + schema::NodePtr* out); + +PARQUET_EXPORT +::arrow::Status ToParquetSchema(const ::arrow::Schema* arrow_schema, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + std::shared_ptr* out); + +PARQUET_EXPORT +::arrow::Status ToParquetSchema(const ::arrow::Schema* arrow_schema, + const WriterProperties& properties, + std::shared_ptr* out); + +/// @} + +/// \defgroup parquet-to-arrow-schema-conversion Functions to convert a Parquet +/// schema into an Arrow schema. +/// +/// @{ + +PARQUET_EXPORT +::arrow::Status FromParquetSchema( + const SchemaDescriptor* parquet_schema, const ArrowReaderProperties& properties, + const std::shared_ptr& key_value_metadata, + std::shared_ptr<::arrow::Schema>* out); + +PARQUET_EXPORT +::arrow::Status FromParquetSchema(const SchemaDescriptor* parquet_schema, + const ArrowReaderProperties& properties, + std::shared_ptr<::arrow::Schema>* out); + +PARQUET_EXPORT +::arrow::Status FromParquetSchema(const SchemaDescriptor* parquet_schema, + std::shared_ptr<::arrow::Schema>* out); + +/// @} + +/// \brief Bridge between an arrow::Field and parquet column indices. +struct PARQUET_EXPORT SchemaField { + std::shared_ptr<::arrow::Field> field; + std::vector children; + + // Only set for leaf nodes + int column_index = -1; + + parquet::internal::LevelInfo level_info; + + bool is_leaf() const { return column_index != -1; } +}; + +/// \brief Bridge between a parquet Schema and an arrow Schema. +/// +/// Expose parquet columns as a tree structure. Useful traverse and link +/// between arrow's Schema and parquet's Schema. +struct PARQUET_EXPORT SchemaManifest { + static ::arrow::Status Make( + const SchemaDescriptor* schema, + const std::shared_ptr& metadata, + const ArrowReaderProperties& properties, SchemaManifest* manifest); + + const SchemaDescriptor* descr; + std::shared_ptr<::arrow::Schema> origin_schema; + std::shared_ptr schema_metadata; + std::vector schema_fields; + + std::unordered_map column_index_to_field; + std::unordered_map child_to_parent; + + ::arrow::Status GetColumnField(int column_index, const SchemaField** out) const { + auto it = column_index_to_field.find(column_index); + if (it == column_index_to_field.end()) { + return ::arrow::Status::KeyError("Column index ", column_index, + " not found in schema manifest, may be malformed"); + } + *out = it->second; + return ::arrow::Status::OK(); + } + + const SchemaField* GetParent(const SchemaField* field) const { + // Returns nullptr also if not found + auto it = child_to_parent.find(field); + if (it == child_to_parent.end()) { + return NULLPTR; + } + return it->second; + } + + /// Coalesce a list of field indices (relative to the equivalent arrow::Schema) which + /// correspond to the column root (first node below the parquet schema's root group) of + /// each leaf referenced in column_indices. + /// + /// For example, for leaves `a.b.c`, `a.b.d.e`, and `i.j.k` (column_indices=[0,1,3]) + /// the roots are `a` and `i` (return=[0,2]). + /// + /// root + /// -- a <------ + /// -- -- b | | + /// -- -- -- c | + /// -- -- -- d | + /// -- -- -- -- e + /// -- f + /// -- -- g + /// -- -- -- h + /// -- i <--- + /// -- -- j | + /// -- -- -- k + ::arrow::Result> GetFieldIndices( + const std::vector& column_indices) const { + const schema::GroupNode* group = descr->group_node(); + std::unordered_set already_added; + + std::vector out; + for (int column_idx : column_indices) { + if (column_idx < 0 || column_idx >= descr->num_columns()) { + return ::arrow::Status::IndexError("Column index ", column_idx, " is not valid"); + } + + auto field_node = descr->GetColumnRoot(column_idx); + auto field_idx = group->FieldIndex(*field_node); + if (field_idx == -1) { + return ::arrow::Status::IndexError("Column index ", column_idx, " is not valid"); + } + + if (already_added.insert(field_idx).second) { + out.push_back(field_idx); + } + } + return out; + } +}; + +} // namespace arrow +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/test_util.h b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/test_util.h new file mode 100644 index 0000000000000000000000000000000000000000..05f6fd24ac038e4609f5f182dbff64dae4d21e4d --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/test_util.h @@ -0,0 +1,487 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" +#include "arrow/array/builder_binary.h" +#include "arrow/array/builder_decimal.h" +#include "arrow/array/builder_primitive.h" +#include "arrow/testing/gtest_util.h" +#include "arrow/testing/random.h" +#include "arrow/type_fwd.h" +#include "arrow/type_traits.h" +#include "arrow/util/decimal.h" +#include "arrow/util/float16.h" +#include "parquet/column_reader.h" +#include "parquet/test_util.h" + +namespace parquet { + +using internal::RecordReader; + +namespace arrow { + +using ::arrow::Array; +using ::arrow::ChunkedArray; +using ::arrow::Status; + +template > +struct DecimalWithPrecisionAndScale { + using type = T; + static_assert(PRECISION >= T::kMinPrecision && PRECISION <= T::kMaxPrecision, + "Invalid precision value"); + static constexpr ::arrow::Type::type type_id = T::type_id; + static constexpr int32_t precision = PRECISION; + static constexpr int32_t scale = PRECISION - 1; +}; +template +using Decimal32WithPrecisionAndScale = + DecimalWithPrecisionAndScale<::arrow::Decimal32Type, PRECISION>; +template +using Decimal64WithPrecisionAndScale = + DecimalWithPrecisionAndScale<::arrow::Decimal64Type, PRECISION>; +template +using Decimal128WithPrecisionAndScale = + DecimalWithPrecisionAndScale<::arrow::Decimal128Type, PRECISION>; +template +using Decimal256WithPrecisionAndScale = + DecimalWithPrecisionAndScale<::arrow::Decimal256Type, PRECISION>; + +template +::arrow::enable_if_floating_point NonNullArray( + size_t size, std::shared_ptr* out) { + using c_type = typename ArrowType::c_type; + std::vector values; + if constexpr (::arrow::is_half_float_type::value) { + values.resize(size); + test::random_float16_numbers(static_cast(size), 0, ::arrow::util::Float16(0.0f), + ::arrow::util::Float16(1.0f), values.data()); + } else { + ::arrow::random_real(size, 0, static_cast(0), static_cast(1), + &values); + } + ::arrow::NumericBuilder builder; + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size())); + return builder.Finish(out); +} + +template +::arrow::enable_if_integer NonNullArray(size_t size, + std::shared_ptr* out) { + std::vector values; + ::arrow::randint(size, 0, 64, &values); + + // Passing data type so this will work with TimestampType too + ::arrow::NumericBuilder builder(std::make_shared(), + ::arrow::default_memory_pool()); + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size())); + return builder.Finish(out); +} + +template +::arrow::enable_if_date NonNullArray(size_t size, + std::shared_ptr* out) { + std::vector values; + ::arrow::randint(size, 0, 24, &values); + for (size_t i = 0; i < size; i++) { + values[i] *= 86400000; + } + + // Passing data type so this will work with TimestampType too + ::arrow::NumericBuilder builder(std::make_shared(), + ::arrow::default_memory_pool()); + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size())); + return builder.Finish(out); +} + +template +::arrow::enable_if_base_binary NonNullArray( + size_t size, std::shared_ptr* out) { + using BuilderType = typename ::arrow::TypeTraits::BuilderType; + BuilderType builder; + for (size_t i = 0; i < size; i++) { + RETURN_NOT_OK(builder.Append("test-string")); + } + return builder.Finish(out); +} + +template +::arrow::enable_if_fixed_size_binary NonNullArray( + size_t size, std::shared_ptr* out) { + using BuilderType = typename ::arrow::TypeTraits::BuilderType; + // set byte_width to the length of "fixed": 5 + // todo: find a way to generate test data with more diversity. + BuilderType builder(::arrow::fixed_size_binary(5)); + for (size_t i = 0; i < size; i++) { + RETURN_NOT_OK(builder.Append("fixed")); + } + return builder.Finish(out); +} + +template +static void random_decimals(int64_t n, uint32_t seed, int32_t precision, uint8_t* out) { + auto gen = ::arrow::random::RandomArrayGenerator(seed); + std::shared_ptr decimals; + if constexpr (byte_width == 4) { + decimals = gen.Decimal32(::arrow::decimal32(precision, 0), n); + } else if constexpr (byte_width == 8) { + decimals = gen.Decimal64(::arrow::decimal64(precision, 0), n); + } else if constexpr (byte_width == 16) { + decimals = gen.Decimal128(::arrow::decimal128(precision, 0), n); + } else { + decimals = gen.Decimal256(::arrow::decimal256(precision, 0), n); + } + std::memcpy(out, decimals->data()->GetValues(1, 0), byte_width * n); +} + +template +::arrow::enable_if_t>, + Status> +NonNullArray(size_t size, std::shared_ptr* out) { + constexpr int32_t kDecimalPrecision = precision; + constexpr int32_t kDecimalScale = ArrowType::scale; + + const auto type = + std::make_shared(kDecimalPrecision, kDecimalScale); + const int32_t byte_width = type->byte_width(); + + constexpr int32_t seed = 0; + + ARROW_ASSIGN_OR_RAISE(auto out_buf, ::arrow::AllocateBuffer(size * byte_width)); + random_decimals(size, seed, kDecimalPrecision, + out_buf->mutable_data()); + + using Builder = typename ::arrow::TypeTraits::BuilderType; + Builder builder(type); + RETURN_NOT_OK(builder.AppendValues(out_buf->data(), size)); + return builder.Finish(out); +} + +template +::arrow::enable_if_boolean NonNullArray(size_t size, + std::shared_ptr* out) { + std::vector values; + ::arrow::randint(size, 0, 1, &values); + ::arrow::BooleanBuilder builder; + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size())); + return builder.Finish(out); +} + +// This helper function only supports (size/2) nulls. +template +::arrow::enable_if_floating_point NullableArray( + size_t size, size_t num_nulls, uint32_t seed, std::shared_ptr* out) { + using c_type = typename ArrowType::c_type; + std::vector values; + if constexpr (::arrow::is_half_float_type::value) { + values.resize(size); + test::random_float16_numbers(static_cast(size), 0, ::arrow::util::Float16(-1e4f), + ::arrow::util::Float16(1e4f), values.data()); + } else { + ::arrow::random_real(size, seed, static_cast(-1e10), + static_cast(1e10), &values); + } + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + ::arrow::NumericBuilder builder; + if (values.size() > 0) { + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size(), valid_bytes.data())); + } + return builder.Finish(out); +} + +// This helper function only supports (size/2) nulls. +template +::arrow::enable_if_integer NullableArray(size_t size, size_t num_nulls, + uint32_t seed, + std::shared_ptr* out) { + std::vector values; + + // Seed is random in Arrow right now + (void)seed; + ::arrow::randint(size, 0, 64, &values); + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + // Passing data type so this will work with TimestampType too + ::arrow::NumericBuilder builder(std::make_shared(), + ::arrow::default_memory_pool()); + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size(), valid_bytes.data())); + return builder.Finish(out); +} + +template +::arrow::enable_if_date NullableArray(size_t size, size_t num_nulls, + uint32_t seed, + std::shared_ptr* out) { + std::vector values; + + // Seed is random in Arrow right now + (void)seed; + ::arrow::randint(size, 0, 24, &values); + for (size_t i = 0; i < size; i++) { + values[i] *= 86400000; + } + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + // Passing data type so this will work with TimestampType too + ::arrow::NumericBuilder builder(std::make_shared(), + ::arrow::default_memory_pool()); + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size(), valid_bytes.data())); + return builder.Finish(out); +} + +// This helper function only supports (size/2) nulls yet. +template +::arrow::enable_if_base_binary NullableArray( + size_t size, size_t num_nulls, uint32_t seed, std::shared_ptr<::arrow::Array>* out) { + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + using BuilderType = typename ::arrow::TypeTraits::BuilderType; + BuilderType builder; + + const int kBufferSize = 10; + uint8_t buffer[kBufferSize]; + for (size_t i = 0; i < size; i++) { + if (!valid_bytes[i]) { + RETURN_NOT_OK(builder.AppendNull()); + } else { + ::arrow::random_bytes(kBufferSize, seed + static_cast(i), buffer); + if (ArrowType::is_utf8) { + // Trivially force data to be valid UTF8 by making it all ASCII + for (auto& byte : buffer) { + byte &= 0x7f; + } + } + RETURN_NOT_OK(builder.Append(buffer, kBufferSize)); + } + } + return builder.Finish(out); +} + +// This helper function only supports (size/2) nulls yet, +// same as NullableArray(..) +template +::arrow::enable_if_fixed_size_binary NullableArray( + size_t size, size_t num_nulls, uint32_t seed, std::shared_ptr<::arrow::Array>* out) { + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + using BuilderType = typename ::arrow::TypeTraits::BuilderType; + const int byte_width = 10; + BuilderType builder(::arrow::fixed_size_binary(byte_width)); + + const int kBufferSize = byte_width; + uint8_t buffer[kBufferSize]; + for (size_t i = 0; i < size; i++) { + if (!valid_bytes[i]) { + RETURN_NOT_OK(builder.AppendNull()); + } else { + ::arrow::random_bytes(kBufferSize, seed + static_cast(i), buffer); + RETURN_NOT_OK(builder.Append(buffer)); + } + } + return builder.Finish(out); +} + +template +::arrow::enable_if_t>, + Status> +NullableArray(size_t size, size_t num_nulls, uint32_t seed, + std::shared_ptr<::arrow::Array>* out) { + std::vector valid_bytes(size, '\1'); + + for (size_t i = 0; i < num_nulls; ++i) { + valid_bytes[i * 2] = '\0'; + } + + constexpr int32_t kDecimalPrecision = precision; + constexpr int32_t kDecimalScale = ArrowType::scale; + + const auto type = + std::make_shared(kDecimalPrecision, kDecimalScale); + const int32_t byte_width = type->byte_width(); + + ARROW_ASSIGN_OR_RAISE(auto out_buf, ::arrow::AllocateBuffer(size * byte_width)); + random_decimals(size, seed, precision, + out_buf->mutable_data()); + + using Builder = typename ::arrow::TypeTraits::BuilderType; + Builder builder(type); + RETURN_NOT_OK(builder.AppendValues(out_buf->data(), size, valid_bytes.data())); + return builder.Finish(out); +} + +// This helper function only supports (size/2) nulls yet. +template +::arrow::enable_if_boolean NullableArray(size_t size, size_t num_nulls, + uint32_t seed, + std::shared_ptr* out) { + std::vector values; + + // Seed is random in Arrow right now + (void)seed; + + ::arrow::randint(size, 0, 1, &values); + std::vector valid_bytes(size, 1); + + for (size_t i = 0; i < num_nulls; i++) { + valid_bytes[i * 2] = 0; + } + + ::arrow::BooleanBuilder builder; + RETURN_NOT_OK(builder.AppendValues(values.data(), values.size(), valid_bytes.data())); + return builder.Finish(out); +} + +/// Wrap an Array into a ListArray by splitting it up into size lists. +/// +/// This helper function only supports (size/2) nulls. +Status MakeListArray(const std::shared_ptr& values, int64_t size, + int64_t null_count, const std::string& item_name, + bool nullable_values, std::shared_ptr<::arrow::ListArray>* out) { + // We always include an empty list + int64_t non_null_entries = size - null_count - 1; + int64_t length_per_entry = values->length() / non_null_entries; + + auto offsets = AllocateBuffer(); + RETURN_NOT_OK(offsets->Resize((size + 1) * sizeof(int32_t))); + int32_t* offsets_ptr = reinterpret_cast(offsets->mutable_data()); + + auto null_bitmap = AllocateBuffer(); + int64_t bitmap_size = ::arrow::bit_util::BytesForBits(size); + RETURN_NOT_OK(null_bitmap->Resize(bitmap_size)); + uint8_t* null_bitmap_ptr = null_bitmap->mutable_data(); + memset(null_bitmap_ptr, 0, bitmap_size); + + int32_t current_offset = 0; + for (int64_t i = 0; i < size; i++) { + offsets_ptr[i] = current_offset; + if (!(((i % 2) == 0) && ((i / 2) < null_count))) { + // Non-null list (list with index 1 is always empty). + ::arrow::bit_util::SetBit(null_bitmap_ptr, i); + if (i != 1) { + current_offset += static_cast(length_per_entry); + } + } + } + offsets_ptr[size] = static_cast(values->length()); + + auto value_field = ::arrow::field(item_name, values->type(), nullable_values); + *out = std::make_shared<::arrow::ListArray>(::arrow::list(value_field), size, offsets, + values, null_bitmap, null_count); + + return Status::OK(); +} + +// Make an array containing only empty lists, with a null values array +Status MakeEmptyListsArray(int64_t size, std::shared_ptr* out_array) { + // Allocate an offsets buffer containing only zeroes + const int64_t offsets_nbytes = (size + 1) * sizeof(int32_t); + ARROW_ASSIGN_OR_RAISE(auto offsets_buffer, ::arrow::AllocateBuffer(offsets_nbytes)); + memset(offsets_buffer->mutable_data(), 0, offsets_nbytes); + + auto value_field = + ::arrow::field("item", ::arrow::float64(), false /* nullable_values */); + auto list_type = ::arrow::list(value_field); + + std::vector> child_buffers = {nullptr /* null bitmap */, + nullptr /* values */}; + auto child_data = + ::arrow::ArrayData::Make(value_field->type(), 0, std::move(child_buffers)); + + std::vector> buffers = {nullptr /* bitmap */, + std::move(offsets_buffer)}; + auto array_data = ::arrow::ArrayData::Make(list_type, size, std::move(buffers)); + array_data->child_data.push_back(child_data); + + *out_array = ::arrow::MakeArray(array_data); + return Status::OK(); +} + +std::shared_ptr<::arrow::Table> MakeSimpleTable( + const std::shared_ptr& values, bool nullable) { + auto schema = ::arrow::schema({::arrow::field("col", values->type(), nullable)}); + return ::arrow::Table::Make(schema, {values}); +} + +std::shared_ptr<::arrow::Table> MakeSimpleTable(const std::shared_ptr& values, + bool nullable) { + auto carr = std::make_shared<::arrow::ChunkedArray>(values); + return MakeSimpleTable(carr, nullable); +} + +template +void ExpectArray(T* expected, Array* result) { + auto p_array = static_cast<::arrow::PrimitiveArray*>(result); + for (int i = 0; i < result->length(); i++) { + EXPECT_EQ(expected[i], reinterpret_cast(p_array->values()->data())[i]); + } +} + +template +void ExpectArrayT(void* expected, Array* result) { + ::arrow::PrimitiveArray* p_array = static_cast<::arrow::PrimitiveArray*>(result); + for (int64_t i = 0; i < result->length(); i++) { + EXPECT_EQ(reinterpret_cast(expected)[i], + reinterpret_cast( + p_array->values()->data())[i]); + } +} + +template <> +void ExpectArrayT<::arrow::BooleanType>(void* expected, Array* result) { + ::arrow::BooleanBuilder builder; + ARROW_EXPECT_OK( + builder.AppendValues(reinterpret_cast(expected), result->length())); + + std::shared_ptr expected_array; + ARROW_EXPECT_OK(builder.Finish(&expected_array)); + EXPECT_TRUE(result->Equals(*expected_array)); +} + +} // namespace arrow + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/writer.h b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/writer.h new file mode 100644 index 0000000000000000000000000000000000000000..8ec8796ffd17516cef6d24eac481a7b44703d562 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/arrow/writer.h @@ -0,0 +1,176 @@ +// 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. + +#pragma once + +#include +#include + +#include "parquet/platform.h" +#include "parquet/properties.h" + +namespace arrow { + +class Array; +class ChunkedArray; +class RecordBatch; +class Schema; +class Table; + +} // namespace arrow + +namespace parquet { + +class FileMetaData; +class ParquetFileWriter; + +namespace arrow { + +/// \brief Iterative FileWriter class +/// +/// For basic usage, can write a Table at a time, creating one or more row +/// groups per write call. +/// +/// For advanced usage, can write column-by-column: Start a new RowGroup or +/// Chunk with NewRowGroup, then write column-by-column the whole column chunk. +/// +/// If PARQUET:field_id is present as a metadata key on a field, and the corresponding +/// value is a nonnegative integer, then it will be used as the field_id in the parquet +/// file. +class PARQUET_EXPORT FileWriter { + public: + static ::arrow::Status Make(MemoryPool* pool, std::unique_ptr writer, + std::shared_ptr<::arrow::Schema> schema, + std::shared_ptr arrow_properties, + std::unique_ptr* out); + + /// \brief Try to create an Arrow to Parquet file writer. + /// + /// \param schema schema of data that will be passed. + /// \param pool memory pool to use. + /// \param sink output stream to write Parquet data. + /// \param properties general Parquet writer properties. + /// \param arrow_properties Arrow-specific writer properties. + /// + /// \since 11.0.0 + static ::arrow::Result> Open( + const ::arrow::Schema& schema, MemoryPool* pool, + std::shared_ptr<::arrow::io::OutputStream> sink, + std::shared_ptr properties = default_writer_properties(), + std::shared_ptr arrow_properties = + default_arrow_writer_properties()); + + /// Return the Arrow schema to be written to. + virtual std::shared_ptr<::arrow::Schema> schema() const = 0; + + /// \brief Write a Table to Parquet. + /// + /// \param table Arrow table to write. + /// \param chunk_size maximum number of rows to write per row group. + virtual ::arrow::Status WriteTable( + const ::arrow::Table& table, int64_t chunk_size = DEFAULT_MAX_ROW_GROUP_LENGTH) = 0; + + /// \brief Start a new row group. + /// + /// Returns an error if not all columns have been written. + virtual ::arrow::Status NewRowGroup() = 0; + + /// \brief Write ColumnChunk in row group using an array. + virtual ::arrow::Status WriteColumnChunk(const ::arrow::Array& data) = 0; + + /// \brief Write ColumnChunk in row group using slice of a ChunkedArray + virtual ::arrow::Status WriteColumnChunk( + const std::shared_ptr<::arrow::ChunkedArray>& data, int64_t offset, + int64_t size) = 0; + + /// \brief Write ColumnChunk in a row group using a ChunkedArray + virtual ::arrow::Status WriteColumnChunk( + const std::shared_ptr<::arrow::ChunkedArray>& data) = 0; + + /// \brief Start a new buffered row group. + /// + /// Returns an error if not all columns have been written. + virtual ::arrow::Status NewBufferedRowGroup() = 0; + + /// \brief Write a RecordBatch into the buffered row group. + /// + /// Multiple RecordBatches can be written into the same row group + /// through this method. + /// + /// WriterProperties.max_row_group_length() is respected and a new + /// row group will be created if the current row group exceeds the + /// limit. + /// + /// Batches get flushed to the output stream once NewBufferedRowGroup() + /// or Close() is called. + /// + /// WARNING: If you are writing multiple files in parallel in the same + /// executor, deadlock may occur if ArrowWriterProperties::use_threads + /// is set to true to write columns in parallel. Please disable use_threads + /// option in this case. + virtual ::arrow::Status WriteRecordBatch(const ::arrow::RecordBatch& batch) = 0; + + /// \brief Write the footer and close the file. + virtual ::arrow::Status Close() = 0; + virtual ~FileWriter(); + + virtual MemoryPool* memory_pool() const = 0; + /// \brief Add key-value metadata to the file. + /// \param[in] key_value_metadata the metadata to add. + /// \note This will overwrite any existing metadata with the same key. + /// \return Error if Close() has been called. + /// + /// WARNING: If `store_schema` is enabled, `ARROW:schema` would be stored + /// in the key-value metadata. Overwriting this key would result in + /// `store_schema` being unusable during read. + virtual ::arrow::Status AddKeyValueMetadata( + const std::shared_ptr& key_value_metadata) = 0; + /// \brief Return the file metadata, only available after calling Close(). + virtual const std::shared_ptr metadata() const = 0; +}; + +/// \brief Write Parquet file metadata only to indicated Arrow OutputStream +PARQUET_EXPORT +::arrow::Status WriteFileMetaData(const FileMetaData& file_metadata, + ::arrow::io::OutputStream* sink); + +/// \brief Write metadata-only Parquet file to indicated Arrow OutputStream +PARQUET_EXPORT +::arrow::Status WriteMetaDataFile(const FileMetaData& file_metadata, + ::arrow::io::OutputStream* sink); + +/// \brief Write a Table to Parquet. +/// +/// This writes one table in a single shot. To write a Parquet file with +/// multiple tables iteratively, see parquet::arrow::FileWriter. +/// +/// \param table Table to write. +/// \param pool memory pool to use. +/// \param sink output stream to write Parquet data. +/// \param chunk_size maximum number of rows to write per row group. +/// \param properties general Parquet writer properties. +/// \param arrow_properties Arrow-specific writer properties. +::arrow::Status PARQUET_EXPORT +WriteTable(const ::arrow::Table& table, MemoryPool* pool, + std::shared_ptr<::arrow::io::OutputStream> sink, + int64_t chunk_size = DEFAULT_MAX_ROW_GROUP_LENGTH, + std::shared_ptr properties = default_writer_properties(), + std::shared_ptr arrow_properties = + default_arrow_writer_properties()); + +} // namespace arrow +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/benchmark_util.h b/lib/python3.12/site-packages/pyarrow/include/parquet/benchmark_util.h new file mode 100644 index 0000000000000000000000000000000000000000..abfb33bd5ba9e9d3699f0350a110672036949534 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/benchmark_util.h @@ -0,0 +1,47 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "parquet/types.h" + +namespace parquet::benchmarks { + +template +void GenerateBenchmarkData(uint32_t size, uint32_t seed, T* data, + std::vector* heap, uint32_t data_string_length); + +#define _GENERATE_BENCHMARK_DATA_DECL(KLASS) \ + template <> \ + void GenerateBenchmarkData(uint32_t size, uint32_t seed, KLASS* data, \ + std::vector* heap, uint32_t data_string_length); + +_GENERATE_BENCHMARK_DATA_DECL(int32_t) +_GENERATE_BENCHMARK_DATA_DECL(int64_t) +_GENERATE_BENCHMARK_DATA_DECL(float) +_GENERATE_BENCHMARK_DATA_DECL(double) +_GENERATE_BENCHMARK_DATA_DECL(ByteArray) +_GENERATE_BENCHMARK_DATA_DECL(FLBA) +_GENERATE_BENCHMARK_DATA_DECL(Int96) + +#undef _GENERATE_BENCHMARK_DATA_DECL + +} // namespace parquet::benchmarks diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter.h b/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter.h new file mode 100644 index 0000000000000000000000000000000000000000..82172f363ba7ee920ea47f562688e81d704f4b8e --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter.h @@ -0,0 +1,363 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/util/bit_util.h" +#include "arrow/util/logging.h" +#include "parquet/hasher.h" +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace parquet { + +// A Bloom filter is a compact structure to indicate whether an item is not in a set or +// probably in a set. The Bloom filter usually consists of a bit set that represents a +// set of elements, a hash strategy and a Bloom filter algorithm. +class PARQUET_EXPORT BloomFilter { + public: + // Maximum Bloom filter size, it sets to HDFS default block size 128MB + // This value will be reconsidered when implementing Bloom filter producer. + static constexpr uint32_t kMaximumBloomFilterBytes = 128 * 1024 * 1024; + + /// Determine whether an element exist in set or not. + /// + /// @param hash the element to contain. + /// @return false if value is definitely not in set, and true means PROBABLY + /// in set. + virtual bool FindHash(uint64_t hash) const = 0; + + /// Insert element to set represented by Bloom filter bitset. + /// @param hash the hash of value to insert into Bloom filter. + virtual void InsertHash(uint64_t hash) = 0; + + /// Insert elements to set represented by Bloom filter bitset. + /// @param hashes the hash values to insert into Bloom filter. + /// @param num_values the number of hash values to insert. + virtual void InsertHashes(const uint64_t* hashes, int num_values) = 0; + + /// Write this Bloom filter to an output stream. A Bloom filter structure should + /// include bitset length, hash strategy, algorithm, and bitset. + /// + /// @param sink the output stream to write + virtual void WriteTo(ArrowOutputStream* sink) const = 0; + + /// Get the number of bytes of bitset + virtual uint32_t GetBitsetSize() const = 0; + + /// Compute hash for 32 bits value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(int32_t value) const = 0; + + /// Compute hash for 64 bits value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(int64_t value) const = 0; + + /// Compute hash for float value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(float value) const = 0; + + /// Compute hash for double value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(double value) const = 0; + + /// Compute hash for Int96 value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(const Int96* value) const = 0; + + /// Compute hash for ByteArray value by using its plain encoding result. + /// + /// @param value the value to hash. + /// @return hash result. + virtual uint64_t Hash(const ByteArray* value) const = 0; + + /// Compute hash for fixed byte array value by using its plain encoding result. + /// + /// @param value the value address. + /// @param len the value length. + /// @return hash result. + virtual uint64_t Hash(const FLBA* value, uint32_t len) const = 0; + + /// Batch compute hashes for 32 bits values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const int32_t* values, int num_values, uint64_t* hashes) const = 0; + + /// Batch compute hashes for 64 bits values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const int64_t* values, int num_values, uint64_t* hashes) const = 0; + + /// Batch compute hashes for float values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const float* values, int num_values, uint64_t* hashes) const = 0; + + /// Batch compute hashes for double values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const double* values, int num_values, uint64_t* hashes) const = 0; + + /// Batch compute hashes for Int96 values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const Int96* values, int num_values, uint64_t* hashes) const = 0; + + /// Batch compute hashes for ByteArray values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const ByteArray* values, int num_values, + uint64_t* hashes) const = 0; + + /// Batch compute hashes for fixed byte array values by using its plain encoding result. + /// + /// @param values values a pointer to the values to hash. + /// @param type_len the value length. + /// @param num_values the number of values to hash. + /// @param hashes a pointer to the output hash values, its length should be equal to + /// num_values. + virtual void Hashes(const FLBA* values, uint32_t type_len, int num_values, + uint64_t* hashes) const = 0; + + virtual ~BloomFilter() = default; + + protected: + // Hash strategy available for Bloom filter. + enum class HashStrategy : uint32_t { XXHASH = 0 }; + + // Bloom filter algorithm. + enum class Algorithm : uint32_t { BLOCK = 0 }; + + enum class CompressionStrategy : uint32_t { UNCOMPRESSED = 0 }; +}; + +/// The BlockSplitBloomFilter is implemented using block-based Bloom filters from +/// Putze et al.'s "Cache-,Hash- and Space-Efficient Bloom filters". The basic idea is to +/// hash the item to a tiny Bloom filter which size fit a single cache line or smaller. +/// +/// This implementation sets 8 bits in each tiny Bloom filter. Each tiny Bloom +/// filter is 32 bytes to take advantage of 32-byte SIMD instructions. +class PARQUET_EXPORT BlockSplitBloomFilter : public BloomFilter { + public: + /// The constructor of BlockSplitBloomFilter. It uses XXH64 as hash function. + /// + /// \param pool memory pool to use. + explicit BlockSplitBloomFilter( + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + /// Initialize the BlockSplitBloomFilter. The range of num_bytes should be within + /// [kMinimumBloomFilterBytes, kMaximumBloomFilterBytes], it will be + /// rounded up/down to lower/upper bound if num_bytes is out of range and also + /// will be rounded up to a power of 2. + /// + /// @param num_bytes The number of bytes to store Bloom filter bitset. + void Init(uint32_t num_bytes); + + /// Initialize the BlockSplitBloomFilter. It copies the bitset as underlying + /// bitset because the given bitset may not satisfy the 32-byte alignment requirement + /// which may lead to segfault when performing SIMD instructions. It is the caller's + /// responsibility to free the bitset passed in. This is used when reconstructing + /// a Bloom filter from a parquet file. + /// + /// @param bitset The given bitset to initialize the Bloom filter. + /// @param num_bytes The number of bytes of given bitset. + void Init(const uint8_t* bitset, uint32_t num_bytes); + + /// Minimum Bloom filter size, it sets to 32 bytes to fit a tiny Bloom filter. + static constexpr uint32_t kMinimumBloomFilterBytes = 32; + + /// Calculate optimal size according to the number of distinct values and false + /// positive probability. + /// + /// @param ndv The number of distinct values. + /// @param fpp The false positive probability. + /// @return it always return a value between kMinimumBloomFilterBytes and + /// kMaximumBloomFilterBytes, and the return value is always a power of 2 + static uint32_t OptimalNumOfBytes(uint32_t ndv, double fpp) { + uint32_t optimal_num_of_bits = OptimalNumOfBits(ndv, fpp); + ARROW_DCHECK(::arrow::bit_util::IsMultipleOf8(optimal_num_of_bits)); + return optimal_num_of_bits >> 3; + } + + /// Calculate optimal size according to the number of distinct values and false + /// positive probability. + /// + /// @param ndv The number of distinct values. + /// @param fpp The false positive probability. + /// @return it always return a value between kMinimumBloomFilterBytes * 8 and + /// kMaximumBloomFilterBytes * 8, and the return value is always a power of 16 + static uint32_t OptimalNumOfBits(uint32_t ndv, double fpp) { + ARROW_DCHECK(fpp > 0.0 && fpp < 1.0); + const double m = -8.0 * ndv / log(1 - pow(fpp, 1.0 / 8)); + uint32_t num_bits; + + // Handle overflow. + if (m < 0 || m > kMaximumBloomFilterBytes << 3) { + num_bits = static_cast(kMaximumBloomFilterBytes << 3); + } else { + num_bits = static_cast(m); + } + + // Round up to lower bound + if (num_bits < kMinimumBloomFilterBytes << 3) { + num_bits = kMinimumBloomFilterBytes << 3; + } + + // Get next power of 2 if bits is not power of 2. + if ((num_bits & (num_bits - 1)) != 0) { + num_bits = static_cast(::arrow::bit_util::NextPower2(num_bits)); + } + + // Round down to upper bound + if (num_bits > kMaximumBloomFilterBytes << 3) { + num_bits = kMaximumBloomFilterBytes << 3; + } + + return num_bits; + } + + bool FindHash(uint64_t hash) const override; + void InsertHash(uint64_t hash) override; + void InsertHashes(const uint64_t* hashes, int num_values) override; + void WriteTo(ArrowOutputStream* sink) const override; + uint32_t GetBitsetSize() const override { return num_bytes_; } + + uint64_t Hash(int32_t value) const override { return hasher_->Hash(value); } + uint64_t Hash(int64_t value) const override { return hasher_->Hash(value); } + uint64_t Hash(float value) const override { return hasher_->Hash(value); } + uint64_t Hash(double value) const override { return hasher_->Hash(value); } + uint64_t Hash(const Int96* value) const override { return hasher_->Hash(value); } + uint64_t Hash(const ByteArray* value) const override { return hasher_->Hash(value); } + uint64_t Hash(const FLBA* value, uint32_t len) const override { + return hasher_->Hash(value, len); + } + + void Hashes(const int32_t* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const int64_t* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const float* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const double* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const Int96* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const ByteArray* values, int num_values, uint64_t* hashes) const override { + hasher_->Hashes(values, num_values, hashes); + } + void Hashes(const FLBA* values, uint32_t type_len, int num_values, + uint64_t* hashes) const override { + hasher_->Hashes(values, type_len, num_values, hashes); + } + + uint64_t Hash(const int32_t* value) const { return hasher_->Hash(*value); } + uint64_t Hash(const int64_t* value) const { return hasher_->Hash(*value); } + uint64_t Hash(const float* value) const { return hasher_->Hash(*value); } + uint64_t Hash(const double* value) const { return hasher_->Hash(*value); } + + /// Deserialize the Bloom filter from an input stream. It is used when reconstructing + /// a Bloom filter from a parquet filter. + /// + /// @param properties The parquet reader properties. + /// @param input_stream The input stream from which to construct the bloom filter. + /// @param bloom_filter_length The length of the serialized bloom filter including + /// header. + /// @return The BlockSplitBloomFilter. + static BlockSplitBloomFilter Deserialize( + const ReaderProperties& properties, ArrowInputStream* input_stream, + std::optional bloom_filter_length = std::nullopt); + + private: + inline void InsertHashImpl(uint64_t hash); + + // Bytes in a tiny Bloom filter block. + static constexpr int kBytesPerFilterBlock = 32; + + // The number of bits to be set in each tiny Bloom filter + static constexpr int kBitsSetPerBlock = 8; + + // A mask structure used to set bits in each tiny Bloom filter. + struct BlockMask { + uint32_t item[kBitsSetPerBlock]; + }; + + // The block-based algorithm needs eight odd SALT values to calculate eight indexes + // of bit to set, one bit in each 32-bit word. + static constexpr uint32_t SALT[kBitsSetPerBlock] = { + 0x47b6137bU, 0x44974d91U, 0x8824ad5bU, 0xa2b7289dU, + 0x705495c7U, 0x2df1424bU, 0x9efc4947U, 0x5c6bfb31U}; + + // Memory pool to allocate aligned buffer for bitset + ::arrow::MemoryPool* pool_; + + // The underlying buffer of bitset. + std::shared_ptr data_; + + // The number of bytes of Bloom filter bitset. + uint32_t num_bytes_; + + // Hash strategy used in this Bloom filter. + HashStrategy hash_strategy_; + + // Algorithm used in this Bloom filter. + Algorithm algorithm_; + + // Compression used in this Bloom filter. + CompressionStrategy compression_strategy_; + + // The hash pointer points to actual hash class used. + std::unique_ptr hasher_; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter_reader.h b/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter_reader.h new file mode 100644 index 0000000000000000000000000000000000000000..cbd267dd1972dcde98382dda3c84a6a544ddc3e3 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/bloom_filter_reader.h @@ -0,0 +1,68 @@ +// 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. + +#pragma once + +#include "arrow/io/interfaces.h" +#include "parquet/properties.h" +#include "parquet/type_fwd.h" + +namespace parquet { + +class InternalFileDecryptor; +class BloomFilter; + +class PARQUET_EXPORT RowGroupBloomFilterReader { + public: + virtual ~RowGroupBloomFilterReader() = default; + + /// \brief Read bloom filter of a column chunk. + /// + /// \param[in] i column ordinal of the column chunk. + /// \returns bloom filter of the column or nullptr if it does not exist. + /// \throws ParquetException if the index is out of bound, or read bloom + /// filter failed. + virtual std::unique_ptr GetColumnBloomFilter(int i) = 0; +}; + +/// \brief Interface for reading the bloom filter for a Parquet file. +class PARQUET_EXPORT BloomFilterReader { + public: + virtual ~BloomFilterReader() = default; + + /// \brief Create a BloomFilterReader instance. + /// \returns a BloomFilterReader instance. + /// WARNING: The returned BloomFilterReader references to all the input parameters, so + /// it must not outlive all of the input parameters. Usually these input parameters + /// come from the same ParquetFileReader object, so it must not outlive the reader + /// that creates this BloomFilterReader. + static std::unique_ptr Make( + std::shared_ptr<::arrow::io::RandomAccessFile> input, + std::shared_ptr file_metadata, const ReaderProperties& properties, + std::shared_ptr file_decryptor = NULLPTR); + + /// \brief Get the bloom filter reader of a specific row group. + /// \param[in] i row group ordinal to get bloom filter reader. + /// \returns RowGroupBloomFilterReader of the specified row group. A nullptr may or may + /// not be returned if the bloom filter for the row group is unavailable. It + /// is the caller's responsibility to check the return value of follow-up calls + /// to the RowGroupBloomFilterReader. + /// \throws ParquetException if the index is out of bound. + virtual std::shared_ptr RowGroup(int i) = 0; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/column_page.h b/lib/python3.12/site-packages/pyarrow/include/parquet/column_page.h new file mode 100644 index 0000000000000000000000000000000000000000..111265a842ee7e61753a1067e78e80a868936627 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/column_page.h @@ -0,0 +1,179 @@ +// 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. + +// This module defines an abstract interface for iterating through pages in a +// Parquet column chunk within a row group. It could be extended in the future +// to iterate through all data pages in all chunks in a file. + +#pragma once + +#include +#include +#include +#include + +#include "parquet/size_statistics.h" +#include "parquet/statistics.h" +#include "parquet/types.h" + +namespace parquet { + +// TODO: Parallel processing is not yet safe because of memory-ownership +// semantics (the PageReader may or may not own the memory referenced by a +// page) +// +// TODO(wesm): In the future Parquet implementations may store the crc code +// in format::PageHeader. parquet-mr currently does not, so we also skip it +// here, both on the read and write path +class Page { + public: + Page(const std::shared_ptr& buffer, PageType::type type) + : buffer_(buffer), type_(type) {} + + PageType::type type() const { return type_; } + + std::shared_ptr buffer() const { return buffer_; } + + // @returns: a pointer to the page's data + const uint8_t* data() const { return buffer_->data(); } + + // @returns: the total size in bytes of the page's data buffer + int32_t size() const { return static_cast(buffer_->size()); } + + private: + std::shared_ptr buffer_; + PageType::type type_; +}; + +/// \brief Base type for DataPageV1 and DataPageV2 including common attributes +class DataPage : public Page { + public: + int32_t num_values() const { return num_values_; } + Encoding::type encoding() const { return encoding_; } + int64_t uncompressed_size() const { return uncompressed_size_; } + const EncodedStatistics& statistics() const { return statistics_; } + /// Return the row ordinal within the row group to the first row in the data page. + /// Currently it is only present from data pages created by ColumnWriter in order + /// to collect page index. + std::optional first_row_index() const { return first_row_index_; } + const SizeStatistics& size_statistics() const { return size_statistics_; } + + virtual ~DataPage() = default; + + protected: + DataPage(PageType::type type, const std::shared_ptr& buffer, int32_t num_values, + Encoding::type encoding, int64_t uncompressed_size, + EncodedStatistics statistics, std::optional first_row_index, + SizeStatistics size_statistics) + : Page(buffer, type), + num_values_(num_values), + encoding_(encoding), + uncompressed_size_(uncompressed_size), + statistics_(std::move(statistics)), + first_row_index_(std::move(first_row_index)), + size_statistics_(std::move(size_statistics)) {} + + int32_t num_values_; + Encoding::type encoding_; + int64_t uncompressed_size_; + EncodedStatistics statistics_; + /// Row ordinal within the row group to the first row in the data page. + std::optional first_row_index_; + SizeStatistics size_statistics_; +}; + +class DataPageV1 : public DataPage { + public: + DataPageV1(const std::shared_ptr& buffer, int32_t num_values, + Encoding::type encoding, Encoding::type definition_level_encoding, + Encoding::type repetition_level_encoding, int64_t uncompressed_size, + EncodedStatistics statistics = EncodedStatistics(), + std::optional first_row_index = std::nullopt, + SizeStatistics size_statistics = SizeStatistics()) + : DataPage(PageType::DATA_PAGE, buffer, num_values, encoding, uncompressed_size, + std::move(statistics), std::move(first_row_index), + std::move(size_statistics)), + definition_level_encoding_(definition_level_encoding), + repetition_level_encoding_(repetition_level_encoding) {} + + Encoding::type repetition_level_encoding() const { return repetition_level_encoding_; } + + Encoding::type definition_level_encoding() const { return definition_level_encoding_; } + + private: + Encoding::type definition_level_encoding_; + Encoding::type repetition_level_encoding_; +}; + +class DataPageV2 : public DataPage { + public: + DataPageV2(const std::shared_ptr& buffer, int32_t num_values, int32_t num_nulls, + int32_t num_rows, Encoding::type encoding, + int32_t definition_levels_byte_length, int32_t repetition_levels_byte_length, + int64_t uncompressed_size, bool is_compressed = false, + EncodedStatistics statistics = EncodedStatistics(), + std::optional first_row_index = std::nullopt, + SizeStatistics size_statistics = SizeStatistics()) + : DataPage(PageType::DATA_PAGE_V2, buffer, num_values, encoding, uncompressed_size, + std::move(statistics), std::move(first_row_index), + std::move(size_statistics)), + num_nulls_(num_nulls), + num_rows_(num_rows), + definition_levels_byte_length_(definition_levels_byte_length), + repetition_levels_byte_length_(repetition_levels_byte_length), + is_compressed_(is_compressed) {} + + int32_t num_nulls() const { return num_nulls_; } + + int32_t num_rows() const { return num_rows_; } + + int32_t definition_levels_byte_length() const { return definition_levels_byte_length_; } + + int32_t repetition_levels_byte_length() const { return repetition_levels_byte_length_; } + + bool is_compressed() const { return is_compressed_; } + + private: + int32_t num_nulls_; + int32_t num_rows_; + int32_t definition_levels_byte_length_; + int32_t repetition_levels_byte_length_; + bool is_compressed_; +}; + +class DictionaryPage : public Page { + public: + DictionaryPage(const std::shared_ptr& buffer, int32_t num_values, + Encoding::type encoding, bool is_sorted = false) + : Page(buffer, PageType::DICTIONARY_PAGE), + num_values_(num_values), + encoding_(encoding), + is_sorted_(is_sorted) {} + + int32_t num_values() const { return num_values_; } + + Encoding::type encoding() const { return encoding_; } + + bool is_sorted() const { return is_sorted_; } + + private: + int32_t num_values_; + Encoding::type encoding_; + bool is_sorted_; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/column_scanner.h b/lib/python3.12/site-packages/pyarrow/include/parquet/column_scanner.h new file mode 100644 index 0000000000000000000000000000000000000000..a9953866fab22ee6db13a92578f85556ea6f99ba --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/column_scanner.h @@ -0,0 +1,264 @@ +// 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. + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +#include "parquet/column_reader.h" +#include "parquet/exception.h" +#include "parquet/platform.h" +#include "parquet/schema.h" +#include "parquet/types.h" + +namespace parquet { + +static constexpr int64_t DEFAULT_SCANNER_BATCH_SIZE = 128; + +class PARQUET_EXPORT Scanner { + public: + explicit Scanner(std::shared_ptr reader, + int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) + : batch_size_(batch_size), + level_offset_(0), + levels_buffered_(0), + value_buffer_(AllocateBuffer(pool)), + value_offset_(0), + values_buffered_(0), + reader_(std::move(reader)) { + def_levels_.resize( + descr()->max_definition_level() > 0 ? static_cast(batch_size_) : 0); + rep_levels_.resize( + descr()->max_repetition_level() > 0 ? static_cast(batch_size_) : 0); + } + + virtual ~Scanner() {} + + static std::shared_ptr Make( + std::shared_ptr col_reader, + int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + virtual void PrintNext(std::ostream& out, int width, bool with_levels = false) = 0; + + bool HasNext() { return level_offset_ < levels_buffered_ || reader_->HasNext(); } + + const ColumnDescriptor* descr() const { return reader_->descr(); } + + int64_t batch_size() const { return batch_size_; } + + void SetBatchSize(int64_t batch_size) { batch_size_ = batch_size; } + + protected: + int64_t batch_size_; + + std::vector def_levels_; + std::vector rep_levels_; + int level_offset_; + int levels_buffered_; + + std::shared_ptr value_buffer_; + int value_offset_; + int64_t values_buffered_; + std::shared_ptr reader_; +}; + +template +class PARQUET_TEMPLATE_CLASS_EXPORT TypedScanner : public Scanner { + public: + typedef typename DType::c_type T; + + explicit TypedScanner(std::shared_ptr reader, + int64_t batch_size = DEFAULT_SCANNER_BATCH_SIZE, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) + : Scanner(std::move(reader), batch_size, pool) { + typed_reader_ = static_cast*>(reader_.get()); + int value_byte_size = type_traits::value_byte_size; + PARQUET_THROW_NOT_OK(value_buffer_->Resize(batch_size_ * value_byte_size)); + values_ = reinterpret_cast(value_buffer_->mutable_data()); + } + + virtual ~TypedScanner() {} + + bool NextLevels(int16_t* def_level, int16_t* rep_level) { + if (level_offset_ == levels_buffered_) { + levels_buffered_ = static_cast( + typed_reader_->ReadBatch(static_cast(batch_size_), def_levels_.data(), + rep_levels_.data(), values_, &values_buffered_)); + + value_offset_ = 0; + level_offset_ = 0; + if (!levels_buffered_) { + return false; + } + } + *def_level = descr()->max_definition_level() > 0 ? def_levels_[level_offset_] : 0; + *rep_level = descr()->max_repetition_level() > 0 ? rep_levels_[level_offset_] : 0; + level_offset_++; + return true; + } + + bool Next(T* val, int16_t* def_level, int16_t* rep_level, bool* is_null) { + if (level_offset_ == levels_buffered_) { + if (!HasNext()) { + // Out of data pages + return false; + } + } + + NextLevels(def_level, rep_level); + *is_null = *def_level < descr()->max_definition_level(); + + if (*is_null) { + return true; + } + + if (value_offset_ == values_buffered_) { + throw ParquetException("Value was non-null, but has not been buffered"); + } + *val = values_[value_offset_++]; + return true; + } + + // Returns true if there is a next value + bool NextValue(T* val, bool* is_null) { + if (level_offset_ == levels_buffered_) { + if (!HasNext()) { + // Out of data pages + return false; + } + } + + // Out of values + int16_t def_level = -1; + int16_t rep_level = -1; + NextLevels(&def_level, &rep_level); + *is_null = def_level < descr()->max_definition_level(); + + if (*is_null) { + return true; + } + + if (value_offset_ == values_buffered_) { + throw ParquetException("Value was non-null, but has not been buffered"); + } + *val = values_[value_offset_++]; + return true; + } + + virtual void PrintNext(std::ostream& out, int width, bool with_levels = false) { + T val{}; + int16_t def_level = -1; + int16_t rep_level = -1; + bool is_null = false; + char buffer[80]; + + if (!Next(&val, &def_level, &rep_level, &is_null)) { + throw ParquetException("No more values buffered"); + } + + if (with_levels) { + out << " D:" << def_level << " R:" << rep_level << " "; + if (!is_null) { + out << "V:"; + } + } + + if (is_null) { + std::string null_fmt = format_fwf(width); + snprintf(buffer, sizeof(buffer), null_fmt.c_str(), "NULL"); + } else { + FormatValue(&val, buffer, sizeof(buffer), width); + } + out << buffer; + } + + private: + // The ownership of this object is expressed through the reader_ variable in the base + TypedColumnReader* typed_reader_; + + inline void FormatValue(void* val, char* buffer, int bufsize, int width); + + T* values_; +}; + +template +inline void TypedScanner::FormatValue(void* val, char* buffer, int bufsize, + int width) { + std::string fmt = format_fwf(width); + snprintf(buffer, bufsize, fmt.c_str(), *reinterpret_cast(val)); +} + +template <> +inline void TypedScanner::FormatValue(void* val, char* buffer, int bufsize, + int width) { + std::string fmt = format_fwf(width); + std::string result = Int96ToString(*reinterpret_cast(val)); + snprintf(buffer, bufsize, fmt.c_str(), result.c_str()); +} + +template <> +inline void TypedScanner::FormatValue(void* val, char* buffer, int bufsize, + int width) { + std::string fmt = format_fwf(width); + std::string result = ByteArrayToString(*reinterpret_cast(val)); + snprintf(buffer, bufsize, fmt.c_str(), result.c_str()); +} + +template <> +inline void TypedScanner::FormatValue(void* val, char* buffer, int bufsize, + int width) { + std::string fmt = format_fwf(width); + std::string result = FixedLenByteArrayToString( + *reinterpret_cast(val), descr()->type_length()); + snprintf(buffer, bufsize, fmt.c_str(), result.c_str()); +} + +typedef TypedScanner BoolScanner; +typedef TypedScanner Int32Scanner; +typedef TypedScanner Int64Scanner; +typedef TypedScanner Int96Scanner; +typedef TypedScanner FloatScanner; +typedef TypedScanner DoubleScanner; +typedef TypedScanner ByteArrayScanner; +typedef TypedScanner FixedLenByteArrayScanner; + +template +int64_t ScanAll(int32_t batch_size, int16_t* def_levels, int16_t* rep_levels, + uint8_t* values, int64_t* values_buffered, + parquet::ColumnReader* reader) { + typedef typename RType::T Type; + auto typed_reader = static_cast(reader); + auto vals = reinterpret_cast(&values[0]); + return typed_reader->ReadBatch(batch_size, def_levels, rep_levels, vals, + values_buffered); +} + +int64_t PARQUET_EXPORT ScanAllValues(int32_t batch_size, int16_t* def_levels, + int16_t* rep_levels, uint8_t* values, + int64_t* values_buffered, + parquet::ColumnReader* reader); + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/column_writer.h b/lib/python3.12/site-packages/pyarrow/include/parquet/column_writer.h new file mode 100644 index 0000000000000000000000000000000000000000..5b56eb010a242c4b32f51f39bc4d1f21dbe054d8 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/column_writer.h @@ -0,0 +1,307 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/type_fwd.h" +#include "arrow/util/compression.h" +#include "parquet/exception.h" +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace arrow { + +class Array; + +namespace bit_util { +class BitWriter; +} // namespace bit_util + +namespace util { +class RleBitPackedEncoder; +class CodecOptions; +} // namespace util + +} // namespace arrow + +namespace parquet { + +struct ArrowWriteContext; +class ColumnChunkMetaDataBuilder; +class ColumnDescriptor; +class ColumnIndexBuilder; +class DataPage; +class DictionaryPage; +class Encryptor; +class OffsetIndexBuilder; +class WriterProperties; + +class PARQUET_EXPORT LevelEncoder { + public: + LevelEncoder(); + ~LevelEncoder(); + + static int MaxBufferSize(Encoding::type encoding, int16_t max_level, + int num_buffered_values); + + // Initialize the LevelEncoder. + void Init(Encoding::type encoding, int16_t max_level, int num_buffered_values, + uint8_t* data, int data_size); + + // Encodes a batch of levels from an array and returns the number of levels encoded + int Encode(int batch_size, const int16_t* levels); + + int32_t len() { + if (encoding_ != Encoding::RLE) { + throw ParquetException("Only implemented for RLE encoding"); + } + return rle_length_; + } + + private: + int bit_width_; + int rle_length_; + Encoding::type encoding_; + std::unique_ptr<::arrow::util::RleBitPackedEncoder> rle_encoder_; + std::unique_ptr<::arrow::bit_util::BitWriter> bit_packed_encoder_; +}; + +class PARQUET_EXPORT PageWriter { + public: + virtual ~PageWriter() {} + + static std::unique_ptr Open( + std::shared_ptr sink, Compression::type codec, + ColumnChunkMetaDataBuilder* metadata, int16_t row_group_ordinal = -1, + int16_t column_chunk_ordinal = -1, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool(), + bool buffered_row_group = false, + std::shared_ptr header_encryptor = NULLPTR, + std::shared_ptr data_encryptor = NULLPTR, + bool page_write_checksum_enabled = false, + // column_index_builder MUST outlive the PageWriter + ColumnIndexBuilder* column_index_builder = NULLPTR, + // offset_index_builder MUST outlive the PageWriter + OffsetIndexBuilder* offset_index_builder = NULLPTR, + const CodecOptions& codec_options = CodecOptions{}); + + // The Column Writer decides if dictionary encoding is used if set and + // if the dictionary encoding has fallen back to default encoding on reaching dictionary + // page limit + virtual void Close(bool has_dictionary, bool fallback) = 0; + + // Return the number of uncompressed bytes written (including header size) + virtual int64_t WriteDataPage(const DataPage& page) = 0; + + // Return the number of uncompressed bytes written (including header size) + virtual int64_t WriteDictionaryPage(const DictionaryPage& page) = 0; + + /// \brief The total number of bytes written as serialized data and + /// dictionary pages to the sink so far. + virtual int64_t total_compressed_bytes_written() const = 0; + + virtual bool has_compressor() = 0; + + virtual void Compress(const Buffer& src_buffer, ResizableBuffer* dest_buffer) = 0; +}; + +class PARQUET_EXPORT ColumnWriter { + public: + virtual ~ColumnWriter() = default; + + static std::shared_ptr Make(ColumnChunkMetaDataBuilder*, + std::unique_ptr, + const WriterProperties* properties); + + /// \brief Closes the ColumnWriter, commits any buffered values to pages. + /// \return Total size of the column in bytes + virtual int64_t Close() = 0; + + /// \brief The physical Parquet type of the column + virtual Type::type type() const = 0; + + /// \brief The schema for the column + virtual const ColumnDescriptor* descr() const = 0; + + /// \brief The number of rows written so far + virtual int64_t rows_written() const = 0; + + /// \brief The total size of the compressed pages + page headers. Values + /// are still buffered and not written to a pager yet + /// + /// So in un-buffered mode, it always returns 0 + virtual int64_t total_compressed_bytes() const = 0; + + /// \brief The total number of bytes written as serialized data and + /// dictionary pages to the ColumnChunk so far + /// These bytes are uncompressed bytes. + virtual int64_t total_bytes_written() const = 0; + + /// \brief The total number of bytes written as serialized data and + /// dictionary pages to the ColumnChunk so far. + /// If the column is uncompressed, the value would be equal to + /// total_bytes_written(). + virtual int64_t total_compressed_bytes_written() const = 0; + + /// \brief Estimated size of the values that are not written to a page yet. + virtual int64_t estimated_buffered_value_bytes() const = 0; + + /// \brief The file-level writer properties + virtual const WriterProperties* properties() = 0; + + /// \brief Add key-value metadata to the ColumnChunk. + /// \param[in] key_value_metadata the metadata to add. + /// \note This will overwrite any existing metadata with the same key. + /// \throw ParquetException if Close() has been called. + virtual void AddKeyValueMetadata( + const std::shared_ptr& key_value_metadata) = 0; + + /// \brief Reset the ColumnChunk key-value metadata. + /// \throw ParquetException if Close() has been called. + virtual void ResetKeyValueMetadata() = 0; + + /// \brief Write Apache Arrow columnar data directly to ColumnWriter. Returns + /// error status if the array data type is not compatible with the concrete + /// writer type. + /// + /// leaf_array is always a primitive (possibly dictionary encoded type). + /// Leaf_field_nullable indicates whether the leaf array is considered nullable + /// according to its schema in a Table or its parent array. + virtual ::arrow::Status WriteArrow(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_levels, const ::arrow::Array& leaf_array, + ArrowWriteContext* ctx, + bool leaf_field_nullable) = 0; +}; + +// API to write values to a single column. This is the main client facing API. +template +class TypedColumnWriter : public ColumnWriter { + public: + using T = typename DType::c_type; + + // Write a batch of repetition levels, definition levels, and values to the + // column. + // `num_values` is the number of logical leaf values. + // `def_levels` (resp. `rep_levels`) can be null if the column's max definition level + // (resp. max repetition level) is 0. + // If not null, each of `def_levels` and `rep_levels` must have at least + // `num_values`. + // + // The number of physical values written (taken from `values`) is returned. + // It can be smaller than `num_values` is there are some undefined values. + virtual int64_t WriteBatch(int64_t num_values, const int16_t* def_levels, + const int16_t* rep_levels, const T* values) = 0; + + /// Write a batch of repetition levels, definition levels, and values to the + /// column. + /// + /// In comparison to WriteBatch the length of repetition and definition levels + /// is the same as of the number of values read for max_definition_level == 1. + /// In the case of max_definition_level > 1, the repetition and definition + /// levels are larger than the values but the values include the null entries + /// with definition_level == (max_definition_level - 1). Thus we have to differentiate + /// in the parameters of this function if the input has the length of num_values or the + /// _number of rows in the lowest nesting level_. + /// + /// In the case that the most inner node in the Parquet is required, the _number of rows + /// in the lowest nesting level_ is equal to the number of non-null values. If the + /// inner-most schema node is optional, the _number of rows in the lowest nesting level_ + /// also includes all values with definition_level == (max_definition_level - 1). + /// + /// @param num_values number of levels to write. + /// @param def_levels The Parquet definition levels, length is num_values + /// @param rep_levels The Parquet repetition levels, length is num_values + /// @param valid_bits Bitmap that indicates if the row is null on the lowest nesting + /// level. The length is number of rows in the lowest nesting level. + /// @param valid_bits_offset The offset in bits of the valid_bits where the + /// first relevant bit resides. + /// @param values The values in the lowest nested level including + /// spacing for nulls on the lowest levels; input has the length + /// of the number of rows on the lowest nesting level. + virtual void WriteBatchSpaced(int64_t num_values, const int16_t* def_levels, + const int16_t* rep_levels, const uint8_t* valid_bits, + int64_t valid_bits_offset, const T* values) = 0; +}; + +using BoolWriter = TypedColumnWriter; +using Int32Writer = TypedColumnWriter; +using Int64Writer = TypedColumnWriter; +using Int96Writer = TypedColumnWriter; +using FloatWriter = TypedColumnWriter; +using DoubleWriter = TypedColumnWriter; +using ByteArrayWriter = TypedColumnWriter; +using FixedLenByteArrayWriter = TypedColumnWriter; + +namespace internal { + +/** + * Timestamp conversion constants + */ +constexpr int64_t kJulianEpochOffsetDays = INT64_C(2440588); + +template +inline void ArrowTimestampToImpalaTimestamp(const int64_t time, Int96* impala_timestamp) { + auto julian_days = static_cast(time / UnitPerDay + kJulianEpochOffsetDays); + int64_t last_day_units = time % UnitPerDay; + if (last_day_units < 0) { + --julian_days; + last_day_units += UnitPerDay; + } + impala_timestamp->value[2] = static_cast(julian_days); + uint64_t last_day_nanos = static_cast(last_day_units) * NanosecondsPerUnit; + // impala_timestamp will be unaligned every other entry so do memcpy instead + // of assign and reinterpret cast to avoid undefined behavior. + std::memcpy(impala_timestamp, &last_day_nanos, sizeof(uint64_t)); +} + +constexpr int64_t kSecondsInNanos = INT64_C(1000000000); + +inline void SecondsToImpalaTimestamp(const int64_t seconds, Int96* impala_timestamp) { + ArrowTimestampToImpalaTimestamp(seconds, + impala_timestamp); +} + +constexpr int64_t kMillisecondsInNanos = kSecondsInNanos / INT64_C(1000); + +inline void MillisecondsToImpalaTimestamp(const int64_t milliseconds, + Int96* impala_timestamp) { + ArrowTimestampToImpalaTimestamp( + milliseconds, impala_timestamp); +} + +constexpr int64_t kMicrosecondsInNanos = kMillisecondsInNanos / INT64_C(1000); + +inline void MicrosecondsToImpalaTimestamp(const int64_t microseconds, + Int96* impala_timestamp) { + ArrowTimestampToImpalaTimestamp( + microseconds, impala_timestamp); +} + +constexpr int64_t kNanosecondsInNanos = INT64_C(1); + +inline void NanosecondsToImpalaTimestamp(const int64_t nanoseconds, + Int96* impala_timestamp) { + ArrowTimestampToImpalaTimestamp( + nanoseconds, impala_timestamp); +} + +} // namespace internal +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encoding.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encoding.h new file mode 100644 index 0000000000000000000000000000000000000000..d80bf0edcae41cf213e497abdf07ad2d6e980ec0 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encoding.h @@ -0,0 +1,458 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "arrow/type_fwd.h" + +#include "parquet/exception.h" +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace arrow { +template +class Dictionary32Builder; +} + +namespace parquet { + +template +class TypedEncoder; + +using BooleanEncoder = TypedEncoder; +using Int32Encoder = TypedEncoder; +using Int64Encoder = TypedEncoder; +using Int96Encoder = TypedEncoder; +using FloatEncoder = TypedEncoder; +using DoubleEncoder = TypedEncoder; +using ByteArrayEncoder = TypedEncoder; +using FLBAEncoder = TypedEncoder; + +template +class TypedDecoder; + +class BooleanDecoder; +using Int32Decoder = TypedDecoder; +using Int64Decoder = TypedDecoder; +using Int96Decoder = TypedDecoder; +using FloatDecoder = TypedDecoder; +using DoubleDecoder = TypedDecoder; +using ByteArrayDecoder = TypedDecoder; +class FLBADecoder; + +template +struct EncodingTraits; + +template <> +struct EncodingTraits { + using Encoder = BooleanEncoder; + using Decoder = BooleanDecoder; + + using ArrowType = ::arrow::BooleanType; + using Accumulator = ::arrow::BooleanBuilder; + struct DictAccumulator {}; +}; + +template <> +struct EncodingTraits { + using Encoder = Int32Encoder; + using Decoder = Int32Decoder; + + using ArrowType = ::arrow::Int32Type; + using Accumulator = ::arrow::NumericBuilder<::arrow::Int32Type>; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::Int32Type>; +}; + +template <> +struct EncodingTraits { + using Encoder = Int64Encoder; + using Decoder = Int64Decoder; + + using ArrowType = ::arrow::Int64Type; + using Accumulator = ::arrow::NumericBuilder<::arrow::Int64Type>; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::Int64Type>; +}; + +template <> +struct EncodingTraits { + using Encoder = Int96Encoder; + using Decoder = Int96Decoder; + + struct Accumulator {}; + struct DictAccumulator {}; +}; + +template <> +struct EncodingTraits { + using Encoder = FloatEncoder; + using Decoder = FloatDecoder; + + using ArrowType = ::arrow::FloatType; + using Accumulator = ::arrow::NumericBuilder<::arrow::FloatType>; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::FloatType>; +}; + +template <> +struct EncodingTraits { + using Encoder = DoubleEncoder; + using Decoder = DoubleDecoder; + + using ArrowType = ::arrow::DoubleType; + using Accumulator = ::arrow::NumericBuilder<::arrow::DoubleType>; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::DoubleType>; +}; + +template <> +struct EncodingTraits { + using Encoder = ByteArrayEncoder; + using Decoder = ByteArrayDecoder; + + /// \brief Internal helper class for decoding BYTE_ARRAY data + /// + /// This class allows the caller to choose the concrete Arrow data type + /// by passing a corresponding `ArrayBuilder`. + /// Supported `ArrayBuilder` classes are `BinaryBuilder`, `LargeBinaryBuilder` + /// and `BinaryViewBuilder`. + /// If the builder is a `BinaryBuilder`, `chunks` can accumulate several + /// arrays as needed to work around the 32-bit offset limit. + struct Accumulator { + std::unique_ptr<::arrow::ArrayBuilder> builder; + std::vector> chunks; + }; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::BinaryType>; +}; + +template <> +struct EncodingTraits { + using Encoder = FLBAEncoder; + using Decoder = FLBADecoder; + + using ArrowType = ::arrow::FixedSizeBinaryType; + using Accumulator = ::arrow::FixedSizeBinaryBuilder; + using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::FixedSizeBinaryType>; +}; + +class ColumnDescriptor; + +// Untyped base for all encoders +class Encoder { + public: + virtual ~Encoder() = default; + + virtual int64_t EstimatedDataEncodedSize() = 0; + virtual std::shared_ptr FlushValues() = 0; + virtual Encoding::type encoding() const = 0; + + virtual void Put(const ::arrow::Array& values) = 0; + + // Report the number of bytes written to the encoder since the last report. + // It only works for BYTE_ARRAY type and throw for other types. + // This call is not idempotent since it resets the internal counter. + virtual int64_t ReportUnencodedDataBytes() = 0; + + virtual MemoryPool* memory_pool() const = 0; +}; + +// Base class for value encoders. Since encoders may or not have state (e.g., +// dictionary encoding) we use a class instance to maintain any state. +// +// Encode interfaces are internal, subject to change without deprecation. +template +class TypedEncoder : virtual public Encoder { + public: + using T = typename DType::c_type; + + using Encoder::Put; + + virtual void Put(const T* src, int num_values) = 0; + + virtual void Put(const std::vector& src, int num_values = -1); + + virtual void PutSpaced(const T* src, int num_values, const uint8_t* valid_bits, + int64_t valid_bits_offset) = 0; +}; + +template +void TypedEncoder::Put(const std::vector& src, int num_values) { + if (num_values == -1) { + num_values = static_cast(src.size()); + } + Put(src.data(), num_values); +} + +template <> +inline void TypedEncoder::Put(const std::vector& src, int num_values) { + // NOTE(wesm): This stub is here only to satisfy the compiler; it is + // overridden later with the actual implementation +} + +// Base class for dictionary encoders +template +class DictEncoder : virtual public TypedEncoder { + public: + /// Writes out any buffered indices to buffer preceded by the bit width of this data. + /// Returns the number of bytes written. + /// If the supplied buffer is not big enough, returns -1. + /// buffer must be preallocated with buffer_len bytes. Use EstimatedDataEncodedSize() + /// to size buffer. + virtual int WriteIndices(uint8_t* buffer, int buffer_len) = 0; + + virtual int dict_encoded_size() const = 0; + + virtual int bit_width() const = 0; + + /// Writes out the encoded dictionary to buffer. buffer must be preallocated to + /// dict_encoded_size() bytes. + virtual void WriteDict(uint8_t* buffer) const = 0; + + virtual int num_entries() const = 0; + + /// \brief EXPERIMENTAL: Append dictionary indices into the encoder. It is + /// assumed (without any boundschecking) that the indices reference + /// preexisting dictionary values + /// \param[in] indices the dictionary index values. Only Int32Array currently + /// supported + virtual void PutIndices(const ::arrow::Array& indices) = 0; + + /// \brief EXPERIMENTAL: Append dictionary into encoder, inserting indices + /// separately. Currently throws exception if the current dictionary memo is + /// non-empty + /// \param[in] values the dictionary values. Only valid for certain + /// Parquet/Arrow type combinations, like BYTE_ARRAY/BinaryArray + virtual void PutDictionary(const ::arrow::Array& values) = 0; +}; + +// ---------------------------------------------------------------------- +// Value decoding + +class Decoder { + public: + virtual ~Decoder() = default; + + // Sets the data for a new page. This will be called multiple times on the same + // decoder and should reset all internal state. + // + // `num_values` comes from the data page header, and may be greater than the number of + // physical values in the data buffer if there are some omitted (null) values. + // `len`, on the other hand, is the size in bytes of the data buffer and + // directly relates to the number of physical values. + virtual void SetData(int num_values, const uint8_t* data, int len) = 0; + + // Returns the number of values left (for the last call to SetData()). This is + // the number of values left in this page. + virtual int values_left() const = 0; + virtual Encoding::type encoding() const = 0; +}; + +template +class TypedDecoder : virtual public Decoder { + public: + using T = typename DType::c_type; + + /// \brief Decode values into a buffer + /// + /// Subclasses may override the more specialized Decode methods below. + /// + /// \param[in] buffer destination for decoded values + /// \param[in] max_values maximum number of values to decode + /// \return The number of values decoded. Should be identical to max_values except + /// at the end of the current data page. + virtual int Decode(T* buffer, int max_values) = 0; + + /// \brief Decode the values in this data page but leave spaces for null entries. + /// + /// \param[in] buffer destination for decoded values + /// \param[in] num_values size of the def_levels and buffer arrays including the number + /// of null slots + /// \param[in] null_count number of null slots + /// \param[in] valid_bits bitmap data indicating position of valid slots + /// \param[in] valid_bits_offset offset into valid_bits + /// \return The number of values decoded, including nulls. + virtual int DecodeSpaced(T* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) = 0; + + /// \brief Decode into an ArrayBuilder or other accumulator + /// + /// This function assumes the definition levels were already decoded + /// as a validity bitmap in the given `valid_bits`. `null_count` + /// is the number of 0s in `valid_bits`. + /// As a space optimization, it is allowed for `valid_bits` to be null + /// if `null_count` is zero. + /// + /// \return number of values decoded + virtual int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset, + typename EncodingTraits::Accumulator* out) = 0; + + /// \brief Decode into an ArrayBuilder or other accumulator ignoring nulls + /// + /// \return number of values decoded + int DecodeArrowNonNull(int num_values, + typename EncodingTraits::Accumulator* out) { + return DecodeArrow(num_values, 0, /*valid_bits=*/NULLPTR, 0, out); + } + + /// \brief Decode into a DictionaryBuilder + /// + /// This function assumes the definition levels were already decoded + /// as a validity bitmap in the given `valid_bits`. `null_count` + /// is the number of 0s in `valid_bits`. + /// As a space optimization, it is allowed for `valid_bits` to be null + /// if `null_count` is zero. + /// + /// \return number of values decoded + virtual int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset, + typename EncodingTraits::DictAccumulator* builder) = 0; + + /// \brief Decode into a DictionaryBuilder ignoring nulls + /// + /// \return number of values decoded + int DecodeArrowNonNull(int num_values, + typename EncodingTraits::DictAccumulator* builder) { + return DecodeArrow(num_values, 0, /*valid_bits=*/NULLPTR, 0, builder); + } +}; + +template +class DictDecoder : virtual public TypedDecoder { + public: + using T = typename DType::c_type; + + virtual void SetDict(TypedDecoder* dictionary) = 0; + + /// \brief Insert dictionary values into the Arrow dictionary builder's memo, + /// but do not append any indices + virtual void InsertDictionary(::arrow::ArrayBuilder* builder) = 0; + + /// \brief Decode only dictionary indices and append to dictionary + /// builder. The builder must have had the dictionary from this decoder + /// inserted already. + /// + /// \warning Remember to reset the builder each time the dict decoder is initialized + /// with a new dictionary page + virtual int DecodeIndicesSpaced(int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + ::arrow::ArrayBuilder* builder) = 0; + + /// \brief Decode only dictionary indices (no nulls) + /// + /// \warning Remember to reset the builder each time the dict decoder is initialized + /// with a new dictionary page + virtual int DecodeIndices(int num_values, ::arrow::ArrayBuilder* builder) = 0; + + /// \brief Decode only dictionary indices (no nulls). Same as above + /// DecodeIndices but target is an array instead of a builder. + /// + /// \note API EXPERIMENTAL + virtual int DecodeIndices(int num_values, int32_t* indices) = 0; + + /// \brief Get dictionary. The reader will call this API when it encounters a + /// new dictionary. + /// + /// @param[out] dictionary The pointer to dictionary values. Dictionary is owned by + /// the decoder and is destroyed when the decoder is destroyed. + /// @param[out] dictionary_length The dictionary length. + /// + /// \note API EXPERIMENTAL + virtual void GetDictionary(const T** dictionary, int32_t* dictionary_length) = 0; +}; + +// ---------------------------------------------------------------------- +// TypedEncoder specializations, traits, and factory functions + +class BooleanDecoder : virtual public TypedDecoder { + public: + using TypedDecoder::Decode; + + /// \brief Decode and bit-pack values into a buffer + /// + /// \param[in] buffer destination for decoded values + /// This buffer will contain bit-packed values. If + /// max_values is not a multiple of 8, the trailing bits + /// of the last byte will be undefined. + /// \param[in] max_values max values to decode. + /// \return The number of values decoded. Should be identical to max_values except + /// at the end of the current data page. + virtual int Decode(uint8_t* buffer, int max_values) = 0; +}; + +class FLBADecoder : virtual public TypedDecoder { + public: + using TypedDecoder::DecodeSpaced; + + // TODO(wesm): As possible follow-up to PARQUET-1508, we should examine if + // there is value in adding specialized read methods for + // FIXED_LEN_BYTE_ARRAY. If only Decimal data can occur with this data type + // then perhaps not +}; + +PARQUET_EXPORT +std::unique_ptr MakeEncoder( + Type::type type_num, Encoding::type encoding, bool use_dictionary = false, + const ColumnDescriptor* descr = NULLPTR, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +template +std::unique_ptr::Encoder> MakeTypedEncoder( + Encoding::type encoding, bool use_dictionary = false, + const ColumnDescriptor* descr = NULLPTR, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + using OutType = typename EncodingTraits::Encoder; + std::unique_ptr base = + MakeEncoder(DType::type_num, encoding, use_dictionary, descr, pool); + return std::unique_ptr(dynamic_cast(base.release())); +} + +PARQUET_EXPORT +std::unique_ptr MakeDecoder( + Type::type type_num, Encoding::type encoding, const ColumnDescriptor* descr = NULLPTR, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +namespace detail { + +PARQUET_EXPORT +std::unique_ptr MakeDictDecoder(Type::type type_num, + const ColumnDescriptor* descr, + ::arrow::MemoryPool* pool); + +} // namespace detail + +template +std::unique_ptr> MakeDictDecoder( + const ColumnDescriptor* descr = NULLPTR, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + using OutType = DictDecoder; + auto decoder = detail::MakeDictDecoder(DType::type_num, descr, pool); + return std::unique_ptr(dynamic_cast(decoder.release())); +} + +template +std::unique_ptr::Decoder> MakeTypedDecoder( + Encoding::type encoding, const ColumnDescriptor* descr = NULLPTR, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + using OutType = typename EncodingTraits::Decoder; + std::unique_ptr base = MakeDecoder(DType::type_num, encoding, descr, pool); + return std::unique_ptr(dynamic_cast(base.release())); +} + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/crypto_factory.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/crypto_factory.h new file mode 100644 index 0000000000000000000000000000000000000000..3c6a5f29ea88eb6849abb78ce334406ba293428f --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/crypto_factory.h @@ -0,0 +1,152 @@ +// 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. + +#pragma once + +#include + +#include "parquet/encryption/encryption.h" +#include "parquet/encryption/file_key_wrapper.h" +#include "parquet/encryption/key_toolkit.h" +#include "parquet/encryption/kms_client_factory.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +static constexpr ParquetCipher::type kDefaultEncryptionAlgorithm = + ParquetCipher::AES_GCM_V1; +static constexpr bool kDefaultPlaintextFooter = false; +static constexpr bool kDefaultDoubleWrapping = true; +static constexpr double kDefaultCacheLifetimeSeconds = 600; // 10 minutes +static constexpr bool kDefaultInternalKeyMaterial = true; +static constexpr bool kDefaultUniformEncryption = false; +static constexpr int32_t kDefaultDataKeyLengthBits = 128; + +struct PARQUET_EXPORT EncryptionConfiguration { + explicit EncryptionConfiguration(const std::string& footer_key) + : footer_key(footer_key) {} + + /// ID of the master key for footer encryption/signing + std::string footer_key; + + /// List of columns to encrypt, with column master key IDs (see HIVE-21848). + /// Format: "columnKeyID:colName,colName;columnKeyID:colName..." + /// Either + /// (1) column_keys must be set + /// or + /// (2) uniform_encryption must be set to true + /// If none of (1) and (2) are true, or if both are true, an exception will be + /// thrown. + std::string column_keys; + + /// Encrypt footer and all columns with the same encryption key. + bool uniform_encryption = kDefaultUniformEncryption; + + /// Parquet encryption algorithm. Can be "AES_GCM_V1" (default), or "AES_GCM_CTR_V1". + ParquetCipher::type encryption_algorithm = kDefaultEncryptionAlgorithm; + + /// Write files with plaintext footer. + /// The default is false - files are written with encrypted footer. + bool plaintext_footer = kDefaultPlaintextFooter; + + /// Use double wrapping - where data encryption keys (DEKs) are encrypted with key + /// encryption keys (KEKs), which in turn are encrypted with master keys. + /// The default is true. If set to false, use single wrapping - where DEKs are + /// encrypted directly with master keys. + bool double_wrapping = kDefaultDoubleWrapping; + + /// Lifetime of cached entities (key encryption keys, local wrapping keys, KMS client + /// objects). + /// The default is 600 (10 minutes). + double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds; + + /// Store key material inside Parquet file footers; this mode doesn’t produce + /// additional files. By default, true. If set to false, key material is stored in + /// separate files in the same folder, which enables key rotation for immutable + /// Parquet files. + bool internal_key_material = kDefaultInternalKeyMaterial; + + /// Length of data encryption keys (DEKs), randomly generated by parquet key + /// management tools. Can be 128, 192 or 256 bits. + /// The default is 128 bits. + int32_t data_key_length_bits = kDefaultDataKeyLengthBits; +}; + +struct PARQUET_EXPORT DecryptionConfiguration { + /// Lifetime of cached entities (key encryption keys, local wrapping keys, KMS client + /// objects). + /// The default is 600 (10 minutes). + double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds; +}; + +/// This is a core class, that translates the parameters of high level encryption (like +/// the names of encrypted columns, names of master keys, etc), into parameters of low +/// level encryption (like the key metadata, DEK, etc). A factory that produces the low +/// level FileEncryptionProperties and FileDecryptionProperties objects, from the high +/// level parameters. +class PARQUET_EXPORT CryptoFactory { + public: + /// a KmsClientFactory object must be registered via this method before calling any of + /// GetFileEncryptionProperties()/GetFileDecryptionProperties() methods. + void RegisterKmsClientFactory(std::shared_ptr kms_client_factory); + + /// Get the encryption properties for a Parquet file. + /// If external key material is used then a file system and path to the + /// parquet file must be provided. + std::shared_ptr GetFileEncryptionProperties( + const KmsConnectionConfig& kms_connection_config, + const EncryptionConfiguration& encryption_config, const std::string& file_path = "", + const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR); + + /// Get decryption properties for a Parquet file. + /// If external key material is used then a file system and path to the + /// parquet file must be provided. + std::shared_ptr GetFileDecryptionProperties( + const KmsConnectionConfig& kms_connection_config, + const DecryptionConfiguration& decryption_config, const std::string& file_path = "", + const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR); + + void RemoveCacheEntriesForToken(const std::string& access_token) { + key_toolkit_->RemoveCacheEntriesForToken(access_token); + } + + void RemoveCacheEntriesForAllTokens() { + key_toolkit_->RemoveCacheEntriesForAllTokens(); + } + + /// Rotates master encryption keys for a Parquet file that uses external key material. + /// In single wrapping mode, data encryption keys are decrypted with the old master keys + /// and then re-encrypted with new master keys. + /// In double wrapping mode, key encryption keys are decrypted with the old master keys + /// and then re-encrypted with new master keys. + /// This relies on the KMS supporting versioning, such that the old master key is + /// used when unwrapping a key, and the latest version is used when wrapping a key. + void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config, + const std::string& parquet_file_path, + const std::shared_ptr<::arrow::fs::FileSystem>& file_system, + bool double_wrapping = kDefaultDoubleWrapping, + double cache_lifetime_seconds = kDefaultCacheLifetimeSeconds); + + private: + ColumnPathToEncryptionPropertiesMap GetColumnEncryptionProperties( + int dek_length, const std::string& column_keys, FileKeyWrapper* key_wrapper); + + /// Key utilities object for kms client initialization and cache control + std::shared_ptr key_toolkit_ = std::make_shared(); +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/encryption.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/encryption.h new file mode 100644 index 0000000000000000000000000000000000000000..b4634b704735a8a34ca8a9073c639fdac2da3521 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/encryption.h @@ -0,0 +1,433 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/util/secure_string.h" +#include "parquet/exception.h" +#include "parquet/schema.h" +#include "parquet/types.h" + +namespace parquet { + +static constexpr ParquetCipher::type kDefaultEncryptionAlgorithm = + ParquetCipher::AES_GCM_V1; +static constexpr int32_t kMaximalAadMetadataLength = 256; +static constexpr bool kDefaultEncryptedFooter = true; +static constexpr bool kDefaultCheckSignature = true; +static constexpr bool kDefaultAllowPlaintextFiles = false; +static constexpr int32_t kAadFileUniqueLength = 8; + +class ColumnDecryptionProperties; +using ColumnPathToDecryptionPropertiesMap = + std::map>; + +class ColumnEncryptionProperties; +using ColumnPathToEncryptionPropertiesMap = + std::map>; + +class PARQUET_EXPORT DecryptionKeyRetriever { + public: + /// \brief Retrieve a key. + virtual ::arrow::util::SecureString GetKey(const std::string& key_id) = 0; + + virtual ~DecryptionKeyRetriever() {} +}; + +/// Simple integer key retriever +class PARQUET_EXPORT IntegerKeyIdRetriever : public DecryptionKeyRetriever { + public: + void PutKey(uint32_t key_id, ::arrow::util::SecureString key); + + ::arrow::util::SecureString GetKey(const std::string& key_id_string) override { + // key_id_string is string but for IntegerKeyIdRetriever it encodes + // a native-endian 32 bit unsigned integer key_id + uint32_t key_id; + assert(key_id_string.size() == sizeof(key_id)); + memcpy(&key_id, key_id_string.data(), sizeof(key_id)); + + return key_map_.at(key_id); + } + + private: + std::map key_map_; +}; + +// Simple string key retriever +class PARQUET_EXPORT StringKeyIdRetriever : public DecryptionKeyRetriever { + public: + void PutKey(std::string key_id, ::arrow::util::SecureString key); + ::arrow::util::SecureString GetKey(const std::string& key_id) override; + + private: + std::map key_map_; +}; + +class PARQUET_EXPORT HiddenColumnException : public ParquetException { + public: + explicit HiddenColumnException(const std::string& columnPath) + : ParquetException(columnPath.c_str()) {} +}; + +class PARQUET_EXPORT KeyAccessDeniedException : public ParquetException { + public: + explicit KeyAccessDeniedException(const std::string& columnPath) + : ParquetException(columnPath.c_str()) {} +}; + +inline ::arrow::util::span str2span(const std::string& str) { + if (str.empty()) { + return {}; + } + + return {reinterpret_cast(str.data()), str.size()}; +} + +class PARQUET_EXPORT ColumnEncryptionProperties { + public: + class PARQUET_EXPORT Builder { + public: + PARQUET_DEPRECATED("name argument is ignored, use default constructor instead") + explicit Builder(const std::string& name) : encrypted_(true) {} + + PARQUET_DEPRECATED("path argument is ignored, use default constructor instead") + explicit Builder(const schema::ColumnPath& path) : encrypted_(true) {} + + Builder() = default; + + /// Set a column-specific key. + /// If key is not set on an encrypted column, the column will + /// be encrypted with the footer key. + /// keyBytes Key length must be either 16, 24 or 32 bytes. + /// Caller is responsible for wiping out the input key array. + Builder* key(::arrow::util::SecureString column_key); + + /// Set a key retrieval metadata. + /// use either key_metadata() or key_id(), not both + Builder* key_metadata(std::string key_metadata); + + /// A convenience function to set key metadata using a string id. + /// Set a key retrieval metadata (converted from String). + /// use either key_metadata() or key_id(), not both + /// key_id will be converted to metadata (UTF-8 array). + Builder* key_id(std::string key_id); + + std::shared_ptr build() { + return std::shared_ptr( + new ColumnEncryptionProperties(encrypted_, key_, key_metadata_)); + } + + private: + bool encrypted_ = true; + ::arrow::util::SecureString key_; + std::string key_metadata_; + }; + + bool is_encrypted() const { return encrypted_; } + bool is_encrypted_with_footer_key() const { return encrypted_with_footer_key_; } + const ::arrow::util::SecureString& key() const { return key_; } + const std::string& key_metadata() const { return key_metadata_; } + + static std::shared_ptr Unencrypted(); + static std::shared_ptr WithFooterKey(); + static std::shared_ptr WithColumnKey( + ::arrow::util::SecureString key, std::string key_metadata = ""); + + private: + bool encrypted_; + bool encrypted_with_footer_key_; + ::arrow::util::SecureString key_; + std::string key_metadata_; + explicit ColumnEncryptionProperties(bool encrypted, ::arrow::util::SecureString key, + std::string key_metadata); +}; + +class PARQUET_EXPORT ColumnDecryptionProperties { + public: + class PARQUET_EXPORT Builder { + public: + explicit Builder(std::string name) : column_path_(std::move(name)) {} + + explicit Builder(const schema::ColumnPath& path) : Builder(path.ToDotString()) {} + + /// Set an explicit column key. If applied on a file that contains + /// key metadata for this column the metadata will be ignored, + /// the column will be decrypted with this key. + /// key length must be either 16, 24 or 32 bytes. + Builder* key(::arrow::util::SecureString key); + + std::shared_ptr build(); + + private: + std::string column_path_; + ::arrow::util::SecureString key_; + }; + + const std::string& column_path() const { return column_path_; } + const ::arrow::util::SecureString& key() const { return key_; } + + private: + std::string column_path_; + ::arrow::util::SecureString key_; + + /// This class is only required for setting explicit column decryption keys - + /// to override key retriever (or to provide keys when key metadata and/or + /// key retriever are not available) + explicit ColumnDecryptionProperties(std::string column_path, + ::arrow::util::SecureString key); +}; + +class PARQUET_EXPORT AADPrefixVerifier { + public: + /// Verifies identity (AAD Prefix) of individual file, + /// or of file collection in a data set. + /// Throws exception if an AAD prefix is wrong. + /// In a data set, AAD Prefixes should be collected, + /// and then checked for missing files. + virtual void Verify(const std::string& aad_prefix) = 0; + virtual ~AADPrefixVerifier() {} +}; + +class PARQUET_EXPORT FileDecryptionProperties { + public: + class PARQUET_EXPORT Builder { + public: + Builder() { + check_plaintext_footer_integrity_ = kDefaultCheckSignature; + plaintext_files_allowed_ = kDefaultAllowPlaintextFiles; + } + + /// Set an explicit footer key. If applied on a file that contains + /// footer key metadata the metadata will be ignored, the footer + /// will be decrypted/verified with this key. + /// If explicit key is not set, footer key will be fetched from + /// key retriever. + /// With explicit keys or AAD prefix, new encryption properties object must be + /// created for each encrypted file. + /// Explicit encryption keys (footer and column) are cloned. + /// Upon completion of file reading, the cloned encryption keys in the properties + /// will be wiped out (array values set to 0). + /// Caller is responsible for wiping out the input key array. + /// param footerKey Key length must be either 16, 24 or 32 bytes. + Builder* footer_key(::arrow::util::SecureString footer_key); + + /// Set explicit column keys (decryption properties). + /// Its also possible to set a key retriever on this property object. + /// Upon file decryption, availability of explicit keys is checked before + /// invocation of the retriever callback. + /// If an explicit key is available for a footer or a column, + /// its key metadata will be ignored. + Builder* column_keys( + ColumnPathToDecryptionPropertiesMap column_decryption_properties); + + /// Set a key retriever callback. Its also possible to + /// set explicit footer or column keys on this file property object. + /// Upon file decryption, availability of explicit keys is checked before + /// invocation of the retriever callback. + /// If an explicit key is available for a footer or a column, + /// its key metadata will be ignored. + Builder* key_retriever(std::shared_ptr key_retriever); + + /// Skip integrity verification of plaintext footers. + /// If not called, integrity of plaintext footers will be checked in runtime, + /// and an exception will be thrown in the following situations: + /// - footer signing key is not available + /// (not passed, or not found by key retriever) + /// - footer content and signature don't match + Builder* disable_footer_signature_verification() { + check_plaintext_footer_integrity_ = false; + return this; + } + + /// Explicitly supply the file AAD prefix. + /// A must when a prefix is used for file encryption, but not stored in file. + /// If AAD prefix is stored in file, it will be compared to the explicitly + /// supplied value and an exception will be thrown if they differ. + Builder* aad_prefix(std::string aad_prefix); + + /// Set callback for verification of AAD Prefixes stored in file. + Builder* aad_prefix_verifier(std::shared_ptr aad_prefix_verifier); + + /// By default, reading plaintext (unencrypted) files is not + /// allowed when using a decryptor + /// - in order to detect files that were not encrypted by mistake. + /// However, the default behavior can be overridden by calling this method. + /// The caller should use then a different method to ensure encryption + /// of files with sensitive data. + Builder* plaintext_files_allowed() { + plaintext_files_allowed_ = true; + return this; + } + + std::shared_ptr build() { + return std::shared_ptr(new FileDecryptionProperties( + footer_key_, key_retriever_, check_plaintext_footer_integrity_, aad_prefix_, + aad_prefix_verifier_, column_decryption_properties_, plaintext_files_allowed_)); + } + + private: + ::arrow::util::SecureString footer_key_; + std::string aad_prefix_; + std::shared_ptr aad_prefix_verifier_; + ColumnPathToDecryptionPropertiesMap column_decryption_properties_; + + std::shared_ptr key_retriever_; + bool check_plaintext_footer_integrity_; + bool plaintext_files_allowed_; + }; + + const ::arrow::util::SecureString& column_key(const std::string& column_path) const; + + const ::arrow::util::SecureString& footer_key() const { return footer_key_; } + + const std::string& aad_prefix() const { return aad_prefix_; } + + const std::shared_ptr& key_retriever() const { + return key_retriever_; + } + + bool check_plaintext_footer_integrity() const { + return check_plaintext_footer_integrity_; + } + + bool plaintext_files_allowed() const { return plaintext_files_allowed_; } + + const std::shared_ptr& aad_prefix_verifier() const { + return aad_prefix_verifier_; + } + + private: + ::arrow::util::SecureString footer_key_; + std::string aad_prefix_; + std::shared_ptr aad_prefix_verifier_; + ColumnPathToDecryptionPropertiesMap column_decryption_properties_; + std::shared_ptr key_retriever_; + bool check_plaintext_footer_integrity_; + bool plaintext_files_allowed_; + + FileDecryptionProperties( + ::arrow::util::SecureString footer_key, + std::shared_ptr key_retriever, + bool check_plaintext_footer_integrity, std::string aad_prefix, + std::shared_ptr aad_prefix_verifier, + ColumnPathToDecryptionPropertiesMap column_decryption_properties, + bool plaintext_files_allowed); +}; + +class PARQUET_EXPORT FileEncryptionProperties { + public: + class PARQUET_EXPORT Builder { + public: + explicit Builder(::arrow::util::SecureString footer_key) + : parquet_cipher_(kDefaultEncryptionAlgorithm), + encrypted_footer_(kDefaultEncryptedFooter), + footer_key_(std::move(footer_key)) { + store_aad_prefix_in_file_ = false; + } + + /// Create files with plaintext footer. + /// If not called, the files will be created with encrypted footer (default). + Builder* set_plaintext_footer() { + encrypted_footer_ = false; + return this; + } + + /// Set encryption algorithm. + /// If not called, files will be encrypted with AES_GCM_V1 (default). + Builder* algorithm(ParquetCipher::type parquet_cipher) { + parquet_cipher_ = parquet_cipher; + return this; + } + + /// Set a key retrieval metadata (converted from String). + /// use either footer_key_metadata or footer_key_id, not both. + Builder* footer_key_id(std::string key_id); + + /// Set a key retrieval metadata. + /// use either footer_key_metadata or footer_key_id, not both. + Builder* footer_key_metadata(std::string footer_key_metadata); + + /// Set the file AAD Prefix. + Builder* aad_prefix(std::string aad_prefix); + + /// Skip storing AAD Prefix in file. + /// If not called, and if AAD Prefix is set, it will be stored. + Builder* disable_aad_prefix_storage(); + + /// Set the list of encrypted columns and their properties (keys etc). + /// If not called, all columns will be encrypted with the footer key. + /// If called, the file columns not in the list will be left unencrypted. + Builder* encrypted_columns(ColumnPathToEncryptionPropertiesMap encrypted_columns); + + std::shared_ptr build() { + return std::shared_ptr(new FileEncryptionProperties( + parquet_cipher_, footer_key_, footer_key_metadata_, encrypted_footer_, + aad_prefix_, store_aad_prefix_in_file_, encrypted_columns_)); + } + + private: + ParquetCipher::type parquet_cipher_; + bool encrypted_footer_; + ::arrow::util::SecureString footer_key_; + std::string footer_key_metadata_; + + std::string aad_prefix_; + bool store_aad_prefix_in_file_; + ColumnPathToEncryptionPropertiesMap encrypted_columns_; + }; + + bool encrypted_footer() const { return encrypted_footer_; } + + EncryptionAlgorithm algorithm() const { return algorithm_; } + + const ::arrow::util::SecureString& footer_key() const { return footer_key_; } + + const std::string& footer_key_metadata() const { return footer_key_metadata_; } + + const std::string& file_aad() const { return file_aad_; } + + std::shared_ptr column_encryption_properties( + const std::string& column_path); + + const ColumnPathToEncryptionPropertiesMap& encrypted_columns() const { + return encrypted_columns_; + } + + private: + EncryptionAlgorithm algorithm_; + ::arrow::util::SecureString footer_key_; + std::string footer_key_metadata_; + bool encrypted_footer_; + std::string file_aad_; + std::string aad_prefix_; + bool store_aad_prefix_in_file_; + ColumnPathToEncryptionPropertiesMap encrypted_columns_; + + FileEncryptionProperties(ParquetCipher::type cipher, + ::arrow::util::SecureString footer_key, + std::string footer_key_metadata, bool encrypted_footer, + std::string aad_prefix, bool store_aad_prefix_in_file, + ColumnPathToEncryptionPropertiesMap encrypted_columns); +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_material_store.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_material_store.h new file mode 100644 index 0000000000000000000000000000000000000000..83f028a4bc1e9e0d24e21e7acfb785af0e5b37f7 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_material_store.h @@ -0,0 +1,57 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/filesystem/filesystem.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +/// Stores encryption key material outside the Parquet file, for example in a separate +/// small file in the same folder. This is important for “key rotation”, when MEKs have to +/// be changed (if compromised; or periodically, just in case) - without modifying the +/// Parquet files (often immutable). +class PARQUET_EXPORT FileKeyMaterialStore { + public: + /// Add key material for one encryption key. + virtual void AddKeyMaterial(std::string key_id_in_file, std::string key_material) = 0; + + /// Get key material + virtual std::string GetKeyMaterial(std::string key_id_in_file) = 0; + + /// After key material was added for all keys in the given Parquet file, + /// save material in persistent store. + virtual void SaveMaterial() = 0; + + /// Remove key material from persistent store. Used in key rotation. + virtual void RemoveMaterial() = 0; + + /// Move key material to another store. Used in key rotation. + virtual void MoveMaterialTo(std::shared_ptr target_key_store) = 0; + + /// Returns the Set of all key IDs in this store (for the given Parquet file) + virtual std::vector GetKeyIDSet() = 0; + + virtual ~FileKeyMaterialStore() {} +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_unwrapper.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_unwrapper.h new file mode 100644 index 0000000000000000000000000000000000000000..c86f68121c87231549f129fb5714d0b266b8dd1b --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_unwrapper.h @@ -0,0 +1,96 @@ +// 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. + +#pragma once + +#include "arrow/util/concurrent_map.h" +#include "arrow/util/secure_string.h" + +#include "parquet/encryption/encryption.h" +#include "parquet/encryption/file_system_key_material_store.h" +#include "parquet/encryption/key_material.h" +#include "parquet/encryption/key_toolkit.h" +#include "parquet/encryption/key_toolkit_internal.h" +#include "parquet/encryption/kms_client.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +// This class will retrieve the key from "key metadata", following these steps: +// 1. Parse "key metadata" (see structure in KeyMetadata class). +// 2. Retrieve "key material" which can be stored inside or outside "key metadata". +// 3. Unwrap the "data encryption key" from "key material". There are 2 modes: +// 3.1. single wrapping: decrypt the wrapped "data encryption key" directly with "master +// encryption key" 3.2. double wrapping: 2 steps: 3.2.1. "key encryption key" is decrypted +// with "master encryption key" 3.2.2. "data encryption key" is decrypted with the above +// "key encryption key" +class PARQUET_EXPORT FileKeyUnwrapper : public DecryptionKeyRetriever { + public: + /// key_toolkit and kms_connection_config is to get KmsClient from cache or create + /// KmsClient if it's not in the cache yet. cache_entry_lifetime_seconds is life time of + /// KmsClient in the cache. + /// If the file uses external key material then the Parquet file path and file + /// system must be specified. + FileKeyUnwrapper(std::shared_ptr key_toolkit, + const KmsConnectionConfig& kms_connection_config, + double cache_lifetime_seconds, const std::string& file_path = "", + const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR); + + /// Constructor overload that takes a raw pointer to the KeyToolkit + FileKeyUnwrapper(KeyToolkit* key_toolkit, + const KmsConnectionConfig& kms_connection_config, + double cache_lifetime_seconds, const std::string& file_path = "", + const std::shared_ptr<::arrow::fs::FileSystem>& file_system = NULLPTR); + + /// Constructor overload that takes a raw pointer to the KeyToolkit and + /// accepts an existing key_material_store rather than using + /// the file path and file system to create one when needed. + FileKeyUnwrapper(KeyToolkit* key_toolkit, + const KmsConnectionConfig& kms_connection_config, + double cache_lifetime_seconds, + std::shared_ptr key_material_store); + + /// Get the data key from key metadata + ::arrow::util::SecureString GetKey(const std::string& key_metadata_bytes) override; + + /// Get the data key along with the master key id from key material + KeyWithMasterId GetDataEncryptionKey(const KeyMaterial& key_material); + + private: + FileKeyUnwrapper(std::shared_ptr key_toolkit_owner, KeyToolkit* key_toolkit, + const KmsConnectionConfig& kms_connection_config, + double cache_lifetime_seconds, + std::shared_ptr key_material_store, + const std::string& file_path, + const std::shared_ptr<::arrow::fs::FileSystem>& file_system); + + std::shared_ptr GetKmsClientFromConfigOrKeyMaterial( + const KeyMaterial& key_material); + + /// A map of Key Encryption Key (KEK) ID -> KEK bytes, for the current token + std::shared_ptr<::arrow::util::ConcurrentMap> + kek_per_kek_id_; + std::shared_ptr key_toolkit_owner_; + KeyToolkit* key_toolkit_; + KmsConnectionConfig kms_connection_config_; + const double cache_entry_lifetime_seconds_; + std::shared_ptr key_material_store_; + const std::string file_path_; + std::shared_ptr<::arrow::fs::FileSystem> file_system_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_wrapper.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_wrapper.h new file mode 100644 index 0000000000000000000000000000000000000000..aa6d878bafe19c73a2bfc099f81df0bba9ce49a9 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_key_wrapper.h @@ -0,0 +1,84 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/util/concurrent_map.h" + +#include "parquet/encryption/file_key_material_store.h" +#include "parquet/encryption/key_encryption_key.h" +#include "parquet/encryption/key_toolkit.h" +#include "parquet/encryption/kms_client.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +// This class will generate "key metadata" from "data encryption key" and "master key", +// following these steps: +// 1. Wrap "data encryption key". There are 2 modes: +// 1.1. single wrapping: encrypt "data encryption key" directly with "master encryption +// key" +// 1.2. double wrapping: 2 steps: +// 1.2.1. "key encryption key" is randomized (see KeyEncryptionKey class) +// 1.2.2. "data encryption key" is encrypted with the above "key encryption key" +// 2. Create "key material" (see structure in KeyMaterial class) +// 3. Create "key metadata" with "key material" inside or a reference to outside "key +// material" (see structure in KeyMetadata class). +class PARQUET_EXPORT FileKeyWrapper { + public: + static constexpr int kKeyEncryptionKeyLength = 16; + static constexpr int kKeyEncryptionKeyIdLength = 16; + + /// key_toolkit and kms_connection_config is to get KmsClient from the cache or create + /// KmsClient if it's not in the cache yet. cache_entry_lifetime_seconds is life time of + /// KmsClient in the cache. key_material_store is to store "key material" outside + /// parquet file, NULL if "key material" is stored inside parquet file. + FileKeyWrapper(KeyToolkit* key_toolkit, + const KmsConnectionConfig& kms_connection_config, + std::shared_ptr key_material_store, + double cache_entry_lifetime_seconds, bool double_wrapping); + + /// Creates key_metadata field for a given data key, via wrapping the key with the + /// master key. + /// When external key material is used, an identifier is usually generated automatically + /// but may be specified explicitly to support key rotation, + /// which requires keeping the same identifiers. + std::string GetEncryptionKeyMetadata(const ::arrow::util::SecureString& data_key, + const std::string& master_key_id, + bool is_footer_key, + std::string key_id_in_file = ""); + + private: + KeyEncryptionKey CreateKeyEncryptionKey(const std::string& master_key_id); + + /// A map of Master Encryption Key ID -> KeyEncryptionKey, for the current token + std::shared_ptr<::arrow::util::ConcurrentMap> + kek_per_master_key_id_; + + std::shared_ptr kms_client_; + KmsConnectionConfig kms_connection_config_; + std::shared_ptr key_material_store_; + const double cache_entry_lifetime_seconds_; + const bool double_wrapping_; + uint16_t key_counter_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_system_key_material_store.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_system_key_material_store.h new file mode 100644 index 0000000000000000000000000000000000000000..ecbadb90a94833881ca851aa43fd213ba7ed93c2 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/file_system_key_material_store.h @@ -0,0 +1,93 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/filesystem/filesystem.h" + +#include "parquet/encryption/file_key_material_store.h" +#include "parquet/exception.h" + +namespace parquet::encryption { + +/// A FileKeyMaterialStore that stores key material in a file system file in the same +/// folder as the Parquet file. +class PARQUET_EXPORT FileSystemKeyMaterialStore : public FileKeyMaterialStore { + public: + static constexpr const char kKeyMaterialFilePrefix[] = "_KEY_MATERIAL_FOR_"; + static constexpr const char kTempFilePrefix[] = "_TMP"; + static constexpr const char kKeyMaterialFileSuffix[] = ".json"; + + FileSystemKeyMaterialStore() {} + FileSystemKeyMaterialStore(std::string key_material_file_path, + std::shared_ptr<::arrow::fs::FileSystem> file_system); + + /// Creates a new file system key material store for a parquet file. + /// When use_tmp_prefix is true, files are saved with an extra _TMP prefix so they don't + /// conflict with existing external material files. This is useful during key rotation + /// so that temporary key material files can be created while using the existing key + /// material, before moving the key material to the non-temporary location. + static std::shared_ptr Make( + std::string parquet_file_path, std::shared_ptr<::arrow::fs::FileSystem> file_system, + bool use_tmp_prefix); + + /// Add key material for one encryption key. + void AddKeyMaterial(std::string key_id_in_file, std::string key_material) { + key_material_map_.emplace(std::move(key_id_in_file), std::move(key_material)); + } + + /// Get key material + std::string GetKeyMaterial(std::string key_id_in_file) { + if (key_material_map_.empty()) { + LoadKeyMaterialMap(); + } + auto found = key_material_map_.find(key_id_in_file); + if (found == key_material_map_.end()) { + throw ParquetException("Invalid key id"); + } + return found->second; + } + + /// After key material was added for all keys in the given Parquet file, + /// save material in persistent store. + void SaveMaterial(); + + /// Remove key material from persistent store. Used in key rotation. + void RemoveMaterial(); + + /// Move key material to another store. Used in key rotation. + void MoveMaterialTo(std::shared_ptr target_key_store); + + /// Returns the Set of all key IDs in this store (for the given Parquet file) + std::vector GetKeyIDSet(); + + private: + std::string GetStorageFilePath() { return key_material_file_path_; } + + std::string BuildKeyMaterialMapJson(); + void LoadKeyMaterialMap(); + std::string key_material_file_path_; + std::shared_ptr<::arrow::fs::FileSystem> file_system_; + /// Maps ID of a key in Parquet file and key material + std::unordered_map key_material_map_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_encryption_key.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_encryption_key.h new file mode 100644 index 0000000000000000000000000000000000000000..1157937632afb1deca6568d75655c6195eac564d --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_encryption_key.h @@ -0,0 +1,58 @@ +// 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. + +#pragma once + +#include +#include + +#include "arrow/util/base64.h" +#include "arrow/util/secure_string.h" + +namespace parquet::encryption { + +// In the double wrapping mode, each "data encryption key" (DEK) is encrypted with a “key +// encryption key” (KEK), that in turn is encrypted with a "master encryption key" (MEK). +// In a writer process, a random KEK is generated for each MEK ID, and cached in a map. This allows to perform an interaction with a KMS server only once for each +// MEK, in order to wrap its KEK. "Data encryption key" (DEK) wrapping is performed +// locally, and does not involve an interaction with a KMS server. +class KeyEncryptionKey { + public: + KeyEncryptionKey(::arrow::util::SecureString kek_bytes, std::string kek_id, + std::string encoded_wrapped_kek) + : kek_bytes_(std::move(kek_bytes)), + kek_id_(std::move(kek_id)), + encoded_kek_id_(::arrow::util::base64_encode(kek_id_)), + encoded_wrapped_kek_(std::move(encoded_wrapped_kek)) {} + + const ::arrow::util::SecureString& kek_bytes() const { return kek_bytes_; } + + const std::string& kek_id() const { return kek_id_; } + + const std::string& encoded_kek_id() const { return encoded_kek_id_; } + + const std::string& encoded_wrapped_kek() const { return encoded_wrapped_kek_; } + + private: + ::arrow::util::SecureString kek_bytes_; + std::string kek_id_; + std::string encoded_kek_id_; + std::string encoded_wrapped_kek_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_material.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_material.h new file mode 100644 index 0000000000000000000000000000000000000000..3e7e862c996d3f0b0c016f3953dc40dcb314a8a0 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_material.h @@ -0,0 +1,129 @@ +// 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. + +#pragma once + +#include + +#include "parquet/platform.h" + +namespace arrow { +namespace json { +namespace internal { +class ObjectParser; +} // namespace internal +} // namespace json +} // namespace arrow + +namespace parquet::encryption { + +// KeyMaterial class represents the "key material", keeping the information that allows +// readers to recover an encryption key (see description of the KeyMetadata class). The +// keytools package (PARQUET-1373) implements the "envelope encryption" pattern, in a +// "single wrapping" or "double wrapping" mode. In the single wrapping mode, the key +// material is generated by encrypting the "data encryption key" (DEK) by a "master key". +// In the double wrapping mode, the key material is generated by encrypting the DEK by a +// "key encryption key" (KEK), that in turn is encrypted by a "master key". +// +// Key material is kept in a flat json object, with the following fields: +// 1. "keyMaterialType" - a String, with the type of key material. In the current +// version, only one value is allowed - "PKMT1" (stands +// for "parquet key management tools, version 1"). For external key material storage, +// this field is written in both "key metadata" and "key material" jsons. For internal +// key material storage, this field is written only once in the common json. +// 2. "isFooterKey" - a boolean. If true, means that the material belongs to a file footer +// key, and keeps additional information (such as +// KMS instance ID and URL). If false, means that the material belongs to a column +// key. +// 3. "kmsInstanceID" - a String, with the KMS Instance ID. Written only in footer key +// material. +// 4. "kmsInstanceURL" - a String, with the KMS Instance URL. Written only in footer key +// material. +// 5. "masterKeyID" - a String, with the ID of the master key used to generate the +// material. +// 6. "wrappedDEK" - a String, with the wrapped DEK (base64 encoding). +// 7. "doubleWrapping" - a boolean. If true, means that the material was generated in +// double wrapping mode. +// If false - in single wrapping mode. +// 8. "keyEncryptionKeyID" - a String, with the ID of the KEK used to generate the +// material. Written only in double wrapping mode. +// 9. "wrappedKEK" - a String, with the wrapped KEK (base64 encoding). Written only in +// double wrapping mode. +class PARQUET_EXPORT KeyMaterial { + public: + // these fields are defined in a specification and should never be changed + static constexpr const char kKeyMaterialTypeField[] = "keyMaterialType"; + static constexpr const char kKeyMaterialType1[] = "PKMT1"; + + static constexpr const char kFooterKeyIdInFile[] = "footerKey"; + static constexpr const char kColumnKeyIdInFilePrefix[] = "columnKey"; + + static constexpr const char kIsFooterKeyField[] = "isFooterKey"; + static constexpr const char kDoubleWrappingField[] = "doubleWrapping"; + static constexpr const char kKmsInstanceIdField[] = "kmsInstanceID"; + static constexpr const char kKmsInstanceUrlField[] = "kmsInstanceURL"; + static constexpr const char kMasterKeyIdField[] = "masterKeyID"; + static constexpr const char kWrappedDataEncryptionKeyField[] = "wrappedDEK"; + static constexpr const char kKeyEncryptionKeyIdField[] = "keyEncryptionKeyID"; + static constexpr const char kWrappedKeyEncryptionKeyField[] = "wrappedKEK"; + + public: + KeyMaterial() = default; + + static KeyMaterial Parse(const std::string& key_material_string); + + static KeyMaterial Parse( + const ::arrow::json::internal::ObjectParser* key_material_json); + + /// This method returns a json string that will be stored either inside a parquet file + /// or in a key material store outside the parquet file. + static std::string SerializeToJson(bool is_footer_key, + const std::string& kms_instance_id, + const std::string& kms_instance_url, + const std::string& master_key_id, + bool is_double_wrapped, const std::string& kek_id, + const std::string& encoded_wrapped_kek, + const std::string& encoded_wrapped_dek, + bool is_internal_storage); + + bool is_footer_key() const { return is_footer_key_; } + bool is_double_wrapped() const { return is_double_wrapped_; } + const std::string& master_key_id() const { return master_key_id_; } + const std::string& wrapped_dek() const { return encoded_wrapped_dek_; } + const std::string& kek_id() const { return kek_id_; } + const std::string& wrapped_kek() const { return encoded_wrapped_kek_; } + const std::string& kms_instance_id() const { return kms_instance_id_; } + const std::string& kms_instance_url() const { return kms_instance_url_; } + + private: + KeyMaterial(bool is_footer_key, const std::string& kms_instance_id, + const std::string& kms_instance_url, const std::string& master_key_id, + bool is_double_wrapped, const std::string& kek_id, + const std::string& encoded_wrapped_kek, + const std::string& encoded_wrapped_dek); + + bool is_footer_key_; + std::string kms_instance_id_; + std::string kms_instance_url_; + std::string master_key_id_; + bool is_double_wrapped_; + std::string kek_id_; + std::string encoded_wrapped_kek_; + std::string encoded_wrapped_dek_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_metadata.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_metadata.h new file mode 100644 index 0000000000000000000000000000000000000000..6fe8ac7ccb9db3fb92da42064f9fe2aeabdbfb52 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_metadata.h @@ -0,0 +1,91 @@ +// 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. + +#pragma once + +#include +#include + +#include "parquet/encryption/key_material.h" +#include "parquet/exception.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +// Parquet encryption specification defines "key metadata" as an arbitrary byte array, +// generated by file writers for each encryption key, and passed to the low level API for +// storage in the file footer. The "key metadata" field is made available to file readers +// to enable recovery of the key. This interface can be utilized for implementation +// of any key management scheme. +// +// The keytools package (PARQUET-1373) implements one approach, of many possible, to key +// management and to generation of the "key metadata" fields. This approach, based on the +// "envelope encryption" pattern, allows integration with KMS servers. It keeps the actual +// material, required to recover a key, in a "key material" object (see the KeyMaterial +// class for details). This class is implemented to support version 1 of the parquet key +// management tools specification. +// +// KeyMetadata writes (and reads) the "key metadata" field as a flat json object, +// with the following fields: +// 1. "keyMaterialType" - a String, with the type of key material. +// 2. "internalStorage" - a boolean. If true, means that "key material" is kept inside the +// "key metadata" field. If false, "key material" is kept externally (outside Parquet +// files) - in this case, "key metadata" keeps a reference to the external "key material". +// 3. "keyReference" - a String, with the reference to the external "key material". +// Written only if internalStorage is false. +// +// If internalStorage is true, "key material" is a part of "key metadata", and the json +// keeps additional fields, described in the KeyMaterial class. +class PARQUET_EXPORT KeyMetadata { + public: + static constexpr const char kKeyMaterialInternalStorageField[] = "internalStorage"; + static constexpr const char kKeyReferenceField[] = "keyReference"; + + /// key_metadata_bytes is the key metadata field stored in the parquet file, + /// in the serialized json object format. + static KeyMetadata Parse(const std::string& key_metadata_bytes); + + static std::string CreateSerializedForExternalMaterial( + const std::string& key_reference); + + bool key_material_stored_internally() const { return is_internal_storage_; } + + const KeyMaterial& key_material() const { + if (!is_internal_storage_) { + throw ParquetException("key material is stored externally."); + } + return ::std::get(key_material_or_reference_); + } + + const std::string& key_reference() const { + if (is_internal_storage_) { + throw ParquetException("key material is stored internally."); + } + return ::std::get(key_material_or_reference_); + } + + private: + explicit KeyMetadata(const KeyMaterial& key_material); + explicit KeyMetadata(const std::string& key_reference); + + bool is_internal_storage_; + /// If is_internal_storage_ is true, KeyMaterial is set, + /// else a string referencing to an outside "key material" is set. + ::std::variant key_material_or_reference_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_toolkit.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_toolkit.h new file mode 100644 index 0000000000000000000000000000000000000000..a0b929877eebd71b416d771d003bfaf1abe43369 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/key_toolkit.h @@ -0,0 +1,106 @@ +// 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. + +#pragma once + +#include +#include + +#include "parquet/encryption/key_encryption_key.h" +#include "parquet/encryption/kms_client.h" +#include "parquet/encryption/kms_client_factory.h" +#include "parquet/encryption/two_level_cache_with_expiration.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +static constexpr uint64_t kCacheCleanPeriodForKeyRotation = 60 * 60; // 1 hour + +// KeyToolkit is a utility that keeps various tools for key management (such as key +// rotation, kms client instantiation, cache control, etc), plus a number of auxiliary +// classes for internal use. +class PARQUET_EXPORT KeyToolkit { + public: + KeyToolkit() { last_cache_clean_for_key_rotation_time_ = {}; } + + /// KMS client two level cache: token -> KMSInstanceId -> KmsClient + TwoLevelCacheWithExpiration>& kms_client_cache_per_token() { + return kms_client_cache_; + } + /// Key encryption key two level cache for wrapping: token -> MasterEncryptionKeyId -> + /// KeyEncryptionKey + TwoLevelCacheWithExpiration& kek_write_cache_per_token() { + return key_encryption_key_write_cache_; + } + + /// Key encryption key two level cache for unwrapping: token -> KeyEncryptionKeyId -> + /// KeyEncryptionKeyBytes + TwoLevelCacheWithExpiration<::arrow::util::SecureString>& kek_read_cache_per_token() { + return key_encryption_key_read_cache_; + } + + std::shared_ptr GetKmsClient( + const KmsConnectionConfig& kms_connection_config, double cache_entry_lifetime_ms); + + /// Flush any caches that are tied to the (compromised) access_token + void RemoveCacheEntriesForToken(const std::string& access_token); + + void RemoveCacheEntriesForAllTokens(); + + void RegisterKmsClientFactory(std::shared_ptr kms_client_factory) { + if (kms_client_factory_ != NULLPTR) { + throw ParquetException("KMS client factory has already been registered."); + } + kms_client_factory_ = std::move(kms_client_factory); + } + + /// Key rotation. In the single wrapping mode, decrypts data keys with old master keys, + /// then encrypts them with new master keys. In the double wrapping mode, decrypts KEKs + /// (key encryption keys) with old master keys, generates new KEKs and encrypts them + /// with new master keys. Works only if key material is not stored internally in file + /// footers. Not supported in local key wrapping mode. Method can be run by multiple + /// threads, but each thread must work on different files. + void RotateMasterKeys(const KmsConnectionConfig& kms_connection_config, + const std::string& parquet_file_path, + const std::shared_ptr<::arrow::fs::FileSystem>& file_system, + bool double_wrapping, double cache_lifetime_seconds); + + private: + TwoLevelCacheWithExpiration> kms_client_cache_; + TwoLevelCacheWithExpiration key_encryption_key_write_cache_; + TwoLevelCacheWithExpiration<::arrow::util::SecureString> key_encryption_key_read_cache_; + std::shared_ptr kms_client_factory_; + mutable ::arrow::util::Mutex last_cache_clean_for_key_rotation_time_mutex_; + internal::TimePoint last_cache_clean_for_key_rotation_time_; +}; + +// "data encryption key" and "master key identifier" are paired together as output when +// parsing from "key material" +class PARQUET_EXPORT KeyWithMasterId { + public: + KeyWithMasterId(::arrow::util::SecureString key_bytes, std::string master_id) + : key_bytes_(std::move(key_bytes)), master_id_(std::move(master_id)) {} + + const ::arrow::util::SecureString& data_key() const { return key_bytes_; } + const std::string& master_id() const { return master_id_; } + + private: + ::arrow::util::SecureString key_bytes_; + std::string master_id_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client.h new file mode 100644 index 0000000000000000000000000000000000000000..9c67e7cae492de85df5ab656b3f0c0e1f92a7b96 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client.h @@ -0,0 +1,97 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "arrow/util/mutex.h" +#include "arrow/util/secure_string.h" + +#include "parquet/exception.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +/// This class wraps the key access token of a KMS server. If your token changes over +/// time, you should keep the reference to the KeyAccessToken object and call Refresh() +/// method every time you have a new token. +class PARQUET_EXPORT KeyAccessToken { + public: + KeyAccessToken() = default; + + explicit KeyAccessToken(const std::string value) : value_(value) {} + + void Refresh(const std::string& new_value) { + auto lock = mutex_.Lock(); + value_ = new_value; + } + + const std::string& value() const { + auto lock = mutex_.Lock(); + return value_; + } + + private: + std::string value_; + mutable ::arrow::util::Mutex mutex_; +}; + +struct PARQUET_EXPORT KmsConnectionConfig { + std::string kms_instance_id; + std::string kms_instance_url; + /// If the access token is changed in the future, you should keep a reference to + /// this object and call Refresh() on it whenever there is a new access token. + std::shared_ptr refreshable_key_access_token; + std::unordered_map custom_kms_conf; + + KmsConnectionConfig(); + + const std::string& key_access_token() const { + if (refreshable_key_access_token == NULLPTR || + refreshable_key_access_token->value().empty()) { + throw ParquetException("key access token is not set!"); + } + return refreshable_key_access_token->value(); + } + + void SetDefaultIfEmpty(); +}; + +class PARQUET_EXPORT KmsClient { + public: + static constexpr const char kKmsInstanceIdDefault[] = "DEFAULT"; + static constexpr const char kKmsInstanceUrlDefault[] = "DEFAULT"; + static constexpr const char kKeyAccessTokenDefault[] = "DEFAULT"; + + /// \brief Wraps a key. + /// + /// Encrypts it with the master key, encodes the result + /// and potentially adds a KMS-specific metadata. + virtual std::string WrapKey(const ::arrow::util::SecureString& key_bytes, + const std::string& master_key_identifier) = 0; + + /// \brief Decrypts (unwraps) a key with the master key. + virtual ::arrow::util::SecureString UnwrapKey( + const std::string& wrapped_key, const std::string& master_key_identifier) = 0; + + virtual ~KmsClient() {} +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client_factory.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client_factory.h new file mode 100644 index 0000000000000000000000000000000000000000..7a7c77c7eebbfbb687575acb12b89c1c2e99461a --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/kms_client_factory.h @@ -0,0 +1,38 @@ +// 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. + +#pragma once + +#include "parquet/encryption/kms_client.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +class PARQUET_EXPORT KmsClientFactory { + public: + explicit KmsClientFactory(bool wrap_locally = false) : wrap_locally_(wrap_locally) {} + + virtual ~KmsClientFactory() = default; + + virtual std::shared_ptr CreateKmsClient( + const KmsConnectionConfig& kms_connection_config) = 0; + + protected: + bool wrap_locally_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/local_wrap_kms_client.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/local_wrap_kms_client.h new file mode 100644 index 0000000000000000000000000000000000000000..607c75a4c2e8828a55a73c85c48c0551d1b72aa2 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/local_wrap_kms_client.h @@ -0,0 +1,95 @@ +// 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. + +#pragma once + +#include +#include + +#include "arrow/util/concurrent_map.h" + +#include "parquet/encryption/kms_client.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +/// This class supports local wrapping mode, master keys will be fetched from the KMS +/// server and used to encrypt other keys (data encryption keys or key encryption keys). +class PARQUET_EXPORT LocalWrapKmsClient : public KmsClient { + public: + static constexpr const char kLocalWrapNoKeyVersion[] = "NO_VERSION"; + + explicit LocalWrapKmsClient(const KmsConnectionConfig& kms_connection_config); + + std::string WrapKey(const ::arrow::util::SecureString& key_bytes, + const std::string& master_key_identifier) override; + + ::arrow::util::SecureString UnwrapKey( + const std::string& wrapped_key, const std::string& master_key_identifier) override; + + protected: + /// Get master key from the remote KMS server. + /// Note: this function might be called by multiple threads + virtual const ::arrow::util::SecureString& GetMasterKeyFromServer( + const std::string& master_key_identifier) = 0; + + private: + /// KMS systems wrap keys by encrypting them by master keys, and attaching additional + /// information (such as the version number of the masker key) to the result of + /// encryption. The master key version is required in key rotation. Currently, the + /// local wrapping mode does not support key rotation (because not all KMS systems allow + /// to fetch a master key by its ID and version number). Still, the local wrapping mode + /// adds a placeholder for the master key version, that will enable support for key + /// rotation in this mode in the future, with appropriate KMS systems. This will also + /// enable backward compatibility, where future readers will be able to extract master + /// key version in the files written by the current code. + /// + /// LocalKeyWrap class writes (and reads) the "key wrap" as a flat json with the + /// following fields: + /// 1. "masterKeyVersion" - a String, with the master key version. In the current + /// version, only one value is allowed - "NO_VERSION". + /// 2. "encryptedKey" - a String, with the key encrypted by the master key + /// (base64-encoded). + class LocalKeyWrap { + public: + static constexpr const char kLocalWrapKeyVersionField[] = "masterKeyVersion"; + static constexpr const char kLocalWrapEncryptedKeyField[] = "encryptedKey"; + + LocalKeyWrap(std::string master_key_version, std::string encrypted_encoded_key); + + static std::string CreateSerialized(const std::string& encrypted_encoded_key); + + static LocalKeyWrap Parse(const std::string& wrapped_key); + + const std::string& master_key_version() const { return master_key_version_; } + + const std::string& encrypted_encoded_key() const { return encrypted_encoded_key_; } + + private: + std::string encrypted_encoded_key_; + std::string master_key_version_; + }; + + const ::arrow::util::SecureString& GetKeyFromServer(const std::string& key_identifier); + + protected: + KmsConnectionConfig kms_connection_config_; + ::arrow::util::ConcurrentMap + master_key_cache_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_encryption_util.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_encryption_util.h new file mode 100644 index 0000000000000000000000000000000000000000..46b84a0d0a33d9a71570403ff4a0561da91b3464 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_encryption_util.h @@ -0,0 +1,136 @@ +// 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. + +// This module defines an abstract interface for iterating through pages in a +// Parquet column chunk within a row group. It could be extended in the future +// to iterate through all data pages in all chunks in a file. + +#pragma once + +#include +#include +#include + +#include + +#include "arrow/filesystem/filesystem.h" +#include "arrow/filesystem/localfs.h" +#include "arrow/status.h" +#include "arrow/util/io_util.h" +#include "arrow/util/secure_string.h" + +#include "parquet/encryption/encryption.h" +#include "parquet/test_util.h" + +namespace parquet { +class ParquetFileReader; +namespace encryption::test { + +using ::arrow::internal::TemporaryDir; +using ::arrow::util::SecureString; + +constexpr int kFixedLength = 10; + +const SecureString kFooterEncryptionKey("0123456789012345"); +const SecureString kColumnEncryptionKey1("1234567890123450"); +const SecureString kColumnEncryptionKey2("1234567890123451"); +const SecureString kColumnEncryptionKey3("1234567890123452"); +const char kFileName[] = "tester"; + +// Get the path of file inside parquet test data directory +std::string data_file(const char* file); + +// A temporary directory that contains the encrypted files generated in the tests. +extern std::unique_ptr temp_dir; + +inline ::arrow::Result> temp_data_dir() { + return TemporaryDir::Make("parquet-encryption-test-"); +} + +const char kDoubleFieldName[] = "double_field"; +const char kFloatFieldName[] = "float_field"; +const char kBooleanFieldName[] = "boolean_field"; +const char kInt32FieldName[] = "int32_field"; +const char kInt64FieldName[] = "int64_field"; +const char kInt96FieldName[] = "int96_field"; +const char kByteArrayFieldName[] = "ba_field"; +const char kFixedLenByteArrayFieldName[] = "flba_field"; + +const char kFooterMasterKey[] = "0123456789012345"; +const char kFooterMasterKeyId[] = "kf"; +const char* const kColumnMasterKeys[] = {"1234567890123450", "1234567890123451", + "1234567890123452", "1234567890123453", + "1234567890123454", "1234567890123455"}; +const char* const kColumnMasterKeyIds[] = {"kc1", "kc2", "kc3", "kc4", "kc5", "kc6"}; + +// New master key values used to simulate key rotation +const char kNewFooterMasterKey[] = "9123456789012345"; +const char* const kNewColumnMasterKeys[] = {"9234567890123450", "9234567890123451", + "9234567890123452", "9234567890123453", + "9234567890123454", "9234567890123455"}; + +// The result of this function will be used to set into TestOnlyInMemoryKmsClientFactory +// as the key mapping to look at. +std::unordered_map BuildKeyMap(const char* const* column_ids, + const char* const* column_keys, + const char* footer_id, + const char* footer_key); + +// The result of this function will be used to set into EncryptionConfiguration +// as column keys. +std::string BuildColumnKeyMapping(); + +// FileEncryptor and FileDecryptor are helper classes to write/read an encrypted parquet +// file corresponding to each pair of FileEncryptionProperties/FileDecryptionProperties. +// FileEncryptor writes the file with fixed data values and FileDecryptor reads the file +// and verify the correctness of data values. +class FileEncryptor { + public: + FileEncryptor(); + + void EncryptFile( + std::string file, + std::shared_ptr encryption_configurations); + + private: + std::shared_ptr SetupEncryptionSchema(); + + int num_rowgroups_ = 5; + int rows_per_rowgroup_ = 50; + std::shared_ptr schema_; +}; + +class FileDecryptor { + public: + void DecryptFile( + const std::string& file_name, + const std::shared_ptr& file_decryption_properties); + void DecryptPageIndex( + const std::string& file_name, + const std::shared_ptr& file_decryption_properties); + + private: + void CheckFile( + parquet::ParquetFileReader* file_reader, + const std::shared_ptr& file_decryption_properties); + void CheckPageIndex( + parquet::ParquetFileReader* file_reader, + const std::shared_ptr& file_decryption_properties); +}; + +} // namespace encryption::test +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_in_memory_kms.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_in_memory_kms.h new file mode 100644 index 0000000000000000000000000000000000000000..b9d4169c6345f9f1d565615919830d94f107479c --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/test_in_memory_kms.h @@ -0,0 +1,101 @@ +// 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. + +#pragma once + +#include + +#include "arrow/util/base64.h" + +#include "parquet/encryption/kms_client_factory.h" +#include "parquet/encryption/local_wrap_kms_client.h" +#include "parquet/platform.h" + +namespace parquet::encryption { + +// This is a mock class, built for testing only. Don't use it as an example of +// LocalWrapKmsClient implementation. +class TestOnlyLocalWrapInMemoryKms : public LocalWrapKmsClient { + public: + explicit TestOnlyLocalWrapInMemoryKms(const KmsConnectionConfig& kms_connection_config); + + static void InitializeMasterKeys( + const std::unordered_map& + master_keys_map); + + protected: + const ::arrow::util::SecureString& GetMasterKeyFromServer( + const std::string& master_key_identifier) override; + + private: + static std::unordered_map master_key_map_; +}; + +// This is a mock class, built for testing only. Don't use it as an example of KmsClient +// implementation. +class TestOnlyInServerWrapKms : public KmsClient { + public: + static void InitializeMasterKeys( + const std::unordered_map& + master_keys_map); + + std::string WrapKey(const ::arrow::util::SecureString& key_bytes, + const std::string& master_key_identifier) override; + + ::arrow::util::SecureString UnwrapKey( + const std::string& wrapped_key, const std::string& master_key_identifier) override; + + static void StartKeyRotation( + const std::unordered_map& + new_master_keys_map); + static void FinishKeyRotation(); + + private: + ::arrow::util::SecureString GetMasterKeyFromServer( + const std::string& master_key_identifier); + + // Different wrapping and unwrapping key maps to imitate versioning + // and support key rotation. + static std::unordered_map + unwrapping_master_key_map_; + static std::unordered_map + wrapping_master_key_map_; +}; + +// This is a mock class, built for testing only. Don't use it as an example of +// KmsClientFactory implementation. +class TestOnlyInMemoryKmsClientFactory : public KmsClientFactory { + public: + TestOnlyInMemoryKmsClientFactory( + bool wrap_locally, + const std::unordered_map& master_keys_map) + : KmsClientFactory(wrap_locally) { + TestOnlyLocalWrapInMemoryKms::InitializeMasterKeys(master_keys_map); + TestOnlyInServerWrapKms::InitializeMasterKeys(master_keys_map); + } + + std::shared_ptr CreateKmsClient( + const KmsConnectionConfig& kms_connection_config) { + if (wrap_locally_) { + return std::make_shared(kms_connection_config); + } else { + return std::make_shared(); + } + } +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/two_level_cache_with_expiration.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/two_level_cache_with_expiration.h new file mode 100644 index 0000000000000000000000000000000000000000..283ebd97b714c14638b43aae7b9acc75c2c7b51b --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/two_level_cache_with_expiration.h @@ -0,0 +1,149 @@ +// 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. + +#pragma once + +#include +#include + +#include "arrow/util/concurrent_map.h" +#include "arrow/util/mutex.h" + +namespace parquet::encryption { + +using ::arrow::util::ConcurrentMap; + +namespace internal { + +using TimePoint = + std::chrono::time_point>; + +inline TimePoint CurrentTimePoint() { return std::chrono::system_clock::now(); } + +template +class ExpiringCacheEntry { + public: + ExpiringCacheEntry() = default; + + ExpiringCacheEntry(E cached_item, double expiration_interval_seconds) + : expiration_timestamp_(CurrentTimePoint() + + std::chrono::duration(expiration_interval_seconds)), + cached_item_(std::move(cached_item)) {} + + bool IsExpired() const { + const auto now = CurrentTimePoint(); + return (now > expiration_timestamp_); + } + + E cached_item() { return cached_item_; } + + private: + const TimePoint expiration_timestamp_; + E cached_item_; +}; + +// This class is to avoid the below warning when compiling KeyToolkit class with VS2015 +// warning C4503: decorated name length exceeded, name was truncated +template +class ExpiringCacheMapEntry { + public: + ExpiringCacheMapEntry() = default; + + explicit ExpiringCacheMapEntry( + std::shared_ptr> cached_item, + double expiration_interval_seconds) + : map_cache_(cached_item, expiration_interval_seconds) {} + + bool IsExpired() { return map_cache_.IsExpired(); } + + std::shared_ptr> cached_item() { + return map_cache_.cached_item(); + } + + private: + // ConcurrentMap object may be accessed and modified at many places at the same time, + // from multiple threads, or even removed from cache. + ExpiringCacheEntry>> map_cache_; +}; + +} // namespace internal + +// Two-level cache with expiration of internal caches according to token lifetime. +// External cache is per token, internal is per string key. +// Wrapper class around: +// std::unordered_map>> +// This cache is safe to be shared between threads. +template +class TwoLevelCacheWithExpiration { + public: + TwoLevelCacheWithExpiration() { + last_cache_cleanup_timestamp_ = internal::CurrentTimePoint(); + } + + std::shared_ptr> GetOrCreateInternalCache( + const std::string& access_token, double cache_entry_lifetime_seconds) { + auto lock = mutex_.Lock(); + + auto external_cache_entry = cache_.find(access_token); + if (external_cache_entry == cache_.end() || + external_cache_entry->second.IsExpired()) { + cache_.insert({access_token, internal::ExpiringCacheMapEntry( + std::make_shared>(), + cache_entry_lifetime_seconds)}); + } + + return cache_[access_token].cached_item(); + } + + void CheckCacheForExpiredTokens(double cache_cleanup_period_seconds = 0.0) { + auto lock = mutex_.Lock(); + + const auto now = internal::CurrentTimePoint(); + if (now > (last_cache_cleanup_timestamp_ + + std::chrono::duration(cache_cleanup_period_seconds))) { + RemoveExpiredEntriesNoMutex(); + last_cache_cleanup_timestamp_ = now; + } + } + + void Remove(const std::string& access_token) { + auto lock = mutex_.Lock(); + cache_.erase(access_token); + } + + void Clear() { + auto lock = mutex_.Lock(); + cache_.clear(); + } + + private: + void RemoveExpiredEntriesNoMutex() { + for (auto it = cache_.begin(); it != cache_.end();) { + if (it->second.IsExpired()) { + it = cache_.erase(it); + } else { + ++it; + } + } + } + std::unordered_map> cache_; + internal::TimePoint last_cache_cleanup_timestamp_; + ::arrow::util::Mutex mutex_; +}; + +} // namespace parquet::encryption diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/type_fwd.h b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/type_fwd.h new file mode 100644 index 0000000000000000000000000000000000000000..623811718482c591e708a297dff9eb35ae0c85a9 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/encryption/type_fwd.h @@ -0,0 +1,28 @@ +// 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. + +#pragma once + +namespace parquet { + +class Decryptor; +class Encryptor; + +class InternalFileDecryptor; +class InternalFileEncryptor; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/file_writer.h b/lib/python3.12/site-packages/pyarrow/include/parquet/file_writer.h new file mode 100644 index 0000000000000000000000000000000000000000..d5ea1d7c98a0ef8509b4821ab111007d7601996b --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/file_writer.h @@ -0,0 +1,245 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "parquet/metadata.h" +#include "parquet/platform.h" +#include "parquet/properties.h" +#include "parquet/schema.h" + +namespace parquet { + +class ColumnWriter; + +// FIXME: copied from reader-internal.cc +static constexpr uint8_t kParquetMagic[4] = {'P', 'A', 'R', '1'}; +static constexpr uint8_t kParquetEMagic[4] = {'P', 'A', 'R', 'E'}; + +class PARQUET_EXPORT RowGroupWriter { + public: + // Forward declare a virtual class 'Contents' to aid dependency injection and more + // easily create test fixtures + // An implementation of the Contents class is defined in the .cc file + struct Contents { + virtual ~Contents() = default; + virtual int num_columns() const = 0; + virtual int64_t num_rows() const = 0; + + // to be used only with ParquetFileWriter::AppendRowGroup + virtual ColumnWriter* NextColumn() = 0; + // to be used only with ParquetFileWriter::AppendBufferedRowGroup + virtual ColumnWriter* column(int i) = 0; + + virtual int current_column() const = 0; + virtual void Close() = 0; + + /// \brief total uncompressed bytes written by the page writer + virtual int64_t total_bytes_written() const = 0; + /// \brief total bytes still compressed but not written by the page writer + virtual int64_t total_compressed_bytes() const = 0; + /// \brief total compressed bytes written by the page writer + virtual int64_t total_compressed_bytes_written() const = 0; + + virtual bool buffered() const = 0; + }; + + explicit RowGroupWriter(std::unique_ptr contents); + + /// Construct a ColumnWriter for the indicated row group-relative column. + /// + /// To be used only with ParquetFileWriter::AppendRowGroup + /// Ownership is solely within the RowGroupWriter. The ColumnWriter is only + /// valid until the next call to NextColumn or Close. As the contents are + /// directly written to the sink, once a new column is started, the contents + /// of the previous one cannot be modified anymore. + ColumnWriter* NextColumn(); + /// Index of currently written column. Equal to -1 if NextColumn() + /// has not been called yet. + int current_column(); + void Close(); + + int num_columns() const; + + /// Construct a ColumnWriter for the indicated row group column. + /// + /// To be used only with ParquetFileWriter::AppendBufferedRowGroup + /// Ownership is solely within the RowGroupWriter. The ColumnWriter is + /// valid until Close. The contents are buffered in memory and written to sink + /// on Close + ColumnWriter* column(int i); + + /** + * Number of rows that shall be written as part of this RowGroup. + */ + int64_t num_rows() const; + + /// \brief total uncompressed bytes written by the page writer + int64_t total_bytes_written() const; + /// \brief total bytes still compressed but not written by the page writer. + /// It will always return 0 from the SerializedPageWriter. + int64_t total_compressed_bytes() const; + /// \brief total compressed bytes written by the page writer + int64_t total_compressed_bytes_written() const; + + /// Returns whether the current RowGroupWriter is in the buffered mode and is created + /// by calling ParquetFileWriter::AppendBufferedRowGroup. + bool buffered() const; + + private: + // Holds a pointer to an instance of Contents implementation + std::unique_ptr contents_; +}; + +PARQUET_EXPORT +void WriteFileMetaData(const FileMetaData& file_metadata, + ::arrow::io::OutputStream* sink); + +PARQUET_EXPORT +void WriteMetaDataFile(const FileMetaData& file_metadata, + ::arrow::io::OutputStream* sink); + +PARQUET_EXPORT +void WriteEncryptedFileMetadata(const FileMetaData& file_metadata, + ArrowOutputStream* sink, + const std::shared_ptr& encryptor, + bool encrypt_footer); + +PARQUET_EXPORT +void WriteEncryptedFileMetadata(const FileMetaData& file_metadata, + ::arrow::io::OutputStream* sink, + const std::shared_ptr& encryptor = NULLPTR, + bool encrypt_footer = false); +PARQUET_EXPORT +void WriteFileCryptoMetaData(const FileCryptoMetaData& crypto_metadata, + ::arrow::io::OutputStream* sink); + +class PARQUET_EXPORT ParquetFileWriter { + public: + // Forward declare a virtual class 'Contents' to aid dependency injection and more + // easily create test fixtures + // An implementation of the Contents class is defined in the .cc file + struct Contents { + Contents(std::shared_ptr<::parquet::schema::GroupNode> schema, + std::shared_ptr key_value_metadata) + : schema_(), key_value_metadata_(std::move(key_value_metadata)) { + schema_.Init(std::move(schema)); + } + virtual ~Contents() {} + // Perform any cleanup associated with the file contents + virtual void Close() = 0; + + virtual RowGroupWriter* AppendRowGroup() = 0; + virtual RowGroupWriter* AppendBufferedRowGroup() = 0; + + virtual int64_t num_rows() const = 0; + virtual int num_columns() const = 0; + virtual int num_row_groups() const = 0; + + virtual const std::shared_ptr& properties() const = 0; + + const std::shared_ptr& key_value_metadata() const { + return key_value_metadata_; + } + + virtual void AddKeyValueMetadata( + const std::shared_ptr& key_value_metadata) = 0; + + // Return const-pointer to make it clear that this object is not to be copied + const SchemaDescriptor* schema() const { return &schema_; } + + SchemaDescriptor schema_; + + /// This should be the only place this is stored. Everything else is a const reference + std::shared_ptr key_value_metadata_; + + const std::shared_ptr& metadata() const { return file_metadata_; } + std::shared_ptr file_metadata_; + }; + + ParquetFileWriter(); + ~ParquetFileWriter(); + + static std::unique_ptr Open( + std::shared_ptr<::arrow::io::OutputStream> sink, + std::shared_ptr schema, + std::shared_ptr properties = default_writer_properties(), + std::shared_ptr key_value_metadata = NULLPTR); + + void Open(std::unique_ptr contents); + void Close(); + + /// Construct a RowGroupWriter with an arbitrary number of rows. + /// + /// Ownership is solely within the ParquetFileWriter. The RowGroupWriter is only valid + /// until the next call to AppendRowGroup or AppendBufferedRowGroup or Close. + RowGroupWriter* AppendRowGroup(); + + /// Construct a RowGroupWriter that buffers all the values until the RowGroup is ready. + /// Use this if you want to write a RowGroup based on a certain size + /// + /// Ownership is solely within the ParquetFileWriter. The RowGroupWriter is only valid + /// until the next call to AppendRowGroup or AppendBufferedRowGroup or Close. + RowGroupWriter* AppendBufferedRowGroup(); + + /// \brief Add key-value metadata to the file. + /// \param[in] key_value_metadata the metadata to add. + /// \note This will overwrite any existing metadata with the same key(s). + /// \throw ParquetException if Close() has been called. + void AddKeyValueMetadata( + const std::shared_ptr& key_value_metadata); + + /// Number of columns. + /// + /// This number is fixed during the lifetime of the writer as it is determined via + /// the schema. + int num_columns() const; + + /// Number of rows in the yet started RowGroups. + /// + /// Changes on the addition of a new RowGroup. + int64_t num_rows() const; + + /// Number of started RowGroups. + int num_row_groups() const; + + /// Configuration passed to the writer, e.g. the used Parquet format version. + const std::shared_ptr& properties() const; + + /// Returns the file schema descriptor + const SchemaDescriptor* schema() const; + + /// Returns a column descriptor in schema + const ColumnDescriptor* descr(int i) const; + + /// Returns the file custom metadata + const std::shared_ptr& key_value_metadata() const; + + /// Returns the file metadata, only available after calling Close(). + const std::shared_ptr metadata() const; + + private: + // Holds a pointer to an instance of Contents implementation + std::unique_ptr contents_; + std::shared_ptr file_metadata_; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/geospatial/statistics.h b/lib/python3.12/site-packages/pyarrow/include/parquet/geospatial/statistics.h new file mode 100644 index 0000000000000000000000000000000000000000..fb4c31af1b7d9ab45ef012d6b9dedaed05194bfc --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/geospatial/statistics.h @@ -0,0 +1,198 @@ +// 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. + +#pragma once + +#include +#include +#include + +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace parquet::geospatial { + +/// \brief The maximum number of dimensions represented by a geospatial type +/// (i.e., X, Y, Z, and M) +inline constexpr int kMaxDimensions = 4; + +/// \brief NaN, used to represent bounds for which predicate pushdown cannnot +/// be applied (e.g., because a writer did not provide bounds for a given dimension) +inline constexpr double kNaN = std::numeric_limits::quiet_NaN(); + +/// \brief Structure represented encoded statistics to be written to and read from Parquet +/// serialized metadata. +/// +/// See the Parquet Thrift definition and GeoStatistics for the specific definition +/// of field values. +struct PARQUET_EXPORT EncodedGeoStatistics { + bool xy_bounds_present{false}; + double xmin{kNaN}; + double xmax{kNaN}; + double ymin{kNaN}; + double ymax{kNaN}; + + bool z_bounds_present{false}; + double zmin{kNaN}; + double zmax{kNaN}; + + bool m_bounds_present{false}; + double mmin{kNaN}; + double mmax{kNaN}; + + bool geospatial_types_present() const { return !geospatial_types.empty(); } + std::vector geospatial_types; +}; + +class GeoStatisticsImpl; + +/// \brief Base type for computing geospatial column statistics while writing a file +/// or representing them when reading a file +/// +/// These statistics track the minimum and maximum value (omitting NaN values) of the +/// four possible dimensions (X, Y, Z, and M) and the distinct set of geometry +/// type/dimension combinations (e.g., point XY, linestring XYZM) present in the data. +/// Any of these individual components may be "invalid": for example, when reading a +/// Parquet file, information about individual components obtained from the column +/// chunk metadata may have been missing or deemed unusable. Orthogonally, +/// any of these individual components may be "empty": for example, when using +/// GeoStatistics to accumulate bounds whilst writing, if all geometries in a column chunk +/// are null, all ranges (X, Y, Z, and M) will be empty. If all geometries in a column +/// chunk contain only XY coordinates (the most common case), the Z and M ranges will +/// be empty but the X and Y ranges will contain finite bounds. Empty ranges are +/// considered "valid" because they are known to represent exactly zero values (in +/// contrast to an invalid range, whose contents is completely unknown). These concepts +/// are all necessary for this object to accurately represent (1) accumulated or partially +/// accumulated statistics during the writing process and (2) deserialized statistics read +/// from the column chunk metadata during the reading process. +/// +/// EXPERIMENTAL +class PARQUET_EXPORT GeoStatistics { + public: + GeoStatistics(); + explicit GeoStatistics(const EncodedGeoStatistics& encoded); + + ~GeoStatistics(); + + /// \brief Return true if bounds, geometry types, and validity are identical + bool Equals(const GeoStatistics& other) const; + + /// \brief Update these statistics based on previously calculated or decoded statistics + /// + /// Merging statistics with wraparound X values is not currently supported. Merging + /// two GeoStatistics where one or both has a wraparound X range will result in these + /// statistics having an X dimension marked as invalid. + void Merge(const GeoStatistics& other); + + /// \brief Update these statistics based on values + void Update(const ByteArray* values, int64_t num_values); + + /// \brief Update these statistics based on the non-null elements of values + void UpdateSpaced(const ByteArray* values, const uint8_t* valid_bits, + int64_t valid_bits_offset, int64_t num_spaced_values, + int64_t num_values); + + /// \brief Update these statistics based on the non-null elements of values + /// + /// Currently, BinaryArray and LargeBinaryArray input is supported. + void Update(const ::arrow::Array& values); + + /// \brief Return these statistics to an empty state + void Reset(); + + /// \brief Encode the statistics for serializing to Thrift + /// + /// If invalid WKB was encountered or if the statistics contain NaN + /// for any reason, Encode() will return nullopt to indicate that + /// statistics should not be written to thrift. + std::optional Encode() const; + + /// \brief Returns false if invalid WKB was encountered + bool is_valid() const; + + /// \brief Reset existing statistics and populate them from previously-encoded ones + void Decode(const EncodedGeoStatistics& encoded); + + /// \brief Minimum values in XYZM order + /// + /// For dimensions where dimension_valid() is false, the value will be NaN. For + /// dimensions where dimension_empty() is true, the value will be +Inf. + /// + /// For the first dimension (X) only, wraparound bounds apply where xmin > xmax. In this + /// case, these bounds represent the union of the intervals [xmax, Inf] and [-Inf, + /// xmin]. This implementation does not yet generate these types of bounds but they may + /// be encountered in statistics when reading a Parquet file. + std::array lower_bound() const; + + /// \brief Maximum values in XYZM order + /// + /// For dimensions where dimension_valid() is false, the value will be NaN. For + /// dimensions where dimension_empty() is true, the value will be -Inf. + /// + /// For the first dimension (X) only, wraparound bounds apply where xmin > xmax. In this + /// case, these bounds represent the union of the intervals [xmax, Inf] and [-Inf, + /// xmin]. This implementation does not yet generate these types of bounds but they may + /// be encountered in statistics when reading a Parquet file. + std::array upper_bound() const; + + /// \brief Dimension emptiness in XYZM order + /// + /// True for a given dimension if and only if zero non-NaN values were encountered + /// in that dimension and dimension_valid() is true for that dimension. + /// + /// When calculating statistics, zero or more of these values may be true because + /// this implementation calculates bounds for all dimensions; however, it may be + /// true that zero coordinates were encountered in a given dimension. For example, + /// dimension_empty() will return four true values if Update() was not called + /// or if Update() was called with only null values. If Update() was provided + /// one or more geometries with X and Y dimensions but not Z or M dimensions, + /// dimension_empty() will return true, true, false, false. + /// + /// For statistics read from a Parquet file, dimension_empty() will always contain + /// false values because there is no mechanism to communicate an empty interval + /// in the Thrift metadata. + std::array dimension_empty() const; + + /// \brief Dimension validity (i.e. presence) in XYZM order + /// + /// When calculating statistics, this will always be true because this implementation + /// calculates statistics for all dimensions. When reading a Parquet file, one or more + /// of these values may be false because the file may not have provided bounds for all + /// dimensions. + /// + /// See documentation for dimension_empty(), lower_bound(), and/or upper_bound() for the + /// canonical values of those outputs for the dimensions where dimension_valid() is + /// false. + std::array dimension_valid() const; + + /// \brief Return the geometry type codes + /// + /// This implementation always returns sorted output with no duplicates. When + /// calculating statistics, a value will always be returned (although the returned + /// vector may be empty if Update() was never called or was only called with null + /// values). When reading a Parquet file, std::nullopt may be returned because + /// the file may not have provided this information. + std::optional> geometry_types() const; + + /// \brief Return a string representation of these statistics + std::string ToString() const; + + private: + std::unique_ptr impl_; +}; + +} // namespace parquet::geospatial diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison.h b/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison.h new file mode 100644 index 0000000000000000000000000000000000000000..3ae442dd46e57b7f86b405d9502442d3195719e8 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison.h @@ -0,0 +1,38 @@ +// 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. +#pragma once + +#include +#include + +#include "parquet/platform.h" + +namespace parquet::internal { + +/// Builds a bitmap where each set bit indicates the corresponding level is greater +/// than rhs. +uint64_t PARQUET_EXPORT GreaterThanBitmap(const int16_t* levels, int64_t num_levels, + int16_t rhs); + +struct MinMax { + int16_t min; + int16_t max; +}; + +MinMax FindMinMax(const int16_t* levels, int64_t num_levels); + +} // namespace parquet::internal diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison_inc.h b/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison_inc.h new file mode 100644 index 0000000000000000000000000000000000000000..04f628d53311166233bd24158de07a7a12003f61 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/level_comparison_inc.h @@ -0,0 +1,61 @@ +// 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. +#pragma once + +#include "arrow/util/bit_util.h" +#include "arrow/util/endian.h" +#include "parquet/level_comparison.h" + +// Used to make sure ODR rule isn't violated. +#ifndef PARQUET_IMPL_NAMESPACE +# error "PARQUET_IMPL_NAMESPACE must be defined" +#endif +namespace parquet::internal::PARQUET_IMPL_NAMESPACE { +/// Builds a bitmap by applying predicate to the level vector provided. +/// +/// \param[in] levels Rep or def level array. +/// \param[in] num_levels The number of levels to process (must be [0, 64]) +/// \param[in] predicate The predicate to apply (must have the signature `bool +/// predicate(int16_t)`. +/// \returns The bitmap using least significant "bit" ordering. +/// +template +inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, + Predicate predicate) { + // Both clang and GCC can vectorize this automatically with SSE4/AVX2. + uint64_t mask = 0; + for (int x = 0; x < num_levels; x++) { + mask |= static_cast(predicate(levels[x]) ? 1 : 0) << x; + } + return ::arrow::bit_util::ToLittleEndian(mask); +} + +inline MinMax FindMinMaxImpl(const int16_t* levels, int64_t num_levels) { + MinMax out{std::numeric_limits::max(), std::numeric_limits::min()}; + for (int x = 0; x < num_levels; x++) { + out.min = std::min(levels[x], out.min); + out.max = std::max(levels[x], out.max); + } + return out; +} + +inline uint64_t GreaterThanBitmapImpl(const int16_t* levels, int64_t num_levels, + int16_t rhs) { + return LevelsToBitmap(levels, num_levels, [rhs](int16_t value) { return value > rhs; }); +} + +} // namespace parquet::internal::PARQUET_IMPL_NAMESPACE diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion.h b/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion.h new file mode 100644 index 0000000000000000000000000000000000000000..31de95be41c473814c52cd9d2f5902d63f1b944b --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion.h @@ -0,0 +1,216 @@ +// 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. + +#pragma once + +#include + +#include "arrow/util/endian.h" +#include "parquet/platform.h" +#include "parquet/schema.h" + +namespace parquet::internal { + +struct PARQUET_EXPORT LevelInfo { + LevelInfo() + : null_slot_usage(1), def_level(0), rep_level(0), repeated_ancestor_def_level(0) {} + LevelInfo(int32_t null_slots, int32_t definition_level, int32_t repetition_level, + int32_t repeated_ancestor_definition_level) + : null_slot_usage(null_slots), + def_level(static_cast(definition_level)), + rep_level(static_cast(repetition_level)), + repeated_ancestor_def_level( + static_cast(repeated_ancestor_definition_level)) {} + + bool operator==(const LevelInfo& b) const { + return null_slot_usage == b.null_slot_usage && def_level == b.def_level && + rep_level == b.rep_level && + repeated_ancestor_def_level == b.repeated_ancestor_def_level; + } + + bool HasNullableValues() const { return repeated_ancestor_def_level < def_level; } + + // How many slots an undefined but present (i.e. null) element in + // parquet consumes when decoding to Arrow. + // "Slot" is used in the same context as the Arrow specification + // (i.e. a value holder). + // This is only ever >1 for descendents of FixedSizeList. + int32_t null_slot_usage = 1; + + // The definition level at which the value for the field + // is considered not null (definition levels greater than + // or equal to this value indicate a not-null + // value for the field). For list fields definition levels + // greater than or equal to this field indicate a present, + // possibly null, child value. + int16_t def_level = 0; + + // The repetition level corresponding to this element + // or the closest repeated ancestor. Any repetition + // level less than this indicates either a new list OR + // an empty list (which is determined in conjunction + // with definition levels). + int16_t rep_level = 0; + + // The definition level indicating the level at which the closest + // repeated ancestor is not empty. This is used to discriminate + // between a value less than |def_level| being null or excluded entirely. + // For instance if we have an arrow schema like: + // list(struct(f0: int)). Then then there are the following + // definition levels: + // 0 = null list + // 1 = present but empty list. + // 2 = a null value in the list + // 3 = a non null struct but null integer. + // 4 = a present integer. + // When reconstructing, the struct and integer arrays' + // repeated_ancestor_def_level would be 2. Any + // def_level < 2 indicates that there isn't a corresponding + // child value in the list. + // i.e. [null, [], [null], [{f0: null}], [{f0: 1}]] + // has the def levels [0, 1, 2, 3, 4]. The actual + // struct array is only of length 3: [not-set, set, set] and + // the int array is also of length 3: [N/A, null, 1]. + // + int16_t repeated_ancestor_def_level = 0; + + /// Increments levels according to the cardinality of node. + void Increment(const schema::Node& node) { + if (node.is_repeated()) { + IncrementRepeated(); + return; + } + if (node.is_optional()) { + IncrementOptional(); + return; + } + } + + /// Increments level for a optional node. + void IncrementOptional() { def_level++; } + + /// Increments levels for the repeated node. Returns + /// the previous ancestor_list_def_level. + int16_t IncrementRepeated() { + int16_t last_repeated_ancestor = repeated_ancestor_def_level; + + // Repeated fields add both a repetition and definition level. This is used + // to distinguish between an empty list and a list with an item in it. + ++rep_level; + ++def_level; + // For levels >= repeated_ancestor_def_level it indicates the list was + // non-null and had at least one element. This is important + // for later decoding because we need to add a slot for these + // values. for levels < current_def_level no slots are added + // to arrays. + repeated_ancestor_def_level = def_level; + return last_repeated_ancestor; + } + + // Calculates and returns LevelInfo for a column descriptor. + static LevelInfo ComputeLevelInfo(const ColumnDescriptor* descr) { + LevelInfo level_info; + level_info.def_level = descr->max_definition_level(); + level_info.rep_level = descr->max_repetition_level(); + + int16_t min_spaced_def_level = descr->max_definition_level(); + const ::parquet::schema::Node* node = descr->schema_node().get(); + while (node && !node->is_repeated()) { + if (node->is_optional()) { + min_spaced_def_level--; + } + node = node->parent(); + } + level_info.repeated_ancestor_def_level = min_spaced_def_level; + return level_info; + } + + friend std::ostream& operator<<(std::ostream& os, const LevelInfo& levels) { + // This print method is to silence valgrind issues. What's printed + // is not important because all asserts happen directly on + // members. + os << "{def=" << levels.def_level << ", rep=" << levels.rep_level + << ", repeated_ancestor_def=" << levels.repeated_ancestor_def_level; + if (levels.null_slot_usage > 1) { + os << ", null_slot_usage=" << levels.null_slot_usage; + } + os << "}"; + return os; + } +}; + +// Input/Output structure for reconstructed validity bitmaps. +struct PARQUET_EXPORT ValidityBitmapInputOutput { + // Input only. + // The maximum number of values_read expected (actual + // values read must be less than or equal to this value). + // If this number is exceeded methods will throw a + // ParquetException. Exceeding this limit indicates + // either a corrupt or incorrectly written file. + int64_t values_read_upper_bound = 0; + // Output only. The number of values added to the encountered + // (this is logically the count of the number of elements + // for an Arrow array). + int64_t values_read = 0; + // Input/Output. The number of nulls encountered. + int64_t null_count = 0; + // Output only. The validity bitmap to populate. Maybe be null only + // for DefRepLevelsToListInfo (if all that is needed is list offsets). + uint8_t* valid_bits = NULLPTR; + // Input only, offset into valid_bits to start at. + int64_t valid_bits_offset = 0; +}; + +// Converts def_levels to validity bitmaps for non-list arrays and structs that have +// at least one member that is not a list and has no list descendents. +// For lists use DefRepLevelsToList and structs where all descendants contain +// a list use DefRepLevelsToBitmap. +void PARQUET_EXPORT DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output); + +// Reconstructs a validity bitmap and list offsets for a list arrays based on +// def/rep levels. The first element of offsets will not be modified if rep_levels +// starts with a new list. The first element of offsets will be used when calculating +// the next offset. See documentation onf DefLevelsToBitmap for when to use this +// method vs the other ones in this file for reconstruction. +// +// Offsets must be sized to 1 + values_read_upper_bound. +void PARQUET_EXPORT DefRepLevelsToList(const int16_t* def_levels, + const int16_t* rep_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output, + int32_t* offsets); +void PARQUET_EXPORT DefRepLevelsToList(const int16_t* def_levels, + const int16_t* rep_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output, + int64_t* offsets); + +// Reconstructs a validity bitmap for a struct every member is a list or has +// a list descendant. See documentation on DefLevelsToBitmap for when more +// details on this method compared to the other ones defined above. +void PARQUET_EXPORT DefRepLevelsToBitmap(const int16_t* def_levels, + const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output); + +// This is exposed to ensure we can properly test a software simulated pext function +// (i.e. it isn't hidden by runtime dispatch). +uint64_t PARQUET_EXPORT TestOnlyExtractBitsSoftware(uint64_t bitmap, uint64_t selection); + +} // namespace parquet::internal diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion_inc.h b/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion_inc.h new file mode 100644 index 0000000000000000000000000000000000000000..335f5b92154b389babb15ac644b6d3b51076fcaa --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/level_conversion_inc.h @@ -0,0 +1,355 @@ +// 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. +#pragma once + +#include "parquet/level_conversion.h" + +#include +#include +#include + +#include "arrow/util/bit_run_reader.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/bitmap_writer.h" +#include "arrow/util/logging.h" +#include "arrow/util/simd.h" +#include "parquet/exception.h" +#include "parquet/level_comparison.h" + +#ifndef PARQUET_IMPL_NAMESPACE +# error "PARQUET_IMPL_NAMESPACE must be defined" +#endif + +namespace parquet::internal::PARQUET_IMPL_NAMESPACE { + +// clang-format off +/* Python code to generate lookup table: + +kLookupBits = 5 +count = 0 +print('constexpr int kLookupBits = {};'.format(kLookupBits)) +print('constexpr uint8_t kPextTable[1 << kLookupBits][1 << kLookupBits] = {') +print(' ', end = '') +for mask in range(1 << kLookupBits): + for data in range(1 << kLookupBits): + bit_value = 0 + bit_len = 0 + for i in range(kLookupBits): + if mask & (1 << i): + bit_value |= (((data >> i) & 1) << bit_len) + bit_len += 1 + out = '0x{:02X},'.format(bit_value) + count += 1 + if count % (1 << kLookupBits) == 1: + print(' {') + if count % 8 == 1: + print(' ', end = '') + if count % 8 == 0: + print(out, end = '\n') + else: + print(out, end = ' ') + if count % (1 << kLookupBits) == 0: + print(' },', end = '') +print('\n};') + +*/ +// clang-format on + +constexpr int kLookupBits = 5; +constexpr uint8_t kPextTable[1 << kLookupBits][1 << kLookupBits] = { + { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, + 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, + 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, + 0x01, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, + 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, + 0x03, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, + 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, + 0x00, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, + 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x00, 0x01, 0x00, + 0x01, 0x02, 0x03, 0x02, 0x03, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, + 0x02, 0x03, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x00, 0x00, 0x01, + 0x01, 0x02, 0x02, 0x03, 0x03, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, + 0x03, 0x03, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, + 0x03, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, + 0x06, 0x07, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, + 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, + 0x03, 0x02, 0x03, 0x02, 0x03, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, + 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, + 0x03, 0x02, 0x02, 0x03, 0x03, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, + 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x02, 0x02, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, + 0x07, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, + 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x04, 0x05, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, + 0x02, 0x03, 0x03, 0x03, 0x03, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, + 0x01, 0x01, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x04, 0x05, 0x04, + 0x05, 0x06, 0x07, 0x06, 0x07, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, + 0x02, 0x03, 0x04, 0x05, 0x04, 0x05, 0x06, 0x07, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x04, 0x04, 0x05, + 0x05, 0x06, 0x06, 0x07, 0x07, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, + 0x03, 0x03, 0x04, 0x04, 0x05, 0x05, 0x06, 0x06, 0x07, 0x07, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, + 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, + 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, + 0x01, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, + 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, 0x02, 0x03, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, + 0x01, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x02, 0x02, + 0x03, 0x03, 0x02, 0x02, 0x03, 0x03, 0x02, 0x02, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, + 0x03, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x04, 0x05, + 0x06, 0x07, 0x04, 0x05, 0x06, 0x07, 0x04, 0x05, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, + 0x00, 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, + 0x03, 0x03, 0x02, 0x02, 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x00, 0x01, 0x00, + 0x01, 0x02, 0x03, 0x02, 0x03, 0x04, 0x05, 0x04, 0x05, 0x06, 0x07, + 0x06, 0x07, 0x04, 0x05, 0x04, 0x05, 0x06, 0x07, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x00, 0x00, 0x01, + 0x01, 0x02, 0x02, 0x03, 0x03, 0x04, 0x04, 0x05, 0x05, 0x06, 0x06, + 0x07, 0x07, 0x04, 0x04, 0x05, 0x05, 0x06, 0x06, 0x07, 0x07, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x00, 0x01, 0x02, + 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, + 0x0E, 0x0F, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, + 0x01, 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, 0x02, 0x02, 0x02, + 0x02, 0x02, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, 0x03, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, + 0x03, 0x02, 0x03, 0x02, 0x03, 0x04, 0x05, 0x04, 0x05, 0x04, 0x05, + 0x04, 0x05, 0x06, 0x07, 0x06, 0x07, 0x06, 0x07, 0x06, 0x07, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, + 0x03, 0x02, 0x02, 0x03, 0x03, 0x04, 0x04, 0x05, 0x05, 0x04, 0x04, + 0x05, 0x05, 0x06, 0x06, 0x07, 0x07, 0x06, 0x06, 0x07, 0x07, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, + 0x07, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x08, 0x09, + 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x0C, 0x0D, 0x0E, 0x0F, + }, + { + 0x00, 0x00, 0x00, 0x00, 0x01, 0x01, 0x01, 0x01, 0x02, 0x02, 0x02, + 0x02, 0x03, 0x03, 0x03, 0x03, 0x04, 0x04, 0x04, 0x04, 0x05, 0x05, + 0x05, 0x05, 0x06, 0x06, 0x06, 0x06, 0x07, 0x07, 0x07, 0x07, + }, + { + 0x00, 0x01, 0x00, 0x01, 0x02, 0x03, 0x02, 0x03, 0x04, 0x05, 0x04, + 0x05, 0x06, 0x07, 0x06, 0x07, 0x08, 0x09, 0x08, 0x09, 0x0A, 0x0B, + 0x0A, 0x0B, 0x0C, 0x0D, 0x0C, 0x0D, 0x0E, 0x0F, 0x0E, 0x0F, + }, + { + 0x00, 0x00, 0x01, 0x01, 0x02, 0x02, 0x03, 0x03, 0x04, 0x04, 0x05, + 0x05, 0x06, 0x06, 0x07, 0x07, 0x08, 0x08, 0x09, 0x09, 0x0A, 0x0A, + 0x0B, 0x0B, 0x0C, 0x0C, 0x0D, 0x0D, 0x0E, 0x0E, 0x0F, 0x0F, + }, + { + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, + 0x0B, 0x0C, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, + 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F, + }, +}; + +inline uint64_t ExtractBitsSoftware(uint64_t bitmap, uint64_t select_bitmap) { + // A software emulation of _pext_u64 + + // These checks should be inline and are likely to be common cases. + if (select_bitmap == ~uint64_t{0}) { + return bitmap; + } else if (select_bitmap == 0) { + return 0; + } + + // Fallback to lookup table method + uint64_t bit_value = 0; + int bit_len = 0; + constexpr uint8_t kLookupMask = (1U << kLookupBits) - 1; + while (select_bitmap != 0) { + const auto mask_len = ARROW_POPCOUNT32(select_bitmap & kLookupMask); + const uint64_t value = kPextTable[select_bitmap & kLookupMask][bitmap & kLookupMask]; + bit_value |= (value << bit_len); + bit_len += mask_len; + bitmap >>= kLookupBits; + select_bitmap >>= kLookupBits; + } + return bit_value; +} + +#ifdef ARROW_HAVE_BMI2 + +// Use _pext_u64 on 64-bit builds, _pext_u32 on 32-bit builds, +# if UINTPTR_MAX == 0xFFFFFFFF + +using extract_bitmap_t = uint32_t; +inline extract_bitmap_t ExtractBits(extract_bitmap_t bitmap, + extract_bitmap_t select_bitmap) { + return _pext_u32(bitmap, select_bitmap); +} + +# else + +using extract_bitmap_t = uint64_t; +inline extract_bitmap_t ExtractBits(extract_bitmap_t bitmap, + extract_bitmap_t select_bitmap) { + return _pext_u64(bitmap, select_bitmap); +} + +# endif + +#else // !defined(ARROW_HAVE_BMI2) + +// Use 64-bit pext emulation when BMI2 isn't available. +using extract_bitmap_t = uint64_t; +inline extract_bitmap_t ExtractBits(extract_bitmap_t bitmap, + extract_bitmap_t select_bitmap) { + return ExtractBitsSoftware(bitmap, select_bitmap); +} + +#endif + +static constexpr int64_t kExtractBitsSize = 8 * sizeof(extract_bitmap_t); + +template +int64_t DefLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, + int64_t upper_bound_remaining, LevelInfo level_info, + ::arrow::internal::FirstTimeBitmapWriter* writer) { + ARROW_DCHECK_LE(batch_size, kExtractBitsSize); + + // Greater than level_info.def_level - 1 implies >= the def_level + auto defined_bitmap = static_cast(::arrow::bit_util::FromLittleEndian( + internal::GreaterThanBitmap(def_levels, batch_size, level_info.def_level - 1))); + + if (has_repeated_parent) { + // Greater than level_info.repeated_ancestor_def_level - 1 implies >= the + // repeated_ancestor_def_level + auto present_bitmap = static_cast( + ::arrow::bit_util::FromLittleEndian(internal::GreaterThanBitmap( + def_levels, batch_size, level_info.repeated_ancestor_def_level - 1))); + auto selected_bits = ExtractBits(defined_bitmap, present_bitmap); + int64_t selected_count = ::arrow::bit_util::PopCount(present_bitmap); + if (ARROW_PREDICT_FALSE(selected_count > upper_bound_remaining)) { + throw ParquetException("Values read exceeded upper bound"); + } + writer->AppendWord(selected_bits, selected_count); + return ::arrow::bit_util::PopCount(selected_bits); + } else { + if (ARROW_PREDICT_FALSE(batch_size > upper_bound_remaining)) { + std::stringstream ss; + ss << "Values read exceeded upper bound"; + throw ParquetException(ss.str()); + } + + writer->AppendWord(defined_bitmap, batch_size); + return ::arrow::bit_util::PopCount(defined_bitmap); + } +} + +template +void DefLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, ValidityBitmapInputOutput* output) { + ::arrow::internal::FirstTimeBitmapWriter writer( + output->valid_bits, + /*start_offset=*/output->valid_bits_offset, + /*length=*/output->values_read_upper_bound); + int64_t set_count = 0; + output->values_read = 0; + int64_t values_read_remaining = output->values_read_upper_bound; + while (num_def_levels > kExtractBitsSize) { + set_count += DefLevelsBatchToBitmap( + def_levels, kExtractBitsSize, values_read_remaining, level_info, &writer); + def_levels += kExtractBitsSize; + num_def_levels -= kExtractBitsSize; + values_read_remaining = output->values_read_upper_bound - writer.position(); + } + set_count += DefLevelsBatchToBitmap( + def_levels, num_def_levels, values_read_remaining, level_info, &writer); + + output->values_read = writer.position(); + output->null_count += output->values_read - set_count; + writer.Finish(); +} + +} // namespace parquet::internal::PARQUET_IMPL_NAMESPACE diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/metadata.h b/lib/python3.12/site-packages/pyarrow/include/parquet/metadata.h new file mode 100644 index 0000000000000000000000000000000000000000..3380adbf56aa7d3bbdca9410bc1ae9ed74c0bb4d --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/metadata.h @@ -0,0 +1,554 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "parquet/encryption/type_fwd.h" +#include "parquet/platform.h" +#include "parquet/properties.h" +#include "parquet/type_fwd.h" + +namespace parquet { + +using KeyValueMetadata = ::arrow::KeyValueMetadata; + +class PARQUET_EXPORT ApplicationVersion { + public: + // Known Versions with Issues + static const ApplicationVersion& PARQUET_251_FIXED_VERSION(); + static const ApplicationVersion& PARQUET_816_FIXED_VERSION(); + static const ApplicationVersion& PARQUET_CPP_FIXED_STATS_VERSION(); + static const ApplicationVersion& PARQUET_MR_FIXED_STATS_VERSION(); + static const ApplicationVersion& PARQUET_CPP_10353_FIXED_VERSION(); + + // Application that wrote the file. e.g. "IMPALA" + std::string application_; + // Build name + std::string build_; + + // Version of the application that wrote the file, expressed as + // (..). Unmatched parts default to 0. + // "1.2.3" => {1, 2, 3} + // "1.2" => {1, 2, 0} + // "1.2-cdh5" => {1, 2, 0} + struct { + int major; + int minor; + int patch; + std::string unknown; + std::string pre_release; + std::string build_info; + } version; + + ApplicationVersion() = default; + explicit ApplicationVersion(const std::string& created_by); + ApplicationVersion(std::string application, int major, int minor, int patch); + + // Returns true if version is strictly less than other_version + bool VersionLt(const ApplicationVersion& other_version) const; + + // Returns true if version is strictly equal with other_version + bool VersionEq(const ApplicationVersion& other_version) const; + + // Checks if the Version has the correct statistics for a given column + bool HasCorrectStatistics(Type::type primitive, const EncodedStatistics& statistics, + SortOrder::type sort_order = SortOrder::SIGNED) const; +}; + +class PARQUET_EXPORT ColumnCryptoMetaData { + public: + static std::unique_ptr Make(const uint8_t* metadata); + ~ColumnCryptoMetaData(); + + bool Equals(const ColumnCryptoMetaData& other) const; + + std::shared_ptr path_in_schema() const; + bool encrypted_with_footer_key() const; + const std::string& key_metadata() const; + + private: + explicit ColumnCryptoMetaData(const uint8_t* metadata); + + class ColumnCryptoMetaDataImpl; + std::unique_ptr impl_; +}; + +/// \brief Public struct for Thrift PageEncodingStats in ColumnChunkMetaData +struct PageEncodingStats { + PageType::type page_type; + Encoding::type encoding; + int32_t count; +}; + +/// \brief Public struct for location to page index in ColumnChunkMetaData. +struct IndexLocation { + /// File offset of the given index, in bytes + int64_t offset; + /// Length of the given index, in bytes + int32_t length; +}; + +/// \brief ColumnChunkMetaData is a proxy around format::ColumnChunkMetaData. +class PARQUET_EXPORT ColumnChunkMetaData { + public: + // API convenience to get a MetaData accessor + static std::unique_ptr Make( + const void* metadata, const ColumnDescriptor* descr, + const ReaderProperties& properties = default_reader_properties(), + const ApplicationVersion* writer_version = NULLPTR, int16_t row_group_ordinal = -1, + int16_t column_ordinal = -1, + std::shared_ptr file_decryptor = NULLPTR); + + ~ColumnChunkMetaData(); + + bool Equals(const ColumnChunkMetaData& other) const; + + // Byte offset of `ColumnMetaData` in `file_path()`. + // + // Note that the meaning of this field has been inconsistent among implementations + // so its use has since been deprecated in the Parquet specification. Modern + // implementations will set this to `0` to indicate that the `ColumnMetaData` is solely + // contained in the `ColumnChunk` struct. + int64_t file_offset() const; + + // parameter is only used when a dataset is spread across multiple files + const std::string& file_path() const; + + // column metadata + bool is_metadata_set() const; + Type::type type() const; + int64_t num_values() const; + std::shared_ptr path_in_schema() const; + bool is_stats_set() const; + bool is_geo_stats_set() const; + std::shared_ptr statistics() const; + std::shared_ptr encoded_statistics() const; + std::shared_ptr size_statistics() const; + std::shared_ptr geo_statistics() const; + + Compression::type compression() const; + // Indicate if the ColumnChunk compression is supported by the current + // compiled parquet library. + bool can_decompress() const; + + const std::vector& encodings() const; + const std::vector& encoding_stats() const; + std::optional bloom_filter_offset() const; + std::optional bloom_filter_length() const; + bool has_dictionary_page() const; + int64_t dictionary_page_offset() const; + int64_t data_page_offset() const; + bool has_index_page() const; + int64_t index_page_offset() const; + int64_t total_compressed_size() const; + int64_t total_uncompressed_size() const; + std::unique_ptr crypto_metadata() const; + std::optional GetColumnIndexLocation() const; + std::optional GetOffsetIndexLocation() const; + const std::shared_ptr& key_value_metadata() const; + + private: + explicit ColumnChunkMetaData( + const void* metadata, const ColumnDescriptor* descr, int16_t row_group_ordinal, + int16_t column_ordinal, const ReaderProperties& properties, + const ApplicationVersion* writer_version = NULLPTR, + std::shared_ptr file_decryptor = NULLPTR); + // PIMPL Idiom + class ColumnChunkMetaDataImpl; + std::unique_ptr impl_; +}; + +/// \brief RowGroupMetaData is a proxy around format::RowGroupMetaData. +class PARQUET_EXPORT RowGroupMetaData { + public: + /// \brief Create a RowGroupMetaData from a serialized thrift message. + static std::unique_ptr Make( + const void* metadata, const SchemaDescriptor* schema, + const ReaderProperties& properties = default_reader_properties(), + const ApplicationVersion* writer_version = NULLPTR, + std::shared_ptr file_decryptor = NULLPTR); + + ~RowGroupMetaData(); + + bool Equals(const RowGroupMetaData& other) const; + + /// \brief The number of columns in this row group. The order must match the + /// parent's column ordering. + int num_columns() const; + + /// \brief Return the ColumnChunkMetaData of the corresponding column ordinal. + /// + /// WARNING, the returned object references memory location in it's parent + /// (RowGroupMetaData) object. Hence, the parent must outlive the returned + /// object. + /// + /// \param[in] index of the ColumnChunkMetaData to retrieve. + /// + /// \throws ParquetException if the index is out of bound. + std::unique_ptr ColumnChunk(int index) const; + + /// \brief Number of rows in this row group. + int64_t num_rows() const; + + /// \brief Total byte size of all the uncompressed column data in this row group. + int64_t total_byte_size() const; + + /// \brief Total byte size of all the compressed (and potentially encrypted) + /// column data in this row group. + /// + /// This information is optional and may be 0 if omitted. + int64_t total_compressed_size() const; + + /// \brief Byte offset from beginning of file to first page (data or + /// dictionary) in this row group + /// + /// The file_offset field that this method exposes is optional. This method + /// will return 0 if that field is not set to a meaningful value. + int64_t file_offset() const; + // Return const-pointer to make it clear that this object is not to be copied + const SchemaDescriptor* schema() const; + // Indicate if all of the RowGroup's ColumnChunks can be decompressed. + bool can_decompress() const; + // Sorting columns of the row group if any. + std::vector sorting_columns() const; + + private: + explicit RowGroupMetaData( + const void* metadata, const SchemaDescriptor* schema, + const ReaderProperties& properties, + const ApplicationVersion* writer_version = NULLPTR, + std::shared_ptr file_decryptor = NULLPTR); + // PIMPL Idiom + class RowGroupMetaDataImpl; + std::unique_ptr impl_; +}; + +class FileMetaDataBuilder; + +/// \brief FileMetaData is a proxy around format::FileMetaData. +class PARQUET_EXPORT FileMetaData { + public: + /// \brief Create a FileMetaData from a serialized thrift message. + static std::shared_ptr Make( + const void* serialized_metadata, uint32_t* inout_metadata_len, + const ReaderProperties& properties = default_reader_properties(), + std::shared_ptr file_decryptor = NULLPTR); + + ~FileMetaData(); + + bool Equals(const FileMetaData& other) const; + + /// \brief The number of parquet "leaf" columns. + /// + /// Parquet thrift definition requires that nested schema elements are + /// flattened. This method returns the number of columns in the flattened + /// version. + /// For instance, if the schema looks like this : + /// 0 foo.bar + /// foo.bar.baz 0 + /// foo.bar.baz2 1 + /// foo.qux 2 + /// 1 foo2 3 + /// 2 foo3 4 + /// This method will return 5, because there are 5 "leaf" fields (so 5 + /// flattened fields) + int num_columns() const; + + /// \brief The number of flattened schema elements. + /// + /// Parquet thrift definition requires that nested schema elements are + /// flattened. This method returns the total number of elements in the + /// flattened list. + int num_schema_elements() const; + + /// \brief The total number of rows. + /// + /// If the FileMetaData was obtained by calling `SubSet()`, this is the total + /// number of rows in the selected row groups. + int64_t num_rows() const; + + /// \brief The number of row groups in the file. + /// + /// If the FileMetaData was obtained by calling `SubSet()`, this is the number + /// of selected row groups. + int num_row_groups() const; + + /// \brief Return the RowGroupMetaData of the corresponding row group ordinal. + /// + /// WARNING, the returned object references memory location in it's parent + /// (FileMetaData) object. Hence, the parent must outlive the returned object. + /// + /// \param[in] index of the RowGroup to retrieve. + /// + /// \throws ParquetException if the index is out of bound. + std::unique_ptr RowGroup(int index) const; + + /// \brief Return the "version" of the file + /// + /// WARNING: The value returned by this method is unreliable as 1) the Parquet + /// file metadata stores the version as a single integer and 2) some producers + /// are known to always write a hardcoded value. Therefore, you cannot use + /// this value to know which features are used in the file. + ParquetVersion::type version() const; + + /// \brief Return the application's user-agent string of the writer. + const std::string& created_by() const; + + /// \brief Return the application's version of the writer. + const ApplicationVersion& writer_version() const; + + /// \brief Size of the original thrift encoded metadata footer. + uint32_t size() const; + + /// \brief Indicate if all of the FileMetaData's RowGroups can be decompressed. + /// + /// This will return false if any of the RowGroup's page is compressed with a + /// compression format which is not compiled in the current parquet library. + bool can_decompress() const; + + bool is_encryption_algorithm_set() const; + EncryptionAlgorithm encryption_algorithm() const; + const std::string& footer_signing_key_metadata() const; + + /// \brief Verify signature of FileMetaData when file is encrypted but footer + /// is not encrypted (plaintext footer). + bool VerifySignature(const void* signature); + + void WriteTo(::arrow::io::OutputStream* dst, + const std::shared_ptr& encryptor = NULLPTR) const; + + /// \brief Return Thrift-serialized representation of the metadata as a + /// string + std::string SerializeToString() const; + + // Return const-pointer to make it clear that this object is not to be copied + const SchemaDescriptor* schema() const; + + const std::shared_ptr& key_value_metadata() const; + + /// \brief Set a path to all ColumnChunk for all RowGroups. + /// + /// Commonly used by systems (Dask, Spark) who generates an metadata-only + /// parquet file. The path is usually relative to said index file. + /// + /// \param[in] path to set. + void set_file_path(const std::string& path); + + /// \brief Merge row groups from another metadata file into this one. + /// + /// The schema of the input FileMetaData must be equal to the + /// schema of this object. + /// + /// This is used by systems who creates an aggregate metadata-only file by + /// concatenating the row groups of multiple files. This newly created + /// metadata file acts as an index of all available row groups. + /// + /// \param[in] other FileMetaData to merge the row groups from. + /// + /// \throws ParquetException if schemas are not equal. + void AppendRowGroups(const FileMetaData& other); + + /// \brief Return a FileMetaData containing a subset of the row groups in this + /// FileMetaData. + std::shared_ptr Subset(const std::vector& row_groups) const; + + /// \brief Serialize metadata unencrypted as string + /// + /// \param[in] scrub whether to remove sensitive information from the metadata. + /// \param[in] debug whether to serialize the metadata as Thrift (if false) or + /// debug text (if true). + std::string SerializeUnencrypted(bool scrub, bool debug) const; + + private: + friend FileMetaDataBuilder; + friend class SerializedFile; + friend class SerializedRowGroup; + + explicit FileMetaData(const void* serialized_metadata, uint32_t* metadata_len, + const ReaderProperties& properties, + std::shared_ptr file_decryptor = NULLPTR); + + void set_file_decryptor(std::shared_ptr file_decryptor); + const std::shared_ptr& file_decryptor() const; + + // PIMPL Idiom + FileMetaData(); + class FileMetaDataImpl; + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT FileCryptoMetaData { + public: + // API convenience to get a MetaData accessor + static std::shared_ptr Make( + const uint8_t* serialized_metadata, uint32_t* metadata_len, + const ReaderProperties& properties = default_reader_properties()); + ~FileCryptoMetaData(); + + EncryptionAlgorithm encryption_algorithm() const; + const std::string& key_metadata() const; + + void WriteTo(::arrow::io::OutputStream* dst) const; + + private: + friend FileMetaDataBuilder; + FileCryptoMetaData(const uint8_t* serialized_metadata, uint32_t* metadata_len, + const ReaderProperties& properties); + + // PIMPL Idiom + FileCryptoMetaData(); + class FileCryptoMetaDataImpl; + std::unique_ptr impl_; +}; + +// Builder API +class PARQUET_EXPORT ColumnChunkMetaDataBuilder { + public: + // API convenience to get a MetaData reader + static std::unique_ptr Make( + std::shared_ptr props, const ColumnDescriptor* column); + + static std::unique_ptr Make( + std::shared_ptr props, const ColumnDescriptor* column, + void* contents); + + ~ColumnChunkMetaDataBuilder(); + + // column chunk + // Used when a dataset is spread across multiple files + void set_file_path(const std::string& path); + + // column metadata + void SetStatistics(const EncodedStatistics& stats); + void SetSizeStatistics(const SizeStatistics& size_stats); + + // column geometry statistics + void SetGeoStatistics(const geospatial::EncodedGeoStatistics& geo_stats); + + void SetKeyValueMetadata(std::shared_ptr key_value_metadata); + + // get the column descriptor + const ColumnDescriptor* descr() const; + + int64_t total_compressed_size() const; + // commit the metadata + + void Finish(int64_t num_values, int64_t dictionary_page_offset, + int64_t index_page_offset, int64_t data_page_offset, + int64_t compressed_size, int64_t uncompressed_size, bool has_dictionary, + bool dictionary_fallback, + const std::map& dict_encoding_stats_, + const std::map& data_encoding_stats_, + const std::shared_ptr& encryptor = NULLPTR); + + // The metadata contents, suitable for passing to ColumnChunkMetaData::Make + const void* contents() const; + + // For writing metadata at end of column chunk + void WriteTo(::arrow::io::OutputStream* sink); + + private: + explicit ColumnChunkMetaDataBuilder(std::shared_ptr props, + const ColumnDescriptor* column); + explicit ColumnChunkMetaDataBuilder(std::shared_ptr props, + const ColumnDescriptor* column, void* contents); + // PIMPL Idiom + class ColumnChunkMetaDataBuilderImpl; + std::unique_ptr impl_; +}; + +class PARQUET_EXPORT RowGroupMetaDataBuilder { + public: + // API convenience to get a MetaData reader + static std::unique_ptr Make( + std::shared_ptr props, const SchemaDescriptor* schema_, + void* contents); + + ~RowGroupMetaDataBuilder(); + + ColumnChunkMetaDataBuilder* NextColumnChunk(); + int num_columns(); + int64_t num_rows(); + int current_column() const; + + void set_num_rows(int64_t num_rows); + + // commit the metadata + void Finish(int64_t total_bytes_written, int16_t row_group_ordinal = -1); + + private: + explicit RowGroupMetaDataBuilder(std::shared_ptr props, + const SchemaDescriptor* schema_, void* contents); + // PIMPL Idiom + class RowGroupMetaDataBuilderImpl; + std::unique_ptr impl_; +}; + +/// \brief Public struct for location to all page indexes in a parquet file. +struct PageIndexLocation { + /// Alias type of page index location of a row group. The index location + /// is located by column ordinal. If the column does not have the page index, + /// its value is set to std::nullopt. + using RowGroupIndexLocation = std::vector>; + /// Alias type of page index location of a parquet file. The index location + /// is located by the row group ordinal. + using FileIndexLocation = std::map; + /// Row group column index locations which uses row group ordinal as the key. + FileIndexLocation column_index_location; + /// Row group offset index locations which uses row group ordinal as the key. + FileIndexLocation offset_index_location; +}; + +class PARQUET_EXPORT FileMetaDataBuilder { + public: + // API convenience to get a MetaData builder + static std::unique_ptr Make( + const SchemaDescriptor* schema, std::shared_ptr props); + + ~FileMetaDataBuilder(); + + // The prior RowGroupMetaDataBuilder (if any) is destroyed + RowGroupMetaDataBuilder* AppendRowGroup(); + + // Update location to all page indexes in the parquet file + void SetPageIndexLocation(const PageIndexLocation& location); + + // Complete the Thrift structure + std::unique_ptr Finish( + const std::shared_ptr& key_value_metadata = NULLPTR); + + // crypto metadata + std::unique_ptr GetCryptoMetaData(); + + private: + explicit FileMetaDataBuilder(const SchemaDescriptor* schema, + std::shared_ptr props); + // PIMPL Idiom + class FileMetaDataBuilderImpl; + std::unique_ptr impl_; +}; + +PARQUET_EXPORT std::string ParquetVersionToString(ParquetVersion::type ver); + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/page_index.h b/lib/python3.12/site-packages/pyarrow/include/parquet/page_index.h new file mode 100644 index 0000000000000000000000000000000000000000..3083159783ba793763f221c4b01c91ff6efef304 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/page_index.h @@ -0,0 +1,386 @@ +// 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. + +#pragma once + +#include "arrow/io/interfaces.h" +#include "parquet/encryption/type_fwd.h" +#include "parquet/type_fwd.h" +#include "parquet/types.h" + +#include +#include + +namespace parquet { + +/// \brief ColumnIndex is a proxy around format::ColumnIndex. +class PARQUET_EXPORT ColumnIndex { + public: + /// \brief Create a ColumnIndex from a serialized thrift message. + static std::unique_ptr Make(const ColumnDescriptor& descr, + const void* serialized_index, + uint32_t index_len, + const ReaderProperties& properties, + Decryptor* decryptor = NULLPTR); + + virtual ~ColumnIndex() = default; + + /// \brief A bitmap with a bit set for each data page that has only null values. + /// + /// The length of this vector is equal to the number of data pages in the column. + virtual const std::vector& null_pages() const = 0; + + /// \brief A vector of encoded lower bounds for each data page in this column. + /// + /// `null_pages` should be inspected first, as only pages with non-null values + /// may have their lower bounds populated. + virtual const std::vector& encoded_min_values() const = 0; + + /// \brief A vector of encoded upper bounds for each data page in this column. + /// + /// `null_pages` should be inspected first, as only pages with non-null values + /// may have their upper bounds populated. + virtual const std::vector& encoded_max_values() const = 0; + + /// \brief The ordering of lower and upper bounds. + /// + /// The boundary order applies across all lower bounds, and all upper bounds, + /// respectively. However, the order between lower bounds and upper bounds + /// cannot be derived from this. + virtual BoundaryOrder::type boundary_order() const = 0; + + /// \brief Whether per-page null count information is available. + virtual bool has_null_counts() const = 0; + + /// \brief An optional vector with the number of null values in each data page. + /// + /// `has_null_counts` should be called first to determine if this information is + /// available. + virtual const std::vector& null_counts() const = 0; + + /// \brief A vector of page indices for non-null pages. + virtual const std::vector& non_null_page_indices() const = 0; + + /// \brief Whether definition level histogram is available. + virtual bool has_definition_level_histograms() const = 0; + + /// \brief Whether repetition level histogram is available. + virtual bool has_repetition_level_histograms() const = 0; + + /// \brief List of definition level histograms for each page concatenated together. + virtual const std::vector& definition_level_histograms() const = 0; + + /// \brief List of repetition level histograms for each page concatenated together. + virtual const std::vector& repetition_level_histograms() const = 0; +}; + +/// \brief Typed implementation of ColumnIndex. +template +class PARQUET_EXPORT TypedColumnIndex : public ColumnIndex { + public: + using T = typename DType::c_type; + + /// \brief A vector of lower bounds for each data page in this column. + /// + /// This is like `encoded_min_values`, but with the values decoded according to + /// the column's physical type. + /// `min_values` and `max_values` can be used together with `boundary_order` + /// in order to prune some data pages when searching for specific values. + virtual const std::vector& min_values() const = 0; + + /// \brief A vector of upper bounds for each data page in this column. + /// + /// Just like `min_values`, but for upper bounds instead of lower bounds. + virtual const std::vector& max_values() const = 0; +}; + +using BoolColumnIndex = TypedColumnIndex; +using Int32ColumnIndex = TypedColumnIndex; +using Int64ColumnIndex = TypedColumnIndex; +using FloatColumnIndex = TypedColumnIndex; +using DoubleColumnIndex = TypedColumnIndex; +using ByteArrayColumnIndex = TypedColumnIndex; +using FLBAColumnIndex = TypedColumnIndex; + +/// \brief PageLocation is a proxy around format::PageLocation. +struct PARQUET_EXPORT PageLocation { + /// File offset of the data page. + int64_t offset; + /// Total compressed size of the data page and header. + int32_t compressed_page_size; + /// Row id of the first row in the page within the row group. + int64_t first_row_index; +}; + +/// \brief OffsetIndex is a proxy around format::OffsetIndex. +class PARQUET_EXPORT OffsetIndex { + public: + /// \brief Create a OffsetIndex from a serialized thrift message. + static std::unique_ptr Make(const void* serialized_index, + uint32_t index_len, + const ReaderProperties& properties, + Decryptor* decryptor = NULLPTR); + + virtual ~OffsetIndex() = default; + + /// \brief A vector of locations for each data page in this column. + virtual const std::vector& page_locations() const = 0; + + /// \brief A vector of unencoded/uncompressed size of each page for BYTE_ARRAY types, + /// or empty for other types. + virtual const std::vector& unencoded_byte_array_data_bytes() const = 0; +}; + +/// \brief Interface for reading the page index for a Parquet row group. +class PARQUET_EXPORT RowGroupPageIndexReader { + public: + virtual ~RowGroupPageIndexReader() = default; + + /// \brief Read column index of a column chunk. + /// + /// \param[in] i column ordinal of the column chunk. + /// \returns column index of the column or nullptr if it does not exist. + /// \throws ParquetException if the index is out of bound. + virtual std::shared_ptr GetColumnIndex(int32_t i) = 0; + + /// \brief Read offset index of a column chunk. + /// + /// \param[in] i column ordinal of the column chunk. + /// \returns offset index of the column or nullptr if it does not exist. + /// \throws ParquetException if the index is out of bound. + virtual std::shared_ptr GetOffsetIndex(int32_t i) = 0; +}; + +struct PageIndexSelection { + /// Specifies whether to read the column index. + bool column_index = false; + /// Specifies whether to read the offset index. + bool offset_index = false; +}; + +PARQUET_EXPORT +std::ostream& operator<<(std::ostream& out, const PageIndexSelection& params); + +struct RowGroupIndexReadRange { + /// Base start and total size of column index of all column chunks in a row group. + /// If none of the column chunks have column index, it is set to std::nullopt. + std::optional<::arrow::io::ReadRange> column_index = std::nullopt; + /// Base start and total size of offset index of all column chunks in a row group. + /// If none of the column chunks have offset index, it is set to std::nullopt. + std::optional<::arrow::io::ReadRange> offset_index = std::nullopt; +}; + +/// \brief Interface for reading the page index for a Parquet file. +class PARQUET_EXPORT PageIndexReader { + public: + virtual ~PageIndexReader() = default; + + /// \brief Create a PageIndexReader instance. + /// \returns a PageIndexReader instance. + /// WARNING: The returned PageIndexReader references to all the input parameters, so + /// it must not outlive all of the input parameters. Usually these input parameters + /// come from the same ParquetFileReader object, so it must not outlive the reader + /// that creates this PageIndexReader. + static std::shared_ptr Make( + ::arrow::io::RandomAccessFile* input, std::shared_ptr file_metadata, + const ReaderProperties& properties, + InternalFileDecryptor* file_decryptor = NULLPTR); + + /// \brief Get the page index reader of a specific row group. + /// \param[in] i row group ordinal to get page index reader. + /// \returns RowGroupPageIndexReader of the specified row group. A nullptr may or may + /// not be returned if the page index for the row group is unavailable. It is + /// the caller's responsibility to check the return value of follow-up calls + /// to the RowGroupPageIndexReader. + /// \throws ParquetException if the index is out of bound. + virtual std::shared_ptr RowGroup(int i) = 0; + + /// \brief Advise the reader which part of page index will be read later. + /// + /// The PageIndexReader can optionally prefetch and cache page index that + /// may be read later to get better performance. + /// + /// The contract of this function is as below: + /// 1) If WillNeed() has not been called for a specific row group and the page index + /// exists, follow-up calls to get column index or offset index of all columns in + /// this row group SHOULD NOT FAIL, but the performance may not be optimal. + /// 2) If WillNeed() has been called for a specific row group, follow-up calls to get + /// page index are limited to columns and index type requested by WillNeed(). + /// So it MAY FAIL if columns that are not requested by WillNeed() are requested. + /// 3) Later calls to WillNeed() MAY OVERRIDE previous calls of same row groups. + /// For example, + /// 1) If WillNeed() is not called for row group 0, then follow-up calls to read + /// column index and/or offset index of all columns of row group 0 should not + /// fail if its page index exists. + /// 2) If WillNeed() is called for columns 0 and 1 for row group 0, then follow-up + /// call to read page index of column 2 for row group 0 MAY FAIL even if its + /// page index exists. + /// 3) If WillNeed() is called for row group 0 with offset index only, then + /// follow-up call to read column index of row group 0 MAY FAIL even if + /// the column index of this column exists. + /// 4) If WillNeed() is called for columns 0 and 1 for row group 0, then later + /// call to WillNeed() for columns 1 and 2 for row group 0. The later one + /// overrides previous call and only columns 1 and 2 of row group 0 are allowed + /// to access. + /// + /// \param[in] row_group_indices list of row group ordinal to read page index later. + /// \param[in] column_indices list of column ordinal to read page index later. If it is + /// empty, it means all columns in the row group will be read. + /// \param[in] selection which kind of page index is required later. + virtual void WillNeed(const std::vector& row_group_indices, + const std::vector& column_indices, + const PageIndexSelection& selection) = 0; + + /// \brief Advise the reader page index of these row groups will not be read anymore. + /// + /// The PageIndexReader implementation has the opportunity to cancel any prefetch or + /// release resource that are related to these row groups. + /// + /// \param[in] row_group_indices list of row group ordinal that whose page index will + /// not be accessed anymore. + virtual void WillNotNeed(const std::vector& row_group_indices) = 0; + + /// \brief Determine the column index and offset index ranges for the given row group. + /// + /// \param[in] row_group_metadata row group metadata to get column chunk metadata. + /// \param[in] columns list of column ordinals to get page index. If the list is empty, + /// it means all columns in the row group. + /// \returns RowGroupIndexReadRange of the specified row group. Throws ParquetException + /// if the selected column ordinal is out of bound or metadata of page index + /// is corrupted. + static RowGroupIndexReadRange DeterminePageIndexRangesInRowGroup( + const RowGroupMetaData& row_group_metadata, const std::vector& columns); +}; + +/// \brief Interface for collecting column index of data pages in a column chunk. +class PARQUET_EXPORT ColumnIndexBuilder { + public: + /// \brief API convenience to create a ColumnIndexBuilder. + static std::unique_ptr Make(const ColumnDescriptor* descr); + + virtual ~ColumnIndexBuilder() = default; + + /// \brief Add statistics of a data page. + /// + /// If the ColumnIndexBuilder has seen any corrupted statistics, it will + /// not update statistics anymore. + /// + /// \param stats Page statistics in the encoded form. + /// \param size_stats Size statistics of the page if available. + virtual void AddPage(const EncodedStatistics& stats, + const SizeStatistics& size_stats) = 0; + + /// \brief Complete the column index. + /// + /// Once called, AddPage() can no longer be called. + /// WriteTo() and Build() can only called after Finish() has been called. + virtual void Finish() = 0; + + /// \brief Serialize the column index thrift message. + /// + /// If the ColumnIndexBuilder has seen any corrupted statistics, it will + /// not write any data to the sink. + /// + /// \param[out] sink output stream to write the serialized message. + /// \param[in] encryptor encryptor to encrypt the serialized column index. + virtual void WriteTo(::arrow::io::OutputStream* sink, + Encryptor* encryptor = NULLPTR) const = 0; + + /// \brief Create a ColumnIndex directly. + /// + /// \return If the ColumnIndexBuilder has seen any corrupted statistics, it simply + /// returns nullptr. Otherwise the column index is built and returned. + virtual std::unique_ptr Build() const = 0; +}; + +/// \brief Interface for collecting offset index of data pages in a column chunk. +class PARQUET_EXPORT OffsetIndexBuilder { + public: + /// \brief API convenience to create a OffsetIndexBuilder. + static std::unique_ptr Make(); + + virtual ~OffsetIndexBuilder() = default; + + /// \brief Add page location and size stats of a data page. + virtual void AddPage(int64_t offset, int32_t compressed_page_size, + int64_t first_row_index, + std::optional unencoded_byte_array_length = {}) = 0; + + /// \brief Add page location and size stats of a data page. + void AddPage(const PageLocation& page_location, const SizeStatistics& size_stats); + + /// \brief Complete the offset index. + /// + /// In the buffered row group mode, data pages are flushed into memory + /// sink and the OffsetIndexBuilder has only collected the relative offset + /// which requires adjustment once they are flushed to the file. + /// + /// \param final_position Final stream offset to add for page offset adjustment. + virtual void Finish(int64_t final_position) = 0; + + /// \brief Serialize the offset index thrift message. + /// + /// \param[out] sink output stream to write the serialized message. + /// \param[in] encryptor encryptor to encrypt the serialized offset index. + virtual void WriteTo(::arrow::io::OutputStream* sink, + Encryptor* encryptor = NULLPTR) const = 0; + + /// \brief Create an OffsetIndex directly. + virtual std::unique_ptr Build() const = 0; +}; + +/// \brief Interface for collecting page index of a parquet file. +class PARQUET_EXPORT PageIndexBuilder { + public: + /// \brief API convenience to create a PageIndexBuilder. + static std::unique_ptr Make( + const SchemaDescriptor* schema, InternalFileEncryptor* file_encryptor = NULLPTR); + + virtual ~PageIndexBuilder() = default; + + /// \brief Start a new row group. + virtual void AppendRowGroup() = 0; + + /// \brief Get the ColumnIndexBuilder from column ordinal. + /// + /// \param i Column ordinal. + /// \return ColumnIndexBuilder for the column and its memory ownership belongs to + /// the PageIndexBuilder. + virtual ColumnIndexBuilder* GetColumnIndexBuilder(int32_t i) = 0; + + /// \brief Get the OffsetIndexBuilder from column ordinal. + /// + /// \param i Column ordinal. + /// \return OffsetIndexBuilder for the column and its memory ownership belongs to + /// the PageIndexBuilder. + virtual OffsetIndexBuilder* GetOffsetIndexBuilder(int32_t i) = 0; + + /// \brief Complete the page index builder and no more write is allowed. + virtual void Finish() = 0; + + /// \brief Serialize the page index thrift message. + /// + /// Only valid column indexes and offset indexes are serialized and their locations + /// are set. + /// + /// \param[out] sink The output stream to write the page index. + /// \param[out] location The location of all page index to the start of sink. + virtual void WriteTo(::arrow::io::OutputStream* sink, + PageIndexLocation* location) const = 0; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/parquet_version.h b/lib/python3.12/site-packages/pyarrow/include/parquet/parquet_version.h new file mode 100644 index 0000000000000000000000000000000000000000..0486e6f09d400dfb051765d0186d6f179fe07480 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/parquet_version.h @@ -0,0 +1,31 @@ +// 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. + +#ifndef PARQUET_VERSION_H +#define PARQUET_VERSION_H + +#define PARQUET_VERSION_MAJOR 23 +#define PARQUET_VERSION_MINOR 0 +#define PARQUET_VERSION_PATCH 0 + +#define PARQUET_SO_VERSION "2300" +#define PARQUET_FULL_SO_VERSION "2300.0.0" + +// define the parquet created by version +#define CREATED_BY_VERSION "parquet-cpp-arrow version 23.0.0" + +#endif // PARQUET_VERSION_H diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/platform.h b/lib/python3.12/site-packages/pyarrow/include/parquet/platform.h new file mode 100644 index 0000000000000000000000000000000000000000..92849347d4e9dafc6d8b81300626fb5dbefe411c --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/platform.h @@ -0,0 +1,116 @@ +// 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. + +#pragma once + +#include +#include + +#include "arrow/buffer.h" // IWYU pragma: export +#include "arrow/io/interfaces.h" // IWYU pragma: export +#include "arrow/status.h" // IWYU pragma: export +#include "arrow/type_fwd.h" // IWYU pragma: export +#include "arrow/util/macros.h" // IWYU pragma: export + +#if defined(_WIN32) || defined(__CYGWIN__) + +# ifdef _MSC_VER +# pragma warning(push) +// Disable warning for STL types usage in DLL interface +// https://web.archive.org/web/20130317015847/http://connect.microsoft.com/VisualStudio/feedback/details/696593/vc-10-vs-2010-basic-string-exports +# pragma warning(disable : 4275 4251) +// Disable diamond inheritance warnings +# pragma warning(disable : 4250) +// Disable macro redefinition warnings +# pragma warning(disable : 4005) +// Disable extern before exported template warnings +# pragma warning(disable : 4910) +# else +# pragma GCC diagnostic ignored "-Wattributes" +# endif + +# ifdef PARQUET_STATIC +# define PARQUET_EXPORT +# elif defined(PARQUET_EXPORTING) +# define PARQUET_EXPORT __declspec(dllexport) +# else +# define PARQUET_EXPORT __declspec(dllimport) +# endif + +# define PARQUET_NO_EXPORT + +#else // Not Windows +# ifndef PARQUET_EXPORT +# define PARQUET_EXPORT __attribute__((visibility("default"))) +# endif +# ifndef PARQUET_NO_EXPORT +# define PARQUET_NO_EXPORT __attribute__((visibility("hidden"))) +# endif +#endif // Non-Windows + +// This is a complicated topic, some reading on it: +// http://www.codesynthesis.com/~boris/blog/2010/01/18/dll-export-cxx-templates/ +#if defined(_MSC_VER) || defined(__clang__) +# define PARQUET_TEMPLATE_CLASS_EXPORT +# define PARQUET_TEMPLATE_EXPORT PARQUET_EXPORT +#else +# define PARQUET_TEMPLATE_CLASS_EXPORT PARQUET_EXPORT +# define PARQUET_TEMPLATE_EXPORT +#endif + +#define PARQUET_DISALLOW_COPY_AND_ASSIGN ARROW_DISALLOW_COPY_AND_ASSIGN + +#define PARQUET_NORETURN ARROW_NORETURN +#define PARQUET_DEPRECATED ARROW_DEPRECATED + +// If ARROW_VALGRIND set when compiling unit tests, also define +// PARQUET_VALGRIND +#ifdef ARROW_VALGRIND +# define PARQUET_VALGRIND +#endif + +namespace parquet { + +using Buffer = ::arrow::Buffer; +using Codec = ::arrow::util::Codec; +using CodecOptions = ::arrow::util::CodecOptions; +using Compression = ::arrow::Compression; +using MemoryPool = ::arrow::MemoryPool; +using MutableBuffer = ::arrow::MutableBuffer; +using ResizableBuffer = ::arrow::ResizableBuffer; +using ResizableBuffer = ::arrow::ResizableBuffer; +using ArrowInputFile = ::arrow::io::RandomAccessFile; +using ArrowInputStream = ::arrow::io::InputStream; +using ArrowOutputStream = ::arrow::io::OutputStream; + +constexpr int64_t kDefaultOutputStreamSize = 1024; + +constexpr int16_t kNonPageOrdinal = static_cast(-1); + +PARQUET_EXPORT +std::shared_ptr<::arrow::io::BufferOutputStream> CreateOutputStream( + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + +PARQUET_EXPORT +std::shared_ptr AllocateBuffer( + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool(), int64_t size = 0); + +} // namespace parquet + +#ifdef _MSC_VER +# pragma warning(pop) +#endif diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/printer.h b/lib/python3.12/site-packages/pyarrow/include/parquet/printer.h new file mode 100644 index 0000000000000000000000000000000000000000..bb86b107f9f9ba048b07a85638df383db025b390 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/printer.h @@ -0,0 +1,46 @@ +// 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. + +#pragma once + +#include +#include + +#include "parquet/platform.h" + +namespace parquet { + +class ParquetFileReader; + +class PARQUET_EXPORT ParquetFilePrinter { + private: + ParquetFileReader* fileReader; + + public: + explicit ParquetFilePrinter(ParquetFileReader* reader) : fileReader(reader) {} + ~ParquetFilePrinter() = default; + + void DebugPrint(std::ostream& stream, std::list selected_columns, + bool print_values = false, bool format_dump = false, + bool print_key_value_metadata = false, + const char* filename = "No Name"); + + void JSONPrint(std::ostream& stream, std::list selected_columns, + const char* filename = "No Name"); +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/size_statistics.h b/lib/python3.12/site-packages/pyarrow/include/parquet/size_statistics.h new file mode 100644 index 0000000000000000000000000000000000000000..ec79b8c4f8b8c07b69ac92e44d053a3fd821befd --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/size_statistics.h @@ -0,0 +1,102 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "arrow/util/span.h" +#include "parquet/platform.h" +#include "parquet/type_fwd.h" + +namespace parquet { + +/// A structure for capturing metadata for estimating the unencoded, +/// uncompressed size of data written. This is useful for readers to estimate +/// how much memory is needed to reconstruct data in their memory model and for +/// fine-grained filter push down on nested structures (the histograms contained +/// in this structure can help determine the number of nulls at a particular +/// nesting level and maximum length of lists). +struct PARQUET_EXPORT SizeStatistics { + /// When present, there is expected to be one element corresponding to each + /// definition (i.e. size=max definition+1) where each element + /// represents the number of times the definition level was observed in the + /// data. + /// + /// This field may be omitted (a.k.a. zero-length vector) if max_definition_level + /// is 0 without loss of information. + std::vector definition_level_histogram; + + /// Same as definition_level_histogram except for repetition levels. + /// + /// This field may be omitted (a.k.a. zero-length vector) if max_repetition_level + /// is 0 without loss of information. + std::vector repetition_level_histogram; + + /// The number of physical bytes stored for BYTE_ARRAY data values assuming + /// no encoding. This is exclusive of the bytes needed to store the length of + /// each byte array. In other words, this field is equivalent to the `(size + /// of PLAIN-ENCODING the byte array values) - (4 bytes * number of values + /// written)`. To determine unencoded sizes of other types readers can use + /// schema information multiplied by the number of non-null and null values. + /// The number of null/non-null values can be inferred from the histograms + /// below. + /// + /// For example, if a column chunk is dictionary-encoded with dictionary + /// ["a", "bc", "cde"], and a data page contains the indices [0, 0, 1, 2], + /// then this value for that data page should be 7 (1 + 1 + 2 + 3). + /// + /// This field should only be set for types that use BYTE_ARRAY as their + /// physical type. + std::optional unencoded_byte_array_data_bytes; + + /// \brief Check if the SizeStatistics is set. + bool is_set() const { + return !repetition_level_histogram.empty() || !definition_level_histogram.empty() || + unencoded_byte_array_data_bytes.has_value(); + } + + /// \brief Increment the unencoded byte array data bytes. + void IncrementUnencodedByteArrayDataBytes(int64_t value); + + /// \brief Merge two SizeStatistics. + /// \throws ParquetException if SizeStatistics to merge is not compatible. + void Merge(const SizeStatistics& other); + + /// \brief Validate the SizeStatistics + /// \throws ParquetException if the histograms don't have the right length, + /// or if unencoded_byte_array_data_bytes is present for a non-BYTE_ARRAY column. + void Validate(const ColumnDescriptor* descr) const; + + /// \brief Reset the SizeStatistics to be empty. + void Reset(); + + /// \brief Make an empty SizeStatistics object for specific type. + static std::unique_ptr Make(const ColumnDescriptor* descr); +}; + +PARQUET_EXPORT +std::ostream& operator<<(std::ostream&, const SizeStatistics&); + +PARQUET_EXPORT +void UpdateLevelHistogram(::arrow::util::span levels, + ::arrow::util::span histogram); + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/statistics.h b/lib/python3.12/site-packages/pyarrow/include/parquet/statistics.h new file mode 100644 index 0000000000000000000000000000000000000000..c80fb8e3b52847b3803791988c448d58242a8555 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/statistics.h @@ -0,0 +1,441 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace arrow { + +class Array; +class BinaryArray; + +} // namespace arrow + +namespace parquet { + +class ColumnDescriptor; + +// ---------------------------------------------------------------------- +// Value comparator interfaces + +/// \brief Base class for value comparators. Generally used with +/// TypedComparator +class PARQUET_EXPORT Comparator { + public: + virtual ~Comparator() {} + + /// \brief Create a comparator explicitly from physical type and + /// sort order + /// \param[in] physical_type the physical type for the typed + /// comparator + /// \param[in] sort_order either SortOrder::SIGNED or + /// SortOrder::UNSIGNED + /// \param[in] type_length for FIXED_LEN_BYTE_ARRAY only + static std::shared_ptr Make(Type::type physical_type, + SortOrder::type sort_order, + int type_length = -1); + + /// \brief Create typed comparator inferring default sort order from + /// ColumnDescriptor + /// \param[in] descr the Parquet column schema + static std::shared_ptr Make(const ColumnDescriptor* descr); +}; + +/// \brief Interface for comparison of physical types according to the +/// semantics of a particular logical type. +template +class TypedComparator : public Comparator { + public: + using T = typename DType::c_type; + + /// \brief Scalar comparison of two elements, return true if first + /// is strictly less than the second + virtual bool Compare(const T& a, const T& b) const = 0; + + /// \brief Compute maximum and minimum elements in a batch of + /// elements without any nulls + virtual std::pair GetMinMax(const T* values, int64_t length) const = 0; + + /// \brief Compute minimum and maximum elements from an Arrow array. Only + /// valid for certain Parquet Type / Arrow Type combinations, like BYTE_ARRAY + /// / arrow::BinaryArray + virtual std::pair GetMinMax(const ::arrow::Array& values) const = 0; + + /// \brief Compute maximum and minimum elements in a batch of + /// elements with accompanying bitmap indicating which elements are + /// included (bit set) and excluded (bit not set) + /// + /// \param[in] values the sequence of values + /// \param[in] length the length of the sequence + /// \param[in] valid_bits a bitmap indicating which elements are + /// included (1) or excluded (0) + /// \param[in] valid_bits_offset the bit offset into the bitmap of + /// the first element in the sequence + virtual std::pair GetMinMaxSpaced(const T* values, int64_t length, + const uint8_t* valid_bits, + int64_t valid_bits_offset) const = 0; +}; + +/// \brief Typed version of Comparator::Make +template +std::shared_ptr> MakeComparator(Type::type physical_type, + SortOrder::type sort_order, + int type_length = -1) { + return std::static_pointer_cast>( + Comparator::Make(physical_type, sort_order, type_length)); +} + +/// \brief Typed version of Comparator::Make +template +std::shared_ptr> MakeComparator(const ColumnDescriptor* descr) { + return std::static_pointer_cast>(Comparator::Make(descr)); +} + +// ---------------------------------------------------------------------- + +/// \brief Structure represented encoded statistics to be written to +/// and read from Parquet serialized metadata. +class PARQUET_EXPORT EncodedStatistics { + std::string max_, min_; + bool is_signed_ = false; + + public: + EncodedStatistics() = default; + + const std::string& max() const { return max_; } + const std::string& min() const { return min_; } + + std::optional is_max_value_exact; + std::optional is_min_value_exact; + + int64_t null_count = 0; + int64_t distinct_count = 0; + + bool has_min = false; + bool has_max = false; + bool has_null_count = false; + bool has_distinct_count = false; + + // When all values in the statistics are null, it is set to true. + // Otherwise, at least one value is not null, or we are not sure at all. + // Page index requires this information to decide whether a data page + // is a null page or not. + bool all_null_value = false; + + // From parquet-mr + // Don't write stats larger than the max size rather than truncating. The + // rationale is that some engines may use the minimum value in the page as + // the true minimum for aggregations and there is no way to mark that a + // value has been truncated and is a lower bound and not in the page. + void ApplyStatSizeLimits(size_t length) { + if (max_.length() > length) { + has_max = false; + max_.clear(); + is_max_value_exact = std::nullopt; + } + if (min_.length() > length) { + has_min = false; + min_.clear(); + is_min_value_exact = std::nullopt; + } + } + + // Clear Min Max. + void ClearMinMax() { + has_max = false; + max_.clear(); + has_min = false; + min_.clear(); + } + + bool is_set() const { + return has_min || has_max || has_null_count || has_distinct_count; + } + + bool is_signed() const { return is_signed_; } + + void set_is_signed(bool is_signed) { is_signed_ = is_signed; } + + EncodedStatistics& set_max(std::string value) { + max_ = std::move(value); + has_max = true; + return *this; + } + + EncodedStatistics& set_min(std::string value) { + min_ = std::move(value); + has_min = true; + return *this; + } + + EncodedStatistics& set_null_count(int64_t value) { + null_count = value; + has_null_count = true; + return *this; + } + + EncodedStatistics& set_distinct_count(int64_t value) { + distinct_count = value; + has_distinct_count = true; + return *this; + } +}; + +/// \brief Base type for computing column statistics while writing a file +class PARQUET_EXPORT Statistics { + public: + virtual ~Statistics() {} + + /// \brief Create a new statistics instance given a column schema + /// definition + /// \param[in] descr the column schema + /// \param[in] pool a memory pool to use for any memory allocations, optional + static std::shared_ptr Make( + const ColumnDescriptor* descr, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + /// \brief Create a new statistics instance given a column schema + /// definition and preexisting state + /// \param[in] descr the column schema + /// \param[in] encoded_min the encoded minimum value + /// \param[in] encoded_max the encoded maximum value + /// \param[in] num_values total number of values + /// \param[in] null_count number of null values + /// \param[in] distinct_count number of distinct values + /// \param[in] has_min_max whether the min/max statistics are set + /// \param[in] has_null_count whether the null_count statistics are set + /// \param[in] has_distinct_count whether the distinct_count statistics are set + /// \param[in] pool a memory pool to use for any memory allocations, optional + static std::shared_ptr Make( + const ColumnDescriptor* descr, const std::string& encoded_min, + const std::string& encoded_max, int64_t num_values, int64_t null_count, + int64_t distinct_count, bool has_min_max, bool has_null_count, + bool has_distinct_count, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + /// \brief Create a new statistics instance given a column schema + /// definition and preexisting state + /// \param[in] descr the column schema + /// \param[in] encoded_min the encoded minimum value + /// \param[in] encoded_max the encoded maximum value + /// \param[in] num_values total number of values + /// \param[in] null_count number of null values + /// \param[in] distinct_count number of distinct values + /// \param[in] has_min_max whether the min/max statistics are set + /// \param[in] has_null_count whether the null_count statistics are set + /// \param[in] has_distinct_count whether the distinct_count statistics are set + /// \param[in] is_min_value_exact whether the min value is exact + /// \param[in] is_max_value_exact whether the max value is exact + /// \param[in] pool a memory pool to use for any memory allocations, optional + static std::shared_ptr Make( + const ColumnDescriptor* descr, const std::string& encoded_min, + const std::string& encoded_max, int64_t num_values, int64_t null_count, + int64_t distinct_count, bool has_min_max, bool has_null_count, + bool has_distinct_count, std::optional is_min_value_exact, + std::optional is_max_value_exact, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + // Helper function to convert EncodedStatistics to Statistics. + // EncodedStatistics does not contain number of non-null values, and it can be + // passed using the num_values parameter. + static std::shared_ptr Make( + const ColumnDescriptor* descr, const EncodedStatistics* encoded_statistics, + int64_t num_values = -1, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()); + + /// \brief Return true if the count of null values is set + virtual bool HasNullCount() const = 0; + + /// \brief The number of null values, may not be set + virtual int64_t null_count() const = 0; + + /// \brief Return true if the count of distinct values is set + virtual bool HasDistinctCount() const = 0; + + /// \brief The number of distinct values, may not be set + virtual int64_t distinct_count() const = 0; + + /// \brief The number of non-null values in the column + virtual int64_t num_values() const = 0; + + /// \brief Return true if both min and max statistics are set. Obtain + /// with TypedStatistics::min and max + virtual bool HasMinMax() const = 0; + + /// \brief Reset state of object to initial (no data observed) state + virtual void Reset() = 0; + + /// \brief Plain-encoded minimum value + virtual std::string EncodeMin() const = 0; + + /// \brief Plain-encoded maximum value + virtual std::string EncodeMax() const = 0; + + /// \brief Return the minimum value exact flag if set. + /// It will be true if there was no truncation. + virtual std::optional is_min_value_exact() const = 0; + + /// \brief Return the maximum value exact flag if set. + /// It will be true if there was no truncation. + virtual std::optional is_max_value_exact() const = 0; + + /// \brief The finalized encoded form of the statistics for transport + virtual EncodedStatistics Encode() = 0; + + /// \brief The physical type of the column schema + virtual Type::type physical_type() const = 0; + + /// \brief The full type descriptor from the column schema + virtual const ColumnDescriptor* descr() const = 0; + + /// \brief Check two Statistics for equality + virtual bool Equals(const Statistics& other) const = 0; + + protected: + static std::shared_ptr Make(Type::type physical_type, const void* min, + const void* max, int64_t num_values, + int64_t null_count, int64_t distinct_count); +}; + +/// \brief A typed implementation of Statistics +template +class TypedStatistics : public Statistics { + public: + using T = typename DType::c_type; + + /// \brief The current minimum value + virtual const T& min() const = 0; + + /// \brief The current maximum value + virtual const T& max() const = 0; + + /// \brief Update state with state of another Statistics object + virtual void Merge(const TypedStatistics& other) = 0; + + /// \brief Batch statistics update + virtual void Update(const T* values, int64_t num_values, int64_t null_count) = 0; + + /// \brief Batch statistics update with supplied validity bitmap + /// \param[in] values pointer to column values + /// \param[in] valid_bits Pointer to bitmap representing if values are non-null. + /// \param[in] valid_bits_offset Offset offset into valid_bits where the slice of + /// data begins. + /// \param[in] num_spaced_values The length of values in values/valid_bits to inspect + /// when calculating statistics. This can be smaller than + /// num_values+null_count as null_count can include nulls + /// from parents while num_spaced_values does not. + /// \param[in] num_values Number of values that are not null. + /// \param[in] null_count Number of values that are null. + virtual void UpdateSpaced(const T* values, const uint8_t* valid_bits, + int64_t valid_bits_offset, int64_t num_spaced_values, + int64_t num_values, int64_t null_count) = 0; + + /// \brief EXPERIMENTAL: Update statistics with an Arrow array without + /// conversion to a primitive Parquet C type. Only implemented for certain + /// Parquet type / Arrow type combinations like BYTE_ARRAY / + /// arrow::BinaryArray + /// + /// If update_counts is true then the null_count and num_values will be updated + /// based on the null_count of values. Set to false if these are updated + /// elsewhere (e.g. when updating a dictionary where the counts are taken from + /// the indices and not the values) + virtual void Update(const ::arrow::Array& values, bool update_counts = true) = 0; + + /// \brief Set min and max values to particular values + virtual void SetMinMax(const T& min, const T& max) = 0; + + /// \brief Increments the null count directly + /// Use Update to extract the null count from data. Use this if you determine + /// the null count through some other means (e.g. dictionary arrays where the + /// null count is determined from the indices) + virtual void IncrementNullCount(int64_t n) = 0; + + /// \brief Increments the number of values directly + /// The same note on IncrementNullCount applies here + virtual void IncrementNumValues(int64_t n) = 0; +}; + +using BoolStatistics = TypedStatistics; +using Int32Statistics = TypedStatistics; +using Int64Statistics = TypedStatistics; +using FloatStatistics = TypedStatistics; +using DoubleStatistics = TypedStatistics; +using ByteArrayStatistics = TypedStatistics; +using FLBAStatistics = TypedStatistics; + +/// \brief Typed version of Statistics::Make +template +std::shared_ptr> MakeStatistics( + const ColumnDescriptor* descr, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + return std::static_pointer_cast>(Statistics::Make(descr, pool)); +} + +/// \brief Create Statistics initialized to a particular state +/// \param[in] min the minimum value +/// \param[in] max the minimum value +/// \param[in] num_values number of values +/// \param[in] null_count number of null values +/// \param[in] distinct_count number of distinct values +template +std::shared_ptr> MakeStatistics(const typename DType::c_type& min, + const typename DType::c_type& max, + int64_t num_values, + int64_t null_count, + int64_t distinct_count) { + return std::static_pointer_cast>(Statistics::Make( + DType::type_num, &min, &max, num_values, null_count, distinct_count)); +} + +/// \brief Typed version of Statistics::Make +template +std::shared_ptr> MakeStatistics( + const ColumnDescriptor* descr, const std::string& encoded_min, + const std::string& encoded_max, int64_t num_values, int64_t null_count, + int64_t distinct_count, bool has_min_max, bool has_null_count, + bool has_distinct_count, ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + return std::static_pointer_cast>(Statistics::Make( + descr, encoded_min, encoded_max, num_values, null_count, distinct_count, + has_min_max, has_null_count, has_distinct_count, + /*is_min_value_exact=*/std::nullopt, /*is_max_value_exact=*/std::nullopt, pool)); +} + +/// \brief Typed version of Statistics::Make +template +std::shared_ptr> MakeStatistics( + const ColumnDescriptor* descr, const std::string& encoded_min, + const std::string& encoded_max, int64_t num_values, int64_t null_count, + int64_t distinct_count, bool has_min_max, bool has_null_count, + bool has_distinct_count, std::optional is_min_value_exact, + std::optional is_max_value_exact, + ::arrow::MemoryPool* pool = ::arrow::default_memory_pool()) { + return std::static_pointer_cast>( + Statistics::Make(descr, encoded_min, encoded_max, num_values, null_count, + distinct_count, has_min_max, has_null_count, has_distinct_count, + is_min_value_exact, is_max_value_exact, pool)); +} + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/stream_reader.h b/lib/python3.12/site-packages/pyarrow/include/parquet/stream_reader.h new file mode 100644 index 0000000000000000000000000000000000000000..a7dadac92c89277a104e3acc4149a77258177c8c --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/stream_reader.h @@ -0,0 +1,303 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "parquet/column_reader.h" +#include "parquet/file_reader.h" +#include "parquet/stream_writer.h" + +namespace parquet { + +/// \brief A class for reading Parquet files using an output stream type API. +/// +/// The values given must be of the correct type i.e. the type must +/// match the file schema exactly otherwise a ParquetException will be +/// thrown. +/// +/// The user must explicitly advance to the next row using the +/// EndRow() function or EndRow input manipulator. +/// +/// Required and optional fields are supported: +/// - Required fields are read using operator>>(T) +/// - Optional fields are read with +/// operator>>(std::optional) +/// +/// Note that operator>>(std::optional) can be used to read +/// required fields. +/// +/// Similarly operator>>(T) can be used to read optional fields. +/// However, if the value is not present then a ParquetException will +/// be raised. +/// +/// Currently there is no support for repeated fields. +/// +class PARQUET_EXPORT StreamReader { + public: + template + using optional = ::std::optional; + + // N.B. Default constructed objects are not usable. This + // constructor is provided so that the object may be move + // assigned afterwards. + StreamReader() = default; + + explicit StreamReader(std::unique_ptr reader); + + ~StreamReader() = default; + + bool eof() const { return eof_; } + + int current_column() const { return column_index_; } + + int64_t current_row() const { return current_row_; } + + int num_columns() const; + + int64_t num_rows() const; + + // Moving is possible. + StreamReader(StreamReader&&) = default; + StreamReader& operator=(StreamReader&&) = default; + + // Copying is not allowed. + StreamReader(const StreamReader&) = delete; + StreamReader& operator=(const StreamReader&) = delete; + + StreamReader& operator>>(bool& v); + + StreamReader& operator>>(int8_t& v); + + StreamReader& operator>>(uint8_t& v); + + StreamReader& operator>>(int16_t& v); + + StreamReader& operator>>(uint16_t& v); + + StreamReader& operator>>(int32_t& v); + + StreamReader& operator>>(uint32_t& v); + + StreamReader& operator>>(int64_t& v); + + StreamReader& operator>>(uint64_t& v); + + StreamReader& operator>>(std::chrono::milliseconds& v); + + StreamReader& operator>>(std::chrono::microseconds& v); + + StreamReader& operator>>(float& v); + + StreamReader& operator>>(double& v); + + StreamReader& operator>>(char& v); + + template + StreamReader& operator>>(char (&v)[N]) { + ReadFixedLength(v, N); + return *this; + } + + template + StreamReader& operator>>(std::array& v) { + ReadFixedLength(v.data(), static_cast(N)); + return *this; + } + + // N.B. Cannot allow for reading to a arbitrary char pointer as the + // length cannot be verified. Also it would overshadow the + // char[N] input operator. + // StreamReader& operator>>(char * v); + + StreamReader& operator>>(std::string& v); + + StreamReader& operator>>(::arrow::Decimal128& v); + + // Input operators for optional fields. + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional& v); + + StreamReader& operator>>(optional<::arrow::Decimal128>& v); + + template + StreamReader& operator>>(optional>& v) { + CheckColumn(Type::FIXED_LEN_BYTE_ARRAY, ConvertedType::NONE, N); + FixedLenByteArray flba; + if (ReadOptional(&flba)) { + v = std::array{}; + std::memcpy(v->data(), flba.ptr, N); + } else { + v.reset(); + } + return *this; + } + + /// \brief Terminate current row and advance to next one. + /// \throws ParquetException if all columns in the row were not + /// read or skipped. + void EndRow(); + + /// \brief Skip the data in the next columns. + /// If the number of columns exceeds the columns remaining on the + /// current row then skipping is terminated - it does _not_ continue + /// skipping columns on the next row. + /// Skipping of columns still requires the use 'EndRow' even if all + /// remaining columns were skipped. + /// \return Number of columns actually skipped. + int64_t SkipColumns(int64_t num_columns_to_skip); + + /// \brief Skip the data in the next rows. + /// Skipping of rows is not allowed if reading of data for the + /// current row is not finished. + /// Skipping of rows will be terminated if the end of file is + /// reached. + /// \return Number of rows actually skipped. + int64_t SkipRows(int64_t num_rows_to_skip); + + protected: + [[noreturn]] void ThrowReadFailedException( + const std::shared_ptr& node); + + template + void Read(T* v) { + const auto& node = nodes_[column_index_]; + auto reader = static_cast(column_readers_[column_index_++].get()); + int16_t def_level; + int16_t rep_level; + int64_t values_read; + + reader->ReadBatch(kBatchSizeOne, &def_level, &rep_level, v, &values_read); + + if (values_read != 1) { + ThrowReadFailedException(node); + } + } + + template + void Read(T* v) { + const auto& node = nodes_[column_index_]; + auto reader = static_cast(column_readers_[column_index_++].get()); + int16_t def_level; + int16_t rep_level; + ReadType tmp; + int64_t values_read; + + reader->ReadBatch(kBatchSizeOne, &def_level, &rep_level, &tmp, &values_read); + + if (values_read == 1) { + *v = tmp; + } else { + ThrowReadFailedException(node); + } + } + + template + void ReadOptional(optional* v) { + const auto& node = nodes_[column_index_]; + auto reader = static_cast(column_readers_[column_index_++].get()); + int16_t def_level; + int16_t rep_level; + ReadType tmp; + int64_t values_read; + + reader->ReadBatch(kBatchSizeOne, &def_level, &rep_level, &tmp, &values_read); + + if (values_read == 1) { + *v = T(tmp); + } else if ((values_read == 0) && (def_level == 0)) { + v->reset(); + } else { + ThrowReadFailedException(node); + } + } + + void ReadFixedLength(char* ptr, int len); + + void Read(ByteArray* v); + + void Read(FixedLenByteArray* v); + + bool ReadOptional(ByteArray* v); + + bool ReadOptional(FixedLenByteArray* v); + + void NextRowGroup(); + + void CheckColumn(Type::type physical_type, ConvertedType::type converted_type, + int length = 0); + + void SkipRowsInColumn(ColumnReader* reader, int64_t num_rows_to_skip); + + void SetEof(); + + private: + std::unique_ptr file_reader_; + std::shared_ptr file_metadata_; + std::shared_ptr row_group_reader_; + std::vector> column_readers_; + std::vector> nodes_; + + bool eof_{true}; + int row_group_index_{0}; + int column_index_{0}; + int64_t current_row_{0}; + int64_t row_group_row_offset_{0}; + + static constexpr int64_t kBatchSizeOne = 1; +}; // namespace parquet + +PARQUET_EXPORT +StreamReader& operator>>(StreamReader&, EndRowType); + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/type_fwd.h b/lib/python3.12/site-packages/pyarrow/include/parquet/type_fwd.h new file mode 100644 index 0000000000000000000000000000000000000000..02e896598bfa1637634f99a41718804237ae6ce8 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/type_fwd.h @@ -0,0 +1,105 @@ +// 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. + +#pragma once + +namespace parquet { + +/// \brief Feature selection when writing Parquet files +/// +/// `ParquetVersion::type` governs which data types are allowed and how they +/// are represented. For example, uint32_t data will be written differently +/// depending on this value (as INT64 for PARQUET_1_0, as UINT32 for other +/// versions). +/// +/// However, some features - such as compression algorithms, encryption, +/// or the improved "v2" data page format - must be enabled separately in +/// ArrowWriterProperties. +struct ParquetVersion { + enum type : int { + /// Enable only pre-2.2 Parquet format features when writing + /// + /// This setting is useful for maximum compatibility with legacy readers. + /// Note that logical types may still be emitted, as long they have a + /// corresponding converted type. + PARQUET_1_0, + + /// Enable Parquet format 2.4 and earlier features when writing + /// + /// This enables UINT32 as well as logical types which don't have + /// a corresponding converted type. + /// + /// Note: Parquet format 2.4.0 was released in October 2017. + PARQUET_2_4, + + /// Enable Parquet format 2.6 and earlier features when writing + /// + /// This enables the NANOS time unit in addition to the PARQUET_2_4 + /// features. + /// + /// Note: Parquet format 2.6.0 was released in September 2018. + PARQUET_2_6, + + /// Enable latest Parquet format 2.x features + /// + /// This value is equal to the greatest 2.x version supported by + /// this library. + PARQUET_2_LATEST = PARQUET_2_6 + }; +}; + +struct PageIndexLocation; + +class FileMetaData; +class FileCryptoMetaData; +class RowGroupMetaData; + +class ColumnDescriptor; +class SchemaDescriptor; + +class ReaderProperties; +class ArrowReaderProperties; + +class WriterProperties; +class WriterPropertiesBuilder; +class ArrowWriterProperties; +class ArrowWriterPropertiesBuilder; + +class EncodedStatistics; +class Statistics; +struct SizeStatistics; + +namespace geospatial { +class GeoStatistics; +struct EncodedGeoStatistics; +} // namespace geospatial + +class ColumnIndex; +class OffsetIndex; + +namespace arrow { + +class FileWriter; +class FileReader; + +} // namespace arrow + +namespace schema { +class ColumnPath; +} // namespace schema + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/types.h b/lib/python3.12/site-packages/pyarrow/include/parquet/types.h new file mode 100644 index 0000000000000000000000000000000000000000..7e8a18fc94d68ec426bc8765538d35128542c899 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/types.h @@ -0,0 +1,883 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "parquet/platform.h" +#include "parquet/type_fwd.h" +#include "parquet/windows_fixup.h" // for OPTIONAL + +namespace arrow::util { + +class Codec; + +} // namespace arrow::util + +namespace parquet { + +// ---------------------------------------------------------------------- +// Metadata enums to match Thrift metadata +// +// The reason we maintain our own enums is to avoid transitive dependency on +// the compiled Thrift headers (and thus thrift/Thrift.h) for users of the +// public API. After building parquet-cpp, you should not need to include +// Thrift headers in your application. This means some boilerplate to convert +// between our types and Parquet's Thrift types. +// +// We can also add special values like NONE to distinguish between metadata +// values being set and not set. As an example consider ConvertedType and +// CompressionCodec + +// Mirrors parquet::Type +struct Type { + enum type { + BOOLEAN = 0, + INT32 = 1, + INT64 = 2, + INT96 = 3, + FLOAT = 4, + DOUBLE = 5, + BYTE_ARRAY = 6, + FIXED_LEN_BYTE_ARRAY = 7, + // Should always be last element. + UNDEFINED = 8 + }; +}; + +// Mirrors parquet::ConvertedType +struct ConvertedType { + enum type { + NONE, // Not a real converted type, but means no converted type is specified + UTF8, + MAP, + MAP_KEY_VALUE, + LIST, + ENUM, + DECIMAL, + DATE, + TIME_MILLIS, + TIME_MICROS, + TIMESTAMP_MILLIS, + TIMESTAMP_MICROS, + UINT_8, + UINT_16, + UINT_32, + UINT_64, + INT_8, + INT_16, + INT_32, + INT_64, + JSON, + BSON, + INTERVAL, + // DEPRECATED INVALID ConvertedType for all-null data. + // Only useful for reading legacy files written out by interim Parquet C++ releases. + // For writing, always emit LogicalType::Null instead. + // See PARQUET-1990. + NA = 25, + UNDEFINED = 26 // Not a real converted type; should always be last element + }; +}; + +// forward declaration +namespace format { + +class LogicalType; + +} + +// Mirrors parquet::FieldRepetitionType +struct Repetition { + enum type { REQUIRED = 0, OPTIONAL = 1, REPEATED = 2, /*Always last*/ UNDEFINED = 3 }; +}; + +// Reference: +// parquet-mr/parquet-hadoop/src/main/java/org/apache/parquet/ +// format/converter/ParquetMetadataConverter.java +// Sort order for page and column statistics. Types are associated with sort +// orders (e.g., UTF8 columns should use UNSIGNED) and column stats are +// aggregated using a sort order. As of parquet-format version 2.3.1, the +// order used to aggregate stats is always SIGNED and is not stored in the +// Parquet file. These stats are discarded for types that need unsigned. +// See PARQUET-686. +struct SortOrder { + enum type { SIGNED, UNSIGNED, UNKNOWN }; +}; + +namespace schema { + +struct DecimalMetadata { + bool isset; + int32_t scale; + int32_t precision; +}; + +} // namespace schema + +/// \brief Implementation of parquet.thrift LogicalType types. +class PARQUET_EXPORT LogicalType { + public: + struct Type { + enum type { + UNDEFINED = 0, // Not a real logical type + STRING = 1, + MAP, + LIST, + ENUM, + DECIMAL, + DATE, + TIME, + TIMESTAMP, + INTERVAL, + INT, + NIL, // Thrift NullType: annotates data that is always null + JSON, + BSON, + UUID, + FLOAT16, + GEOMETRY, + GEOGRAPHY, + VARIANT, + NONE // Not a real logical type; should always be last element + }; + }; + + struct TimeUnit { + enum unit { UNKNOWN = 0, MILLIS = 1, MICROS, NANOS }; + }; + + enum class EdgeInterpolationAlgorithm { + UNKNOWN = 0, + SPHERICAL = 1, + VINCENTY = 2, + THOMAS = 3, + ANDOYER = 4, + KARNEY = 5 + }; + + /// \brief The latest supported Variant specification version by this library + static constexpr int8_t kVariantSpecVersion = 1; + + /// \brief If possible, return a logical type equivalent to the given legacy + /// converted type (and decimal metadata if applicable). + static std::shared_ptr FromConvertedType( + const parquet::ConvertedType::type converted_type, + const parquet::schema::DecimalMetadata converted_decimal_metadata = {false, -1, + -1}); + + /// \brief Return the logical type represented by the Thrift intermediary object. + static std::shared_ptr FromThrift( + const parquet::format::LogicalType& thrift_logical_type); + + /// \brief Return the explicitly requested logical type. + static std::shared_ptr String(); + static std::shared_ptr Map(); + static std::shared_ptr List(); + static std::shared_ptr Enum(); + static std::shared_ptr Decimal(int32_t precision, int32_t scale = 0); + static std::shared_ptr Date(); + static std::shared_ptr Time(bool is_adjusted_to_utc, + LogicalType::TimeUnit::unit time_unit); + + /// \brief Create a Timestamp logical type + /// \param[in] is_adjusted_to_utc set true if the data is UTC-normalized + /// \param[in] time_unit the resolution of the timestamp + /// \param[in] is_from_converted_type if true, the timestamp was generated + /// by translating a legacy converted type of TIMESTAMP_MILLIS or + /// TIMESTAMP_MICROS. Default is false. + /// \param[in] force_set_converted_type if true, always set the + /// legacy ConvertedType TIMESTAMP_MICROS and TIMESTAMP_MILLIS + /// metadata. Default is false + static std::shared_ptr Timestamp( + bool is_adjusted_to_utc, LogicalType::TimeUnit::unit time_unit, + bool is_from_converted_type = false, bool force_set_converted_type = false); + + static std::shared_ptr Interval(); + static std::shared_ptr Int(int bit_width, bool is_signed); + + /// \brief Create a logical type for data that's always null + /// + /// Any physical type can be annotated with this logical type. + static std::shared_ptr Null(); + + static std::shared_ptr JSON(); + static std::shared_ptr BSON(); + static std::shared_ptr UUID(); + static std::shared_ptr Float16(); + static std::shared_ptr Variant( + int8_t specVersion = kVariantSpecVersion); + + static std::shared_ptr Geometry(std::string crs = ""); + + static std::shared_ptr Geography( + std::string crs = "", LogicalType::EdgeInterpolationAlgorithm algorithm = + EdgeInterpolationAlgorithm::SPHERICAL); + + /// \brief Create a placeholder for when no logical type is specified + static std::shared_ptr None(); + + /// \brief Return true if this logical type is consistent with the given underlying + /// physical type. + bool is_applicable(parquet::Type::type primitive_type, + int32_t primitive_length = -1) const; + + /// \brief Return true if this logical type is equivalent to the given legacy converted + /// type (and decimal metadata if applicable). + bool is_compatible(parquet::ConvertedType::type converted_type, + parquet::schema::DecimalMetadata converted_decimal_metadata = { + false, -1, -1}) const; + + /// \brief If possible, return the legacy converted type (and decimal metadata if + /// applicable) equivalent to this logical type. + parquet::ConvertedType::type ToConvertedType( + parquet::schema::DecimalMetadata* out_decimal_metadata) const; + + /// \brief Return a printable representation of this logical type. + std::string ToString() const; + + /// \brief Return a JSON representation of this logical type. + std::string ToJSON() const; + + /// \brief Return a serializable Thrift object for this logical type. + parquet::format::LogicalType ToThrift() const; + + /// \brief Return true if the given logical type is equivalent to this logical type. + bool Equals(const LogicalType& other) const; + + /// \brief Return the enumerated type of this logical type. + LogicalType::Type::type type() const; + + /// \brief Return the appropriate sort order for this logical type. + SortOrder::type sort_order() const; + + // Type checks ... + bool is_string() const; + bool is_map() const; + bool is_list() const; + bool is_enum() const; + bool is_decimal() const; + bool is_date() const; + bool is_time() const; + bool is_timestamp() const; + bool is_interval() const; + bool is_int() const; + bool is_null() const; + bool is_JSON() const; + bool is_BSON() const; + bool is_UUID() const; + bool is_float16() const; + bool is_geometry() const; + bool is_geography() const; + bool is_variant() const; + bool is_none() const; + /// \brief Return true if this logical type is of a known type. + bool is_valid() const; + bool is_invalid() const; + /// \brief Return true if this logical type is suitable for a schema GroupNode. + bool is_nested() const; + bool is_nonnested() const; + /// \brief Return true if this logical type is included in the Thrift output for its + /// node. + bool is_serialized() const; + + LogicalType(const LogicalType&) = delete; + LogicalType& operator=(const LogicalType&) = delete; + virtual ~LogicalType() noexcept; + + protected: + LogicalType(); + + class Impl; + std::unique_ptr impl_; +}; + +/// \brief Allowed for physical type BYTE_ARRAY, must be encoded as UTF-8. +class PARQUET_EXPORT StringLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + StringLogicalType() = default; +}; + +/// \brief Allowed for group nodes only. +class PARQUET_EXPORT MapLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + MapLogicalType() = default; +}; + +/// \brief Allowed for group nodes only. +class PARQUET_EXPORT ListLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + ListLogicalType() = default; +}; + +/// \brief Allowed for physical type BYTE_ARRAY, must be encoded as UTF-8. +class PARQUET_EXPORT EnumLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + EnumLogicalType() = default; +}; + +/// \brief Allowed for physical type INT32, INT64, FIXED_LEN_BYTE_ARRAY, or BYTE_ARRAY, +/// depending on the precision. +class PARQUET_EXPORT DecimalLogicalType : public LogicalType { + public: + static std::shared_ptr Make(int32_t precision, int32_t scale = 0); + int32_t precision() const; + int32_t scale() const; + + private: + DecimalLogicalType() = default; +}; + +/// \brief Allowed for physical type INT32. +class PARQUET_EXPORT DateLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + DateLogicalType() = default; +}; + +/// \brief Allowed for physical type INT32 (for MILLIS) or INT64 (for MICROS and NANOS). +class PARQUET_EXPORT TimeLogicalType : public LogicalType { + public: + static std::shared_ptr Make(bool is_adjusted_to_utc, + LogicalType::TimeUnit::unit time_unit); + bool is_adjusted_to_utc() const; + LogicalType::TimeUnit::unit time_unit() const; + + private: + TimeLogicalType() = default; +}; + +/// \brief Allowed for physical type INT64. +class PARQUET_EXPORT TimestampLogicalType : public LogicalType { + public: + static std::shared_ptr Make(bool is_adjusted_to_utc, + LogicalType::TimeUnit::unit time_unit, + bool is_from_converted_type = false, + bool force_set_converted_type = false); + bool is_adjusted_to_utc() const; + LogicalType::TimeUnit::unit time_unit() const; + + /// \brief If true, will not set LogicalType in Thrift metadata + bool is_from_converted_type() const; + + /// \brief If true, will set ConvertedType for micros and millis + /// resolution in legacy ConvertedType Thrift metadata + bool force_set_converted_type() const; + + private: + TimestampLogicalType() = default; +}; + +/// \brief Allowed for physical type FIXED_LEN_BYTE_ARRAY with length 12 +class PARQUET_EXPORT IntervalLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + IntervalLogicalType() = default; +}; + +/// \brief Allowed for physical type INT32 (for bit widths 8, 16, and 32) and INT64 +/// (for bit width 64). +class PARQUET_EXPORT IntLogicalType : public LogicalType { + public: + static std::shared_ptr Make(int bit_width, bool is_signed); + int bit_width() const; + bool is_signed() const; + + private: + IntLogicalType() = default; +}; + +/// \brief Allowed for any physical type. +class PARQUET_EXPORT NullLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + NullLogicalType() = default; +}; + +/// \brief Allowed for physical type BYTE_ARRAY. +class PARQUET_EXPORT JSONLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + JSONLogicalType() = default; +}; + +/// \brief Allowed for physical type BYTE_ARRAY. +class PARQUET_EXPORT BSONLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + BSONLogicalType() = default; +}; + +/// \brief Allowed for physical type FIXED_LEN_BYTE_ARRAY with length 16, +/// must encode raw UUID bytes. +class PARQUET_EXPORT UUIDLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + UUIDLogicalType() = default; +}; + +/// \brief Allowed for physical type FIXED_LEN_BYTE_ARRAY with length 2, +/// must encode raw FLOAT16 bytes. +class PARQUET_EXPORT Float16LogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + Float16LogicalType() = default; +}; + +class PARQUET_EXPORT GeometryLogicalType : public LogicalType { + public: + static std::shared_ptr Make(std::string crs = ""); + + const std::string& crs() const; + + private: + GeometryLogicalType() = default; +}; + +class PARQUET_EXPORT GeographyLogicalType : public LogicalType { + public: + static std::shared_ptr Make( + std::string crs = "", LogicalType::EdgeInterpolationAlgorithm algorithm = + EdgeInterpolationAlgorithm::SPHERICAL); + + const std::string& crs() const; + LogicalType::EdgeInterpolationAlgorithm algorithm() const; + std::string_view algorithm_name() const; + + private: + GeographyLogicalType() = default; +}; + +/// \brief Allowed for group nodes only. +class PARQUET_EXPORT VariantLogicalType : public LogicalType { + public: + static std::shared_ptr Make( + int8_t specVersion = kVariantSpecVersion); + + int8_t spec_version() const; + + private: + VariantLogicalType() = default; +}; + +/// \brief Allowed for any physical type. +class PARQUET_EXPORT NoLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + NoLogicalType() = default; +}; + +// Internal API, for unrecognized logical types +class PARQUET_EXPORT UndefinedLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + UndefinedLogicalType() = default; +}; + +// Data encodings. Mirrors parquet::Encoding +struct Encoding { + enum type { + PLAIN = 0, + PLAIN_DICTIONARY = 2, + RLE = 3, + BIT_PACKED = 4, + DELTA_BINARY_PACKED = 5, + DELTA_LENGTH_BYTE_ARRAY = 6, + DELTA_BYTE_ARRAY = 7, + RLE_DICTIONARY = 8, + BYTE_STREAM_SPLIT = 9, + // Should always be last element (except UNKNOWN) + UNDEFINED = 10, + UNKNOWN = 999 + }; +}; + +// Exposed data encodings. It is the encoding of the data read from the file, +// rather than the encoding of the data in the file. E.g., the data encoded as +// RLE_DICTIONARY in the file can be read as dictionary indices by RLE +// decoding, in which case the data read from the file is DICTIONARY encoded. +enum class ExposedEncoding { + NO_ENCODING = 0, // data is not encoded, i.e. already decoded during reading + DICTIONARY = 1 +}; + +/// \brief Return true if Parquet supports indicated compression type +PARQUET_EXPORT +bool IsCodecSupported(Compression::type codec); + +PARQUET_EXPORT +std::unique_ptr GetCodec(Compression::type codec); + +PARQUET_EXPORT +std::unique_ptr GetCodec(Compression::type codec, + const CodecOptions& codec_options); + +PARQUET_EXPORT +std::unique_ptr GetCodec(Compression::type codec, int compression_level); + +struct ParquetCipher { + enum type { AES_GCM_V1 = 0, AES_GCM_CTR_V1 = 1 }; +}; + +struct AadMetadata { + std::string aad_prefix; + std::string aad_file_unique; + bool supply_aad_prefix; +}; + +struct EncryptionAlgorithm { + ParquetCipher::type algorithm; + AadMetadata aad; +}; + +// parquet::PageType +struct PageType { + enum type { + DATA_PAGE, + INDEX_PAGE, + DICTIONARY_PAGE, + DATA_PAGE_V2, + // Should always be last element + UNDEFINED + }; +}; + +bool PageCanUseChecksum(PageType::type pageType); + +class ColumnOrder { + public: + enum type { UNDEFINED, TYPE_DEFINED_ORDER }; + explicit ColumnOrder(ColumnOrder::type column_order) : column_order_(column_order) {} + // Default to Type Defined Order + ColumnOrder() : column_order_(type::TYPE_DEFINED_ORDER) {} + ColumnOrder::type get_order() { return column_order_; } + + static ColumnOrder undefined_; + static ColumnOrder type_defined_; + + private: + ColumnOrder::type column_order_; +}; + +/// \brief BoundaryOrder is a proxy around format::BoundaryOrder. +struct BoundaryOrder { + enum type { + Unordered = 0, + Ascending = 1, + Descending = 2, + // Should always be last element + UNDEFINED = 3 + }; +}; + +/// \brief SortingColumn is a proxy around format::SortingColumn. +struct PARQUET_EXPORT SortingColumn { + // The column index (in this row group) + int32_t column_idx; + + // If true, indicates this column is sorted in descending order. + bool descending; + + // If true, nulls will come before non-null values, otherwise, nulls go at the end. + bool nulls_first; +}; + +inline bool operator==(const SortingColumn& left, const SortingColumn& right) { + return left.nulls_first == right.nulls_first && left.descending == right.descending && + left.column_idx == right.column_idx; +} + +inline bool operator!=(const SortingColumn& left, const SortingColumn& right) { + return !(left == right); +} + +// ---------------------------------------------------------------------- + +struct ByteArray { + ByteArray() : len(0), ptr(NULLPTR) {} + ByteArray(uint32_t len, const uint8_t* ptr) : len(len), ptr(ptr) {} + + ByteArray(::std::string_view view) // NOLINT implicit conversion + : ByteArray(static_cast(view.size()), + reinterpret_cast(view.data())) {} + + explicit operator std::string_view() const { + return std::string_view{reinterpret_cast(ptr), len}; + } + + uint32_t len; + const uint8_t* ptr; +}; + +inline bool operator==(const ByteArray& left, const ByteArray& right) { + return left.len == right.len && + (left.len == 0 || std::memcmp(left.ptr, right.ptr, left.len) == 0); +} + +inline bool operator!=(const ByteArray& left, const ByteArray& right) { + return !(left == right); +} + +struct FixedLenByteArray { + FixedLenByteArray() : ptr(NULLPTR) {} + explicit FixedLenByteArray(const uint8_t* ptr) : ptr(ptr) {} + const uint8_t* ptr; +}; + +using FLBA = FixedLenByteArray; + +// Julian day at unix epoch. +// +// The Julian Day Number (JDN) is the integer assigned to a whole solar day in +// the Julian day count starting from noon Universal time, with Julian day +// number 0 assigned to the day starting at noon on Monday, January 1, 4713 BC, +// proleptic Julian calendar (November 24, 4714 BC, in the proleptic Gregorian +// calendar), +constexpr int64_t kJulianToUnixEpochDays = INT64_C(2440588); +constexpr int64_t kSecondsPerDay = INT64_C(60 * 60 * 24); +constexpr int64_t kMillisecondsPerDay = kSecondsPerDay * INT64_C(1000); +constexpr int64_t kMicrosecondsPerDay = kMillisecondsPerDay * INT64_C(1000); +constexpr int64_t kNanosecondsPerDay = kMicrosecondsPerDay * INT64_C(1000); + +MANUALLY_ALIGNED_STRUCT(1) Int96 { uint32_t value[3]; }; +STRUCT_END(Int96, 12); + +inline bool operator==(const Int96& left, const Int96& right) { + return std::equal(left.value, left.value + 3, right.value); +} + +inline bool operator!=(const Int96& left, const Int96& right) { return !(left == right); } + +static inline std::string ByteArrayToString(const ByteArray& a) { + return std::string(reinterpret_cast(a.ptr), a.len); +} + +static inline void Int96SetNanoSeconds(parquet::Int96& i96, int64_t nanoseconds) { + std::memcpy(&i96.value, &nanoseconds, sizeof(nanoseconds)); +} + +struct DecodedInt96 { + uint64_t days_since_epoch; + uint64_t nanoseconds; +}; + +static inline DecodedInt96 DecodeInt96Timestamp(const parquet::Int96& i96) { + // We do the computations in the unsigned domain to avoid unsigned behaviour + // on overflow. + DecodedInt96 result; + result.days_since_epoch = i96.value[2] - static_cast(kJulianToUnixEpochDays); + result.nanoseconds = 0; + + memcpy(&result.nanoseconds, &i96.value, sizeof(uint64_t)); + return result; +} + +static inline int64_t Int96GetNanoSeconds(const parquet::Int96& i96) { + const auto decoded = DecodeInt96Timestamp(i96); + return static_cast(decoded.days_since_epoch * kNanosecondsPerDay + + decoded.nanoseconds); +} + +static inline int64_t Int96GetMicroSeconds(const parquet::Int96& i96) { + const auto decoded = DecodeInt96Timestamp(i96); + uint64_t microseconds = decoded.nanoseconds / static_cast(1000); + return static_cast(decoded.days_since_epoch * kMicrosecondsPerDay + + microseconds); +} + +static inline int64_t Int96GetMilliSeconds(const parquet::Int96& i96) { + const auto decoded = DecodeInt96Timestamp(i96); + uint64_t milliseconds = decoded.nanoseconds / static_cast(1000000); + return static_cast(decoded.days_since_epoch * kMillisecondsPerDay + + milliseconds); +} + +static inline int64_t Int96GetSeconds(const parquet::Int96& i96) { + const auto decoded = DecodeInt96Timestamp(i96); + uint64_t seconds = decoded.nanoseconds / static_cast(1000000000); + return static_cast(decoded.days_since_epoch * kSecondsPerDay + seconds); +} + +static inline std::string Int96ToString(const Int96& a) { + std::ostringstream result; + std::copy(a.value, a.value + 3, std::ostream_iterator(result, " ")); + return result.str(); +} + +static inline std::string FixedLenByteArrayToString(const FixedLenByteArray& a, int len) { + std::ostringstream result; + std::copy(a.ptr, a.ptr + len, std::ostream_iterator(result, " ")); + return result.str(); +} + +template +struct type_traits {}; + +template <> +struct type_traits { + using value_type = bool; + + static constexpr int value_byte_size = 1; + static constexpr const char* printf_code = "d"; +}; + +template <> +struct type_traits { + using value_type = int32_t; + + static constexpr int value_byte_size = 4; + static constexpr const char* printf_code = "d"; +}; + +template <> +struct type_traits { + using value_type = int64_t; + + static constexpr int value_byte_size = 8; + static constexpr const char* printf_code = + (sizeof(long) == 64) ? "ld" : "lld"; // NOLINT: runtime/int +}; + +template <> +struct type_traits { + using value_type = Int96; + + static constexpr int value_byte_size = 12; + static constexpr const char* printf_code = "s"; +}; + +template <> +struct type_traits { + using value_type = float; + + static constexpr int value_byte_size = 4; + static constexpr const char* printf_code = "f"; +}; + +template <> +struct type_traits { + using value_type = double; + + static constexpr int value_byte_size = 8; + static constexpr const char* printf_code = "lf"; +}; + +template <> +struct type_traits { + using value_type = ByteArray; + + static constexpr int value_byte_size = sizeof(ByteArray); + static constexpr const char* printf_code = "s"; +}; + +template <> +struct type_traits { + using value_type = FixedLenByteArray; + + static constexpr int value_byte_size = sizeof(FixedLenByteArray); + static constexpr const char* printf_code = "s"; +}; + +template +struct PhysicalType { + using c_type = typename type_traits::value_type; + static constexpr Type::type type_num = TYPE; +}; + +using BooleanType = PhysicalType; +using Int32Type = PhysicalType; +using Int64Type = PhysicalType; +using Int96Type = PhysicalType; +using FloatType = PhysicalType; +using DoubleType = PhysicalType; +using ByteArrayType = PhysicalType; +using FLBAType = PhysicalType; + +template +inline std::string format_fwf(int width) { + std::stringstream ss; + ss << "%-" << width << type_traits::printf_code; + return ss.str(); +} + +PARQUET_EXPORT std::string EncodingToString(Encoding::type t); + +PARQUET_EXPORT std::string ConvertedTypeToString(ConvertedType::type t); + +PARQUET_EXPORT std::string TypeToString(Type::type t); + +PARQUET_EXPORT std::string TypeToString(Type::type t, int type_length); + +PARQUET_EXPORT std::string FormatStatValue( + Type::type parquet_type, ::std::string_view val, + const std::shared_ptr& logical_type = NULLPTR); + +PARQUET_EXPORT int GetTypeByteSize(Type::type t); + +PARQUET_EXPORT SortOrder::type DefaultSortOrder(Type::type primitive); + +PARQUET_EXPORT SortOrder::type GetSortOrder(ConvertedType::type converted, + Type::type primitive); + +PARQUET_EXPORT SortOrder::type GetSortOrder( + const std::shared_ptr& logical_type, Type::type primitive); + +// PLAIN_DICTIONARY is deprecated but used to be used as a dictionary index +// encoding. +constexpr bool IsDictionaryIndexEncoding(Encoding::type e) { + return e == Encoding::RLE_DICTIONARY || e == Encoding::PLAIN_DICTIONARY; +} + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/windows_fixup.h b/lib/python3.12/site-packages/pyarrow/include/parquet/windows_fixup.h new file mode 100644 index 0000000000000000000000000000000000000000..feac4e64d19761595105de7cc189e92c9ef6ce09 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/windows_fixup.h @@ -0,0 +1,29 @@ +// 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. + +// This header needs to be included multiple times. + +#include "arrow/util/windows_fixup.h" + +#ifdef _WIN32 + +// parquet.thrift's OPTIONAL RepetitionType conflicts with a Windows #define +# ifdef OPTIONAL +# undef OPTIONAL +# endif + +#endif // _WIN32 diff --git a/lib/python3.12/site-packages/pyarrow/include/parquet/xxhasher.h b/lib/python3.12/site-packages/pyarrow/include/parquet/xxhasher.h new file mode 100644 index 0000000000000000000000000000000000000000..a54f287883e006e9cd6d9aeeb2efeb1d6f9db2df --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/include/parquet/xxhasher.h @@ -0,0 +1,50 @@ +// 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. + +#pragma once + +#include + +#include "parquet/hasher.h" +#include "parquet/platform.h" +#include "parquet/types.h" + +namespace parquet { + +class PARQUET_EXPORT XxHasher : public Hasher { + public: + uint64_t Hash(int32_t value) const override; + uint64_t Hash(int64_t value) const override; + uint64_t Hash(float value) const override; + uint64_t Hash(double value) const override; + uint64_t Hash(const Int96* value) const override; + uint64_t Hash(const ByteArray* value) const override; + uint64_t Hash(const FLBA* val, uint32_t len) const override; + + void Hashes(const int32_t* values, int num_values, uint64_t* hashes) const override; + void Hashes(const int64_t* values, int num_values, uint64_t* hashes) const override; + void Hashes(const float* values, int num_values, uint64_t* hashes) const override; + void Hashes(const double* values, int num_values, uint64_t* hashes) const override; + void Hashes(const Int96* values, int num_values, uint64_t* hashes) const override; + void Hashes(const ByteArray* values, int num_values, uint64_t* hashes) const override; + void Hashes(const FLBA* values, uint32_t type_len, int num_values, + uint64_t* hashes) const override; + + static constexpr int kParquetBloomXxHashSeed = 0; +}; + +} // namespace parquet diff --git a/lib/python3.12/site-packages/pyarrow/includes/__init__.pxd b/lib/python3.12/site-packages/pyarrow/includes/__init__.pxd new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pyarrow/includes/common.pxd b/lib/python3.12/site-packages/pyarrow/includes/common.pxd new file mode 100644 index 0000000000000000000000000000000000000000..f97f2ff1c4732ce8a9491a84e2b1f70ad9c7f9ba --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/common.pxd @@ -0,0 +1,155 @@ +# 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 libc.stdint cimport * + +from libcpp cimport bool as c_bool, nullptr +from libcpp.functional cimport function +from libcpp.memory cimport (shared_ptr, unique_ptr, make_shared, + static_pointer_cast, dynamic_pointer_cast) +from libcpp.optional cimport nullopt, optional +from libcpp.string cimport string as c_string +from libcpp.utility cimport move, pair +from libcpp.vector cimport vector +from libcpp.unordered_map cimport unordered_map +from libcpp.unordered_set cimport unordered_set + +from cpython cimport PyObject +from cpython.datetime cimport PyDateTime_DateTime +cimport cpython + + +cdef extern from "" namespace "std" nogil: + # Needed until https://github.com/cython/cython/issues/6651 is fixed + cdef cppclass cpp_string_view "std::string_view": + string_view() + string_view(const char*) + string_view(c_string&) + size_t size() + bint empty() + const char* data() + + +cdef extern from * namespace "arrow::py" nogil: + """ + #include + #include + #include + #include + + namespace arrow { + namespace py { + + template + std::shared_ptr to_shared(std::unique_ptr& t) { + return std::move(t); + } + template + std::shared_ptr to_shared(std::unique_ptr&& t) { + return std::move(t); + } + + // Needed until https://github.com/cython/cython/issues/6651 is fixed + inline std::string to_string(std::string_view s) { + return std::string(s); + } + + } // namespace py + } // namespace arrow + """ + cdef shared_ptr[T] to_shared" arrow::py::to_shared"[T](unique_ptr[T]) + cdef c_string to_string(cpp_string_view s) + +cdef extern from "arrow/python/platform.h": + pass + +cdef extern from "": + void Py_XDECREF(PyObject* o) + Py_ssize_t Py_REFCNT(PyObject* o) + +cdef extern from "arrow/api.h" namespace "arrow" nogil: + # We can later add more of the common status factory methods as needed + cdef CStatus CStatus_OK "arrow::Status::OK"() + + cdef CStatus CStatus_Invalid "arrow::Status::Invalid"() + cdef CStatus CStatus_NotImplemented \ + "arrow::Status::NotImplemented"(const c_string& msg) + cdef CStatus CStatus_UnknownError \ + "arrow::Status::UnknownError"(const c_string& msg) + + cdef cppclass CStatus "arrow::Status": + CStatus() + + c_string ToString() + c_string message() + shared_ptr[CStatusDetail] detail() + + c_bool ok() + c_bool IsIOError() + c_bool IsOutOfMemory() + c_bool IsInvalid() + c_bool IsKeyError() + c_bool IsNotImplemented() + c_bool IsTypeError() + c_bool IsCapacityError() + c_bool IsIndexError() + c_bool IsSerializationError() + c_bool IsCancelled() + + void Warn() + + cdef cppclass CStatusDetail "arrow::StatusDetail": + c_string ToString() + + +cdef extern from "arrow/result.h" namespace "arrow" nogil: + cdef cppclass CResult "arrow::Result"[T]: + CResult() + CResult(CStatus) + CResult(T) + c_bool ok() + CStatus status() + CStatus Value(T*) + T operator*() + + +cdef extern from "arrow/util/future.h" namespace "arrow" nogil: + cdef cppclass CFuture "arrow::Future"[T]: + CFuture() + + +cdef extern from "arrow/python/async.h" namespace "arrow::py" nogil: + # BindFuture's third argument is really a C++ callable with + # the signature `object(T*)`, but Cython does not allow declaring that. + # We use an ellipsis as a workaround. + # Another possibility is to type-erase the argument by making it + # `object(void*)`, but it would lose compile-time C++ type safety. + void BindFuture[T](CFuture[T], object cb, ...) + + +cdef extern from "arrow/python/common.h" namespace "arrow::py" nogil: + T GetResultValue[T](CResult[T]) except * + cdef function[F] BindFunction[F](void* unbound, object bound, ...) + + +cdef inline object PyObject_to_object(PyObject* o): + # Cast to "object" increments reference count + cdef object result = o + cpython.Py_DECREF(result) + return result diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow.pxd new file mode 100644 index 0000000000000000000000000000000000000000..e96a7d84696d0b78e8236456c2b062d8f1346686 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow.pxd @@ -0,0 +1,3281 @@ +# 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 pyarrow.includes.common cimport * + + +cdef extern from "arrow/util/key_value_metadata.h" namespace "arrow" nogil: + cdef cppclass CKeyValueMetadata" arrow::KeyValueMetadata": + CKeyValueMetadata() + CKeyValueMetadata(const unordered_map[c_string, c_string]&) + CKeyValueMetadata(const vector[c_string]& keys, + const vector[c_string]& values) + + void reserve(int64_t n) + int64_t size() const + c_string key(int64_t i) const + c_string value(int64_t i) const + int FindKey(const c_string& key) const + + shared_ptr[CKeyValueMetadata] Copy() const + c_bool Equals(const CKeyValueMetadata& other) + void Append(const c_string& key, const c_string& value) + void ToUnorderedMap(unordered_map[c_string, c_string]*) const + c_string ToString() const + + CResult[c_string] Get(const c_string& key) const + CStatus Delete(const c_string& key) + CStatus Set(const c_string& key, const c_string& value) + c_bool Contains(const c_string& key) const + + +cdef extern from "arrow/util/decimal.h" namespace "arrow" nogil: + cdef cppclass CDecimal32" arrow::Decimal32": + c_string ToString(int32_t scale) const + + +cdef extern from "arrow/util/decimal.h" namespace "arrow" nogil: + cdef cppclass CDecimal64" arrow::Decimal64": + c_string ToString(int32_t scale) const + + +cdef extern from "arrow/util/decimal.h" namespace "arrow" nogil: + cdef cppclass CDecimal128" arrow::Decimal128": + c_string ToString(int32_t scale) const + + +cdef extern from "arrow/util/decimal.h" namespace "arrow" nogil: + cdef cppclass CDecimal256" arrow::Decimal256": + c_string ToString(int32_t scale) const + + +cdef extern from "arrow/config.h" namespace "arrow" nogil: + cdef cppclass CCppBuildInfo "arrow::BuildInfo": + int version + int version_major + int version_minor + int version_patch + c_string version_string + c_string so_version + c_string full_so_version + c_string compiler_id + c_string compiler_version + c_string compiler_flags + c_string git_id + c_string git_description + c_string package_kind + c_string build_type + + const CCppBuildInfo& GetCppBuildInfo "arrow::GetBuildInfo"() + + cdef cppclass CRuntimeInfo" arrow::RuntimeInfo": + c_string simd_level + c_string detected_simd_level + + CRuntimeInfo GetRuntimeInfo() + + cdef cppclass CGlobalOptions" arrow::GlobalOptions": + optional[c_string] timezone_db_path + + CStatus Initialize(const CGlobalOptions& options) + + +cdef extern from "arrow/util/future.h" namespace "arrow" nogil: + cdef cppclass CFuture_Void" arrow::Future<>": + CStatus status() + + +cdef extern from "" namespace "std" nogil: + cdef cppclass CArrayStatisticsCountType" std::variant": + CArrayStatisticsCountType() + CArrayStatisticsCountType(int64_t) + CArrayStatisticsCountType(double) + + cdef cppclass CArrayStatisticsValueType" std::variant": + CArrayStatisticsValueType() + CArrayStatisticsValueType(c_bool) + CArrayStatisticsValueType(int64_t) + CArrayStatisticsValueType(uint64_t) + CArrayStatisticsValueType(double) + CArrayStatisticsValueType(c_string) + + +cdef extern from "arrow/api.h" namespace "arrow" nogil: + cdef enum Type" arrow::Type::type": + _Type_NA" arrow::Type::NA" + + _Type_BOOL" arrow::Type::BOOL" + + _Type_UINT8" arrow::Type::UINT8" + _Type_INT8" arrow::Type::INT8" + _Type_UINT16" arrow::Type::UINT16" + _Type_INT16" arrow::Type::INT16" + _Type_UINT32" arrow::Type::UINT32" + _Type_INT32" arrow::Type::INT32" + _Type_UINT64" arrow::Type::UINT64" + _Type_INT64" arrow::Type::INT64" + + _Type_HALF_FLOAT" arrow::Type::HALF_FLOAT" + _Type_FLOAT" arrow::Type::FLOAT" + _Type_DOUBLE" arrow::Type::DOUBLE" + + _Type_DECIMAL32" arrow::Type::DECIMAL32" + _Type_DECIMAL64" arrow::Type::DECIMAL64" + _Type_DECIMAL128" arrow::Type::DECIMAL128" + _Type_DECIMAL256" arrow::Type::DECIMAL256" + + _Type_DATE32" arrow::Type::DATE32" + _Type_DATE64" arrow::Type::DATE64" + _Type_TIMESTAMP" arrow::Type::TIMESTAMP" + _Type_TIME32" arrow::Type::TIME32" + _Type_TIME64" arrow::Type::TIME64" + _Type_DURATION" arrow::Type::DURATION" + _Type_INTERVAL_MONTH_DAY_NANO" arrow::Type::INTERVAL_MONTH_DAY_NANO" + _Type_INTERVAL_DAY_TIME" arrow::Type::INTERVAL_DAY_TIME" + _Type_INTERVAL_MONTHS" arrow::Type::INTERVAL_MONTHS" + + _Type_BINARY" arrow::Type::BINARY" + _Type_STRING" arrow::Type::STRING" + _Type_LARGE_BINARY" arrow::Type::LARGE_BINARY" + _Type_LARGE_STRING" arrow::Type::LARGE_STRING" + _Type_FIXED_SIZE_BINARY" arrow::Type::FIXED_SIZE_BINARY" + _Type_BINARY_VIEW" arrow::Type::BINARY_VIEW" + _Type_STRING_VIEW" arrow::Type::STRING_VIEW" + + _Type_LIST" arrow::Type::LIST" + _Type_LARGE_LIST" arrow::Type::LARGE_LIST" + _Type_FIXED_SIZE_LIST" arrow::Type::FIXED_SIZE_LIST" + _Type_LIST_VIEW" arrow::Type::LIST_VIEW" + _Type_LARGE_LIST_VIEW" arrow::Type::LARGE_LIST_VIEW" + _Type_STRUCT" arrow::Type::STRUCT" + _Type_SPARSE_UNION" arrow::Type::SPARSE_UNION" + _Type_DENSE_UNION" arrow::Type::DENSE_UNION" + _Type_DICTIONARY" arrow::Type::DICTIONARY" + _Type_RUN_END_ENCODED" arrow::Type::RUN_END_ENCODED" + _Type_MAP" arrow::Type::MAP" + + _Type_EXTENSION" arrow::Type::EXTENSION" + + cdef enum UnionMode" arrow::UnionMode::type": + _UnionMode_SPARSE" arrow::UnionMode::SPARSE" + _UnionMode_DENSE" arrow::UnionMode::DENSE" + + cdef enum TimeUnit" arrow::TimeUnit::type": + TimeUnit_SECOND" arrow::TimeUnit::SECOND" + TimeUnit_MILLI" arrow::TimeUnit::MILLI" + TimeUnit_MICRO" arrow::TimeUnit::MICRO" + TimeUnit_NANO" arrow::TimeUnit::NANO" + + cdef cppclass CBufferSpec" arrow::DataTypeLayout::BufferSpec": + pass + + cdef cppclass CDataTypeLayout" arrow::DataTypeLayout": + vector[CBufferSpec] buffers + optional[CBufferSpec] variadic_spec + c_bool has_dictionary + + cdef cppclass CDataType" arrow::DataType": + Type id() + + c_bool Equals(const CDataType& other, c_bool check_metadata) + c_bool Equals(const shared_ptr[CDataType]& other, c_bool check_metadata) + + shared_ptr[CField] field(int i) + const vector[shared_ptr[CField]] fields() + int num_fields() + CDataTypeLayout layout() + c_string ToString() + + c_bool is_primitive(Type type) + c_bool is_numeric(Type type) + + cdef cppclass CArrayStatistics" arrow::ArrayStatistics": + optional[CArrayStatisticsCountType] null_count + optional[CArrayStatisticsCountType] distinct_count + optional[CArrayStatisticsValueType] min + c_bool is_min_exact + optional[CArrayStatisticsValueType] max + c_bool is_max_exact + + c_bool Equals(const CArrayStatistics& statistics) const + + cdef cppclass CArrayData" arrow::ArrayData": + shared_ptr[CDataType] type + int64_t length + int64_t null_count + int64_t offset + vector[shared_ptr[CBuffer]] buffers + vector[shared_ptr[CArrayData]] child_data + shared_ptr[CArrayData] dictionary + + @staticmethod + shared_ptr[CArrayData] Make(const shared_ptr[CDataType]& type, + int64_t length, + vector[shared_ptr[CBuffer]]& buffers, + int64_t null_count, + int64_t offset) + + @staticmethod + shared_ptr[CArrayData] MakeWithChildren" Make"( + const shared_ptr[CDataType]& type, + int64_t length, + vector[shared_ptr[CBuffer]]& buffers, + vector[shared_ptr[CArrayData]]& child_data, + int64_t null_count, + int64_t offset) + + @staticmethod + shared_ptr[CArrayData] MakeWithChildrenAndDictionary" Make"( + const shared_ptr[CDataType]& type, + int64_t length, + vector[shared_ptr[CBuffer]]& buffers, + vector[shared_ptr[CArrayData]]& child_data, + shared_ptr[CArrayData]& dictionary, + int64_t null_count, + int64_t offset) + + cdef cppclass CArray" arrow::Array": + shared_ptr[CDataType] type() + + int64_t length() + int64_t null_count() + int64_t offset() + Type type_id() + + int num_fields() + + CResult[shared_ptr[CScalar]] GetScalar(int64_t i) const + + c_string Diff(const CArray& other) + c_bool Equals(const CArray& arr) + c_bool IsNull(int i) + + shared_ptr[CArrayData] data() + + shared_ptr[CArray] Slice(int64_t offset) + shared_ptr[CArray] Slice(int64_t offset, int64_t length) + + CStatus Validate() const + CStatus ValidateFull() const + CResult[shared_ptr[CArray]] View(const shared_ptr[CDataType]& type) + + CDeviceAllocationType device_type() + CResult[shared_ptr[CArray]] CopyTo(const shared_ptr[CMemoryManager]& to) const + + const shared_ptr[CArrayStatistics]& statistics() const + + shared_ptr[CArray] MakeArray(const shared_ptr[CArrayData]& data) + CResult[shared_ptr[CArray]] MakeArrayOfNull( + const shared_ptr[CDataType]& type, int64_t length, CMemoryPool* pool) + + CResult[shared_ptr[CArray]] MakeArrayFromScalar( + const CScalar& scalar, int64_t length, CMemoryPool* pool) + + CStatus DebugPrint(const CArray& arr, int indent) + + cdef cppclass CFixedWidthType" arrow::FixedWidthType"(CDataType): + int bit_width() + int byte_width() + + cdef cppclass CNullArray" arrow::NullArray"(CArray): + CNullArray(int64_t length) + + cdef cppclass CDictionaryArray" arrow::DictionaryArray"(CArray): + CDictionaryArray(const shared_ptr[CDataType]& type, + const shared_ptr[CArray]& indices, + const shared_ptr[CArray]& dictionary) + CDictionaryArray(const shared_ptr[CArrayData]& data) + + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const shared_ptr[CDataType]& type, + const shared_ptr[CArray]& indices, + const shared_ptr[CArray]& dictionary) + + shared_ptr[CArray] indices() + shared_ptr[CArray] dictionary() + + cdef cppclass CDate32Type" arrow::Date32Type"(CFixedWidthType): + pass + + cdef cppclass CDate64Type" arrow::Date64Type"(CFixedWidthType): + pass + + cdef cppclass CTimestampType" arrow::TimestampType"(CFixedWidthType): + CTimestampType(TimeUnit unit) + TimeUnit unit() + const c_string& timezone() + + cdef cppclass CTime32Type" arrow::Time32Type"(CFixedWidthType): + TimeUnit unit() + + cdef cppclass CTime64Type" arrow::Time64Type"(CFixedWidthType): + TimeUnit unit() + + shared_ptr[CDataType] ctime32" arrow::time32"(TimeUnit unit) + shared_ptr[CDataType] ctime64" arrow::time64"(TimeUnit unit) + + cdef cppclass CDurationType" arrow::DurationType"(CFixedWidthType): + TimeUnit unit() + + shared_ptr[CDataType] cduration" arrow::duration"(TimeUnit unit) + + cdef cppclass CDictionaryType" arrow::DictionaryType"(CFixedWidthType): + CDictionaryType(const shared_ptr[CDataType]& index_type, + const shared_ptr[CDataType]& value_type, + c_bool ordered) + + shared_ptr[CDataType] index_type() + shared_ptr[CDataType] value_type() + c_bool ordered() + + shared_ptr[CDataType] ctimestamp" arrow::timestamp"(TimeUnit unit) + shared_ptr[CDataType] ctimestamp" arrow::timestamp"( + TimeUnit unit, const c_string& timezone) + + cdef cppclass CMemoryPool" arrow::MemoryPool": + int64_t bytes_allocated() + int64_t max_memory() + int64_t total_bytes_allocated() + int64_t num_allocations() + c_string backend_name() + void ReleaseUnused() + void PrintStats() + + cdef cppclass CLoggingMemoryPool" arrow::LoggingMemoryPool"(CMemoryPool): + CLoggingMemoryPool(CMemoryPool*) + + cdef cppclass CProxyMemoryPool" arrow::ProxyMemoryPool"(CMemoryPool): + CProxyMemoryPool(CMemoryPool*) + + ctypedef enum CDeviceAllocationType "arrow::DeviceAllocationType": + CDeviceAllocationType_kCPU "arrow::DeviceAllocationType::kCPU" + CDeviceAllocationType_kCUDA "arrow::DeviceAllocationType::kCUDA" + CDeviceAllocationType_kCUDA_HOST "arrow::DeviceAllocationType::kCUDA_HOST" + CDeviceAllocationType_kOPENCL "arrow::DeviceAllocationType::kOPENCL" + CDeviceAllocationType_kVULKAN "arrow::DeviceAllocationType::kVULKAN" + CDeviceAllocationType_kMETAL "arrow::DeviceAllocationType::kMETAL" + CDeviceAllocationType_kVPI "arrow::DeviceAllocationType::kVPI" + CDeviceAllocationType_kROCM "arrow::DeviceAllocationType::kROCM" + CDeviceAllocationType_kROCM_HOST "arrow::DeviceAllocationType::kROCM_HOST" + CDeviceAllocationType_kEXT_DEV "arrow::DeviceAllocationType::kEXT_DEV" + CDeviceAllocationType_kCUDA_MANAGED "arrow::DeviceAllocationType::kCUDA_MANAGED" + CDeviceAllocationType_kONEAPI "arrow::DeviceAllocationType::kONEAPI" + CDeviceAllocationType_kWEBGPU "arrow::DeviceAllocationType::kWEBGPU" + CDeviceAllocationType_kHEXAGON "arrow::DeviceAllocationType::kHEXAGON" + + cdef cppclass CDevice" arrow::Device": + const char* type_name() + c_string ToString() + c_bool Equals(const CDevice& other) + int64_t device_id() + c_bool is_cpu() const + shared_ptr[CMemoryManager] default_memory_manager() + CDeviceAllocationType device_type() + + cdef cppclass CMemoryManager" arrow::MemoryManager": + const shared_ptr[CDevice] device() + c_bool is_cpu() const + + shared_ptr[CMemoryManager] c_default_cpu_memory_manager \ + " arrow::default_cpu_memory_manager"() + + cdef cppclass CBuffer" arrow::Buffer": + CBuffer(const uint8_t* data, int64_t size) + const uint8_t* data() + uint8_t* mutable_data() + uintptr_t address() + uintptr_t mutable_address() + int64_t size() + shared_ptr[CBuffer] parent() + c_bool is_cpu() const + c_bool is_mutable() const + c_string ToHexString() + c_bool Equals(const CBuffer& other) + shared_ptr[CDevice] device() + const shared_ptr[CMemoryManager] memory_manager() + CDeviceAllocationType device_type() + + CResult[shared_ptr[CBuffer]] SliceBufferSafe( + const shared_ptr[CBuffer]& buffer, int64_t offset) + CResult[shared_ptr[CBuffer]] SliceBufferSafe( + const shared_ptr[CBuffer]& buffer, int64_t offset, int64_t length) + + cdef cppclass CMutableBuffer" arrow::MutableBuffer"(CBuffer): + CMutableBuffer(const uint8_t* data, int64_t size) + + cdef cppclass CResizableBuffer" arrow::ResizableBuffer"(CMutableBuffer): + CStatus Resize(const int64_t new_size, c_bool shrink_to_fit) + CStatus Reserve(const int64_t new_size) + + CResult[unique_ptr[CBuffer]] AllocateBuffer(const int64_t size, + CMemoryPool* pool) + + CResult[unique_ptr[CResizableBuffer]] AllocateResizableBuffer( + const int64_t size, CMemoryPool* pool) + + cdef cppclass CSyncEvent" arrow::Device::SyncEvent": + pass + + cdef cppclass CDevice" arrow::Device": + pass + + cdef CMemoryPool* c_default_memory_pool" arrow::default_memory_pool"() + cdef CMemoryPool* c_system_memory_pool" arrow::system_memory_pool"() + cdef CStatus c_jemalloc_memory_pool" arrow::jemalloc_memory_pool"( + CMemoryPool** out) + cdef CStatus c_mimalloc_memory_pool" arrow::mimalloc_memory_pool"( + CMemoryPool** out) + cdef vector[c_string] c_supported_memory_backends \ + " arrow::SupportedMemoryBackendNames"() + + CStatus c_jemalloc_set_decay_ms" arrow::jemalloc_set_decay_ms"(int ms) + + cdef cppclass CListType" arrow::ListType"(CDataType): + CListType(const shared_ptr[CDataType]& value_type) + CListType(const shared_ptr[CField]& field) + shared_ptr[CDataType] value_type() + shared_ptr[CField] value_field() + + cdef cppclass CLargeListType" arrow::LargeListType"(CDataType): + CLargeListType(const shared_ptr[CDataType]& value_type) + CLargeListType(const shared_ptr[CField]& field) + shared_ptr[CDataType] value_type() + shared_ptr[CField] value_field() + + cdef cppclass CListViewType" arrow::ListViewType"(CDataType): + CListViewType(const shared_ptr[CDataType]& value_type) + CListViewType(const shared_ptr[CField]& field) + shared_ptr[CDataType] value_type() + shared_ptr[CField] value_field() + + cdef cppclass CLargeListViewType" arrow::LargeListViewType"(CDataType): + CLargeListViewType(const shared_ptr[CDataType]& value_type) + CLargeListViewType(const shared_ptr[CField]& field) + shared_ptr[CDataType] value_type() + shared_ptr[CField] value_field() + + cdef cppclass CMapType" arrow::MapType"(CDataType): + CMapType(const shared_ptr[CField]& key_field, + const shared_ptr[CField]& item_field, c_bool keys_sorted) + shared_ptr[CDataType] key_type() + shared_ptr[CField] key_field() + shared_ptr[CDataType] item_type() + shared_ptr[CField] item_field() + c_bool keys_sorted() + + cdef cppclass CFixedSizeListType" arrow::FixedSizeListType"(CDataType): + CFixedSizeListType(const shared_ptr[CDataType]& value_type, + int32_t list_size) + CFixedSizeListType(const shared_ptr[CField]& field, int32_t list_size) + shared_ptr[CDataType] value_type() + shared_ptr[CField] value_field() + int32_t list_size() + + cdef cppclass CStringType" arrow::StringType"(CDataType): + pass + + cdef cppclass CFixedSizeBinaryType \ + " arrow::FixedSizeBinaryType"(CFixedWidthType): + CFixedSizeBinaryType(int byte_width) + int byte_width() + int bit_width() + + cdef cppclass CDecimal32Type \ + " arrow::Decimal32Type"(CFixedSizeBinaryType): + CDecimal32Type(int precision, int scale) + int precision() + int scale() + + cdef cppclass CDecimal64Type \ + " arrow::Decimal64Type"(CFixedSizeBinaryType): + CDecimal64Type(int precision, int scale) + int precision() + int scale() + + cdef cppclass CDecimal128Type \ + " arrow::Decimal128Type"(CFixedSizeBinaryType): + CDecimal128Type(int precision, int scale) + int precision() + int scale() + + cdef cppclass CDecimal256Type \ + " arrow::Decimal256Type"(CFixedSizeBinaryType): + CDecimal256Type(int precision, int scale) + int precision() + int scale() + + cdef cppclass CRunEndEncodedType " arrow::RunEndEncodedType"(CDataType): + CRunEndEncodedType(const shared_ptr[CDataType]& run_end_type, + const shared_ptr[CDataType]& value_type) + const shared_ptr[CDataType]& run_end_type() + const shared_ptr[CDataType]& value_type() + + cdef cppclass CField" arrow::Field": + cppclass CMergeOptions "MergeOptions": + CMergeOptions() + c_bool promote_nullability + + @staticmethod + CMergeOptions Defaults() + + @staticmethod + CMergeOptions Permissive() + + const c_string& name() + shared_ptr[CDataType] type() + c_bool nullable() + + c_string ToString() + c_bool Equals(const CField& other, c_bool check_metadata) + + shared_ptr[const CKeyValueMetadata] metadata() + + CField(const c_string& name, const shared_ptr[CDataType]& type, + c_bool nullable) + + CField(const c_string& name, const shared_ptr[CDataType]& type, + c_bool nullable, const shared_ptr[CKeyValueMetadata]& metadata) + + # Removed const in Cython so don't have to cast to get code to generate + shared_ptr[CField] AddMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + shared_ptr[CField] WithMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + shared_ptr[CField] RemoveMetadata() + shared_ptr[CField] WithType(const shared_ptr[CDataType]& type) + shared_ptr[CField] WithName(const c_string& name) + shared_ptr[CField] WithNullable(c_bool nullable) + vector[shared_ptr[CField]] Flatten() + + cdef cppclass CFieldRef" arrow::FieldRef": + CFieldRef() + CFieldRef(c_string name) + CFieldRef(int index) + CFieldRef(vector[CFieldRef]) + + @staticmethod + CResult[CFieldRef] FromDotPath(c_string& dot_path) + const c_string* name() const + + cdef cppclass CFieldRefHash" arrow::FieldRef::Hash": + pass + + cdef cppclass CStructType" arrow::StructType"(CDataType): + CStructType(const vector[shared_ptr[CField]]& fields) + + shared_ptr[CField] GetFieldByName(const c_string& name) + vector[shared_ptr[CField]] GetAllFieldsByName(const c_string& name) + int GetFieldIndex(const c_string& name) + vector[int] GetAllFieldIndices(const c_string& name) + + cdef cppclass CUnionType" arrow::UnionType"(CDataType): + UnionMode mode() + const vector[int8_t]& type_codes() + const vector[int]& child_ids() + + cdef shared_ptr[CDataType] CMakeSparseUnionType" arrow::sparse_union"( + vector[shared_ptr[CField]] fields, + vector[int8_t] type_codes) + + cdef shared_ptr[CDataType] CMakeDenseUnionType" arrow::dense_union"( + vector[shared_ptr[CField]] fields, + vector[int8_t] type_codes) + + cdef shared_ptr[CDataType] CMakeRunEndEncodedType" arrow::run_end_encoded"( + shared_ptr[CDataType] run_end_type, + shared_ptr[CDataType] value_type) + + cdef shared_ptr[CDataType] CMakeListViewType" arrow::list_view"( + shared_ptr[CField] value_type) + + cdef shared_ptr[CDataType] CMakeLargeListViewType" arrow::large_list_view"( + shared_ptr[CField] value_type) + + cdef cppclass CSchema" arrow::Schema": + CSchema(const vector[shared_ptr[CField]]& fields) + CSchema(const vector[shared_ptr[CField]]& fields, + const shared_ptr[const CKeyValueMetadata]& metadata) + + # Does not actually exist, but gets Cython to not complain + CSchema(const vector[shared_ptr[CField]]& fields, + const shared_ptr[CKeyValueMetadata]& metadata) + + c_bool Equals(const CSchema& other, c_bool check_metadata) + + shared_ptr[CField] field(int i) + shared_ptr[const CKeyValueMetadata] metadata() + shared_ptr[CField] GetFieldByName(const c_string& name) + vector[shared_ptr[CField]] GetAllFieldsByName(const c_string& name) + int GetFieldIndex(const c_string& name) + vector[int] GetAllFieldIndices(const c_string& name) + const vector[shared_ptr[CField]] fields() + int num_fields() + c_string ToString() + + CResult[shared_ptr[CSchema]] AddField(int i, + const shared_ptr[CField]& field) + CResult[shared_ptr[CSchema]] RemoveField(int i) + CResult[shared_ptr[CSchema]] SetField(int i, + const shared_ptr[CField]& field) + + # Removed const in Cython so don't have to cast to get code to generate + shared_ptr[CSchema] AddMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + shared_ptr[CSchema] WithMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + shared_ptr[CSchema] RemoveMetadata() + + CResult[shared_ptr[CSchema]] UnifySchemas( + const vector[shared_ptr[CSchema]]& schemas, + CField.CMergeOptions field_merge_options) + + cdef cppclass PrettyPrintOptions: + PrettyPrintOptions() + PrettyPrintOptions(int indent_arg) + PrettyPrintOptions(int indent_arg, int window_arg) + int indent + int indent_size + int window + int container_window + c_string null_rep + c_bool skip_new_lines + c_bool truncate_metadata + c_bool show_field_metadata + c_bool show_schema_metadata + int element_size_limit + + @staticmethod + PrettyPrintOptions Defaults() + + CStatus PrettyPrint(const CArray& schema, + const PrettyPrintOptions& options, + c_string* result) + CStatus PrettyPrint(const CChunkedArray& schema, + const PrettyPrintOptions& options, + c_string* result) + CStatus PrettyPrint(const CSchema& schema, + const PrettyPrintOptions& options, + c_string* result) + + cdef cppclass CBooleanArray" arrow::BooleanArray"(CArray): + c_bool Value(int i) + int64_t false_count() + int64_t true_count() + + cdef cppclass CUInt8Array" arrow::UInt8Array"(CArray): + uint8_t Value(int i) + + cdef cppclass CInt8Array" arrow::Int8Array"(CArray): + int8_t Value(int i) + + cdef cppclass CUInt16Array" arrow::UInt16Array"(CArray): + uint16_t Value(int i) + + cdef cppclass CInt16Array" arrow::Int16Array"(CArray): + int16_t Value(int i) + + cdef cppclass CUInt32Array" arrow::UInt32Array"(CArray): + uint32_t Value(int i) + + cdef cppclass CInt32Array" arrow::Int32Array"(CArray): + int32_t Value(int i) + + cdef cppclass CUInt64Array" arrow::UInt64Array"(CArray): + uint64_t Value(int i) + + cdef cppclass CInt64Array" arrow::Int64Array"(CArray): + int64_t Value(int i) + + cdef cppclass CDate32Array" arrow::Date32Array"(CArray): + int32_t Value(int i) + + cdef cppclass CDate64Array" arrow::Date64Array"(CArray): + int64_t Value(int i) + + cdef cppclass CTime32Array" arrow::Time32Array"(CArray): + int32_t Value(int i) + + cdef cppclass CTime64Array" arrow::Time64Array"(CArray): + int64_t Value(int i) + + cdef cppclass CTimestampArray" arrow::TimestampArray"(CArray): + int64_t Value(int i) + + cdef cppclass CDurationArray" arrow::DurationArray"(CArray): + int64_t Value(int i) + + cdef cppclass CMonthDayNanoIntervalArray \ + "arrow::MonthDayNanoIntervalArray"(CArray): + pass + + cdef cppclass CHalfFloatArray" arrow::HalfFloatArray"(CArray): + uint16_t Value(int i) + + cdef cppclass CFloatArray" arrow::FloatArray"(CArray): + float Value(int i) + + cdef cppclass CDoubleArray" arrow::DoubleArray"(CArray): + double Value(int i) + + cdef cppclass CFixedSizeBinaryArray" arrow::FixedSizeBinaryArray"(CArray): + const uint8_t* GetValue(int i) + + cdef cppclass CDecimal32Array" arrow::Decimal32Array"( + CFixedSizeBinaryArray + ): + c_string FormatValue(int i) + + cdef cppclass CDecimal64Array" arrow::Decimal64Array"( + CFixedSizeBinaryArray + ): + c_string FormatValue(int i) + + cdef cppclass CDecimal128Array" arrow::Decimal128Array"( + CFixedSizeBinaryArray + ): + c_string FormatValue(int i) + + cdef cppclass CDecimal256Array" arrow::Decimal256Array"( + CFixedSizeBinaryArray + ): + c_string FormatValue(int i) + + cdef cppclass CListArray" arrow::ListArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const CArray& offsets, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + shared_ptr[CDataType], + const CArray& offsets, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + const int32_t* raw_value_offsets() + int32_t value_offset(int i) + int32_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CArray] offsets() + shared_ptr[CDataType] value_type() + + cdef cppclass CLargeListArray" arrow::LargeListArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const CArray& offsets, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap + ) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + shared_ptr[CDataType], + const CArray& offsets, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap + ) + + int64_t value_offset(int i) + int64_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CArray] offsets() + shared_ptr[CDataType] value_type() + + cdef cppclass CFixedSizeListArray" arrow::FixedSizeListArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const shared_ptr[CArray]& values, + int32_t list_size, + shared_ptr[CBuffer] null_bitmap) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + const shared_ptr[CArray]& values, + shared_ptr[CDataType], + shared_ptr[CBuffer] null_bitmap) + + int64_t value_offset(int i) + int64_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CDataType] value_type() + + cdef cppclass CListViewArray" arrow::ListViewArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const CArray& offsets, + const CArray& sizes, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + shared_ptr[CDataType], + const CArray& offsets, + const CArray& sizes, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + CResult[shared_ptr[CArray]] Flatten( + CMemoryPool* pool + ) + + const int32_t* raw_value_offsets() + const int32_t* raw_value_sizes() + int32_t value_offset(int i) + int32_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CArray] offsets() + shared_ptr[CArray] sizes() + shared_ptr[CDataType] value_type() + + cdef cppclass CLargeListViewArray" arrow::LargeListViewArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const CArray& offsets, + const CArray& sizes, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + shared_ptr[CDataType], + const CArray& offsets, + const CArray& sizes, + const CArray& values, + CMemoryPool* pool, + shared_ptr[CBuffer] null_bitmap, + ) + + CResult[shared_ptr[CArray]] Flatten( + CMemoryPool* pool + ) + + int64_t value_offset(int i) + int64_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CArray] offsets() + shared_ptr[CArray] sizes() + shared_ptr[CDataType] value_type() + + cdef cppclass CMapArray" arrow::MapArray"(CArray): + @staticmethod + CResult[shared_ptr[CArray]] FromArrays( + const shared_ptr[CArray]& offsets, + const shared_ptr[CArray]& keys, + const shared_ptr[CArray]& items, + CMemoryPool* pool, + const shared_ptr[CBuffer] null_bitmap, + ) + + @staticmethod + CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"( + shared_ptr[CDataType], + const shared_ptr[CArray]& offsets, + const shared_ptr[CArray]& keys, + const shared_ptr[CArray]& items, + CMemoryPool* pool, + const shared_ptr[CBuffer] null_bitmap, + ) + + shared_ptr[CArray] keys() + shared_ptr[CArray] items() + CMapType* map_type() + int64_t value_offset(int i) + int64_t value_length(int i) + shared_ptr[CArray] values() + shared_ptr[CDataType] value_type() + + cdef cppclass CUnionArray" arrow::UnionArray"(CArray): + shared_ptr[CBuffer] type_codes() + int8_t* raw_type_codes() + int child_id(int64_t index) + shared_ptr[CArray] field(int pos) + const CArray* UnsafeField(int pos) + UnionMode mode() + + cdef cppclass CSparseUnionArray" arrow::SparseUnionArray"(CUnionArray): + @staticmethod + CResult[shared_ptr[CArray]] Make( + const CArray& type_codes, + const vector[shared_ptr[CArray]]& children, + const vector[c_string]& field_names, + const vector[int8_t]& type_codes) + + cdef cppclass CDenseUnionArray" arrow::DenseUnionArray"(CUnionArray): + @staticmethod + CResult[shared_ptr[CArray]] Make( + const CArray& type_codes, + const CArray& value_offsets, + const vector[shared_ptr[CArray]]& children, + const vector[c_string]& field_names, + const vector[int8_t]& type_codes) + + int32_t value_offset(int i) + shared_ptr[CBuffer] value_offsets() + + cdef cppclass CBinaryArray" arrow::BinaryArray"(CArray): + const uint8_t* GetValue(int i, int32_t* length) + shared_ptr[CBuffer] value_data() + int32_t value_offset(int64_t i) + int32_t value_length(int64_t i) + int32_t total_values_length() + + cdef cppclass CLargeBinaryArray" arrow::LargeBinaryArray"(CArray): + const uint8_t* GetValue(int i, int64_t* length) + shared_ptr[CBuffer] value_data() + int64_t value_offset(int64_t i) + int64_t value_length(int64_t i) + int64_t total_values_length() + + cdef cppclass CStringArray" arrow::StringArray"(CBinaryArray): + CStringArray(int64_t length, shared_ptr[CBuffer] value_offsets, + shared_ptr[CBuffer] data, + shared_ptr[CBuffer] null_bitmap, + int64_t null_count, + int64_t offset) + + c_string GetString(int i) + + cdef cppclass CLargeStringArray" arrow::LargeStringArray" \ + (CLargeBinaryArray): + CLargeStringArray(int64_t length, shared_ptr[CBuffer] value_offsets, + shared_ptr[CBuffer] data, + shared_ptr[CBuffer] null_bitmap, + int64_t null_count, + int64_t offset) + + c_string GetString(int i) + + cdef cppclass CStructArray" arrow::StructArray"(CArray): + CStructArray(shared_ptr[CDataType]& type, int64_t length, + vector[shared_ptr[CArray]]& children, + shared_ptr[CBuffer] null_bitmap=nullptr, + int64_t null_count=-1, + int64_t offset=0) + + # XXX Cython crashes if default argument values are declared here + # https://github.com/cython/cython/issues/2167 + @staticmethod + CResult[shared_ptr[CArray]] MakeFromFieldNames "Make"( + vector[shared_ptr[CArray]] children, + vector[c_string] field_names, + shared_ptr[CBuffer] null_bitmap, + int64_t null_count, + int64_t offset) + + @staticmethod + CResult[shared_ptr[CArray]] MakeFromFields "Make"( + vector[shared_ptr[CArray]] children, + vector[shared_ptr[CField]] fields, + shared_ptr[CBuffer] null_bitmap, + int64_t null_count, + int64_t offset) + + shared_ptr[CArray] field(int pos) + shared_ptr[CArray] GetFieldByName(const c_string& name) const + CResult[shared_ptr[CArray]] GetFlattenedField(int index, CMemoryPool* pool) const + + CResult[vector[shared_ptr[CArray]]] Flatten(CMemoryPool* pool) + + cdef cppclass CRunEndEncodedArray" arrow::RunEndEncodedArray"(CArray): + @staticmethod + CResult[shared_ptr[CRunEndEncodedArray]] Make( + const shared_ptr[CDataType]& type, + int64_t logical_length, + const shared_ptr[CArray]& run_ends, + const shared_ptr[CArray]& values, + int64_t logical_offset) + + @staticmethod + CResult[shared_ptr[CRunEndEncodedArray]] MakeFromArrays "Make"( + int64_t logical_length, + const shared_ptr[CArray]& run_ends, + const shared_ptr[CArray]& values, + int64_t logical_offset) + + shared_ptr[CArray]& run_ends() + shared_ptr[CArray]& values() + + int64_t FindPhysicalOffset() + int64_t FindPhysicalLength() + + cdef cppclass CChunkedArray" arrow::ChunkedArray": + CChunkedArray(const vector[shared_ptr[CArray]]& arrays) + CChunkedArray(const vector[shared_ptr[CArray]]& arrays, + const shared_ptr[CDataType]& type) + + @staticmethod + CResult[shared_ptr[CChunkedArray]] Make(vector[shared_ptr[CArray]] chunks, + shared_ptr[CDataType] type) + int64_t length() + int64_t null_count() + int num_chunks() + c_bool Equals(const CChunkedArray& other) + + shared_ptr[CArray] chunk(int i) + shared_ptr[CDataType] type() + CResult[shared_ptr[CScalar]] GetScalar(int64_t index) const + shared_ptr[CChunkedArray] Slice(int64_t offset, int64_t length) const + shared_ptr[CChunkedArray] Slice(int64_t offset) const + + CResult[vector[shared_ptr[CChunkedArray]]] Flatten(CMemoryPool* pool) + + c_bool is_cpu() const + + CStatus Validate() const + CStatus ValidateFull() const + + cdef cppclass CRecordBatch" arrow::RecordBatch": + @staticmethod + shared_ptr[CRecordBatch] Make( + const shared_ptr[CSchema]& schema, int64_t num_rows, + const vector[shared_ptr[CArray]]& columns) + + CResult[shared_ptr[CStructArray]] ToStructArray() const + + @staticmethod + CResult[shared_ptr[CRecordBatch]] FromStructArray( + const shared_ptr[CArray]& array) + + c_bool Equals(const CRecordBatch& other, c_bool check_metadata) + + shared_ptr[CSchema] schema() + shared_ptr[CArray] column(int i) + const c_string& column_name(int i) + + CResult[shared_ptr[CRecordBatch]] AddColumn( + int i, shared_ptr[CField] field, shared_ptr[CArray] column) + CResult[shared_ptr[CRecordBatch]] RemoveColumn(int i) + CResult[shared_ptr[CRecordBatch]] SetColumn( + int i, shared_ptr[CField] field, shared_ptr[CArray] column) + + const vector[shared_ptr[CArray]]& columns() + + CResult[shared_ptr[CRecordBatch]] RenameColumns(const vector[c_string]&) + CResult[shared_ptr[CRecordBatch]] SelectColumns(const vector[int]&) + + int num_columns() + int64_t num_rows() + + CDeviceAllocationType device_type() + + CStatus Validate() const + CStatus ValidateFull() const + + shared_ptr[CRecordBatch] ReplaceSchemaMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + + shared_ptr[CRecordBatch] Slice(int64_t offset) + shared_ptr[CRecordBatch] Slice(int64_t offset, int64_t length) + + CResult[shared_ptr[CRecordBatch]] CopyTo(const shared_ptr[CMemoryManager]& to) const + + CResult[shared_ptr[CTensor]] ToTensor(c_bool null_to_nan, c_bool row_major, + CMemoryPool* pool) const + + cdef cppclass CRecordBatchWithMetadata" arrow::RecordBatchWithMetadata": + shared_ptr[CRecordBatch] batch + # The struct in C++ does not actually have these two `const` qualifiers, but + # adding `const` gets Cython to not complain + shared_ptr[const CKeyValueMetadata] custom_metadata + + cdef cppclass CTable" arrow::Table": + CTable(const shared_ptr[CSchema]& schema, + const vector[shared_ptr[CChunkedArray]]& columns) + + @staticmethod + shared_ptr[CTable] Make( + const shared_ptr[CSchema]& schema, + const vector[shared_ptr[CChunkedArray]]& columns) + + @staticmethod + shared_ptr[CTable] MakeWithRows "Make"( + const shared_ptr[CSchema]& schema, + const vector[shared_ptr[CChunkedArray]]& columns, + int64_t num_rows) + + @staticmethod + shared_ptr[CTable] MakeFromArrays" Make"( + const shared_ptr[CSchema]& schema, + const vector[shared_ptr[CArray]]& arrays) + + @staticmethod + CResult[shared_ptr[CTable]] FromRecordBatchReader( + CRecordBatchReader *reader) + + @staticmethod + CResult[shared_ptr[CTable]] FromRecordBatches( + const shared_ptr[CSchema]& schema, + const vector[shared_ptr[CRecordBatch]]& batches) + + int num_columns() + int64_t num_rows() + + c_bool Equals(const CTable& other, c_bool check_metadata) + + shared_ptr[CSchema] schema() + shared_ptr[CChunkedArray] column(int i) + shared_ptr[CField] field(int i) + + CResult[shared_ptr[CTable]] AddColumn( + int i, shared_ptr[CField] field, shared_ptr[CChunkedArray] column) + CResult[shared_ptr[CTable]] RemoveColumn(int i) + CResult[shared_ptr[CTable]] SetColumn( + int i, shared_ptr[CField] field, shared_ptr[CChunkedArray] column) + + vector[c_string] ColumnNames() + CResult[shared_ptr[CTable]] RenameColumns(const vector[c_string]&) + CResult[shared_ptr[CTable]] SelectColumns(const vector[int]&) + + CResult[shared_ptr[CTable]] Flatten(CMemoryPool* pool) + + CResult[shared_ptr[CTable]] CombineChunks(CMemoryPool* pool) + + CStatus Validate() const + CStatus ValidateFull() const + + shared_ptr[CTable] ReplaceSchemaMetadata( + const shared_ptr[CKeyValueMetadata]& metadata) + + shared_ptr[CTable] Slice(int64_t offset) + shared_ptr[CTable] Slice(int64_t offset, int64_t length) + + cdef cppclass CRecordBatchReader" arrow::RecordBatchReader": + shared_ptr[CSchema] schema() + CStatus Close() + CResult[CRecordBatchWithMetadata] ReadNext() + CStatus ReadNext(shared_ptr[CRecordBatch]* batch) + CResult[shared_ptr[CTable]] ToTable() + + cdef cppclass TableBatchReader(CRecordBatchReader): + TableBatchReader(const CTable& table) + TableBatchReader(shared_ptr[CTable] table) + void set_chunksize(int64_t chunksize) + + cdef cppclass CTensor" arrow::Tensor": + shared_ptr[CDataType] type() + shared_ptr[CBuffer] data() + + const vector[int64_t]& shape() + const vector[int64_t]& strides() + int64_t size() + + int ndim() + const vector[c_string]& dim_names() + const c_string& dim_name(int i) + + c_bool is_mutable() + c_bool is_contiguous() + Type type_id() + c_bool Equals(const CTensor& other) + + cdef cppclass CSparseIndex" arrow::SparseIndex": + pass + + cdef cppclass CSparseCOOIndex" arrow::SparseCOOIndex": + c_bool is_canonical() + + cdef cppclass CSparseCOOTensor" arrow::SparseCOOTensor": + shared_ptr[CDataType] type() + shared_ptr[CBuffer] data() + CResult[shared_ptr[CTensor]] ToTensor() + + shared_ptr[CSparseIndex] sparse_index() + + const vector[int64_t]& shape() + int64_t size() + int64_t non_zero_length() + + int ndim() + const vector[c_string]& dim_names() + const c_string& dim_name(int i) + + c_bool is_mutable() + Type type_id() + c_bool Equals(const CSparseCOOTensor& other) + + cdef cppclass CSparseCSRMatrix" arrow::SparseCSRMatrix": + shared_ptr[CDataType] type() + shared_ptr[CBuffer] data() + CResult[shared_ptr[CTensor]] ToTensor() + + const vector[int64_t]& shape() + int64_t size() + int64_t non_zero_length() + + int ndim() + const vector[c_string]& dim_names() + const c_string& dim_name(int i) + + c_bool is_mutable() + Type type_id() + c_bool Equals(const CSparseCSRMatrix& other) + + cdef cppclass CSparseCSCMatrix" arrow::SparseCSCMatrix": + shared_ptr[CDataType] type() + shared_ptr[CBuffer] data() + CResult[shared_ptr[CTensor]] ToTensor() + + const vector[int64_t]& shape() + int64_t size() + int64_t non_zero_length() + + int ndim() + const vector[c_string]& dim_names() + const c_string& dim_name(int i) + + c_bool is_mutable() + Type type_id() + c_bool Equals(const CSparseCSCMatrix& other) + + cdef cppclass CSparseCSFTensor" arrow::SparseCSFTensor": + shared_ptr[CDataType] type() + shared_ptr[CBuffer] data() + CResult[shared_ptr[CTensor]] ToTensor() + + const vector[int64_t]& shape() + int64_t size() + int64_t non_zero_length() + + int ndim() + const vector[c_string]& dim_names() + const c_string& dim_name(int i) + + c_bool is_mutable() + Type type_id() + c_bool Equals(const CSparseCSFTensor& other) + + cdef cppclass CScalar" arrow::Scalar": + CScalar(shared_ptr[CDataType]) + + shared_ptr[CDataType] type + c_bool is_valid + + c_string ToString() const + c_bool Equals(const CScalar& other) const + CStatus Validate() const + CStatus ValidateFull() const + + cdef cppclass CScalarHash" arrow::Scalar::Hash": + size_t operator()(const shared_ptr[CScalar]& scalar) const + + cdef cppclass CNullScalar" arrow::NullScalar"(CScalar): + CNullScalar() + + cdef cppclass CBooleanScalar" arrow::BooleanScalar"(CScalar): + CBooleanScalar(c_bool value) + c_bool value + + cdef cppclass CInt8Scalar" arrow::Int8Scalar"(CScalar): + int8_t value + + cdef cppclass CUInt8Scalar" arrow::UInt8Scalar"(CScalar): + uint8_t value + + cdef cppclass CInt16Scalar" arrow::Int16Scalar"(CScalar): + int16_t value + + cdef cppclass CUInt16Scalar" arrow::UInt16Scalar"(CScalar): + uint16_t value + + cdef cppclass CInt32Scalar" arrow::Int32Scalar"(CScalar): + int32_t value + + cdef cppclass CUInt32Scalar" arrow::UInt32Scalar"(CScalar): + uint32_t value + + cdef cppclass CInt64Scalar" arrow::Int64Scalar"(CScalar): + int64_t value + + cdef cppclass CUInt64Scalar" arrow::UInt64Scalar"(CScalar): + uint64_t value + + cdef cppclass CHalfFloatScalar" arrow::HalfFloatScalar"(CScalar): + uint16_t value + + cdef cppclass CFloatScalar" arrow::FloatScalar"(CScalar): + float value + + cdef cppclass CDoubleScalar" arrow::DoubleScalar"(CScalar): + double value + + cdef cppclass CDecimal32Scalar" arrow::Decimal32Scalar"(CScalar): + CDecimal32 value + + cdef cppclass CDecimal64Scalar" arrow::Decimal64Scalar"(CScalar): + CDecimal64 value + + cdef cppclass CDecimal128Scalar" arrow::Decimal128Scalar"(CScalar): + CDecimal128 value + + cdef cppclass CDecimal256Scalar" arrow::Decimal256Scalar"(CScalar): + CDecimal256 value + + cdef cppclass CDate32Scalar" arrow::Date32Scalar"(CScalar): + int32_t value + + cdef cppclass CDate64Scalar" arrow::Date64Scalar"(CScalar): + int64_t value + + cdef cppclass CTime32Scalar" arrow::Time32Scalar"(CScalar): + int32_t value + + cdef cppclass CTime64Scalar" arrow::Time64Scalar"(CScalar): + int64_t value + + cdef cppclass CTimestampScalar" arrow::TimestampScalar"(CScalar): + int64_t value + + cdef cppclass CDurationScalar" arrow::DurationScalar"(CScalar): + int64_t value + + cdef cppclass CMonthDayNanoIntervalScalar \ + "arrow::MonthDayNanoIntervalScalar"(CScalar): + pass + + cdef cppclass CBaseBinaryScalar" arrow::BaseBinaryScalar"(CScalar): + shared_ptr[CBuffer] value + + cdef cppclass CBaseListScalar" arrow::BaseListScalar"(CScalar): + shared_ptr[CArray] value + + cdef cppclass CListScalar" arrow::ListScalar"(CBaseListScalar): + pass + + cdef cppclass CListViewScalar" arrow::ListViewScalar"(CBaseListScalar): + pass + + cdef cppclass CLargeListViewScalar" arrow::LargeListViewScalar"(CBaseListScalar): + pass + + cdef cppclass CMapScalar" arrow::MapScalar"(CListScalar): + pass + + cdef cppclass CStructScalar" arrow::StructScalar"(CScalar): + vector[shared_ptr[CScalar]] value + CResult[shared_ptr[CScalar]] field(CFieldRef ref) const + + cdef cppclass CDictionaryScalarIndexAndDictionary \ + "arrow::DictionaryScalar::ValueType": + shared_ptr[CScalar] index + shared_ptr[CArray] dictionary + + cdef cppclass CDictionaryScalar" arrow::DictionaryScalar"(CScalar): + CDictionaryScalar(CDictionaryScalarIndexAndDictionary value, + shared_ptr[CDataType], c_bool is_valid) + CDictionaryScalarIndexAndDictionary value + + CResult[shared_ptr[CScalar]] GetEncodedValue() + + cdef cppclass CUnionScalar" arrow::UnionScalar"(CScalar): + int8_t type_code + + cdef cppclass CDenseUnionScalar" arrow::DenseUnionScalar"(CUnionScalar): + shared_ptr[CScalar] value + + cdef cppclass CSparseUnionScalar" arrow::SparseUnionScalar"(CUnionScalar): + vector[shared_ptr[CScalar]] value + int child_id + + cdef cppclass CRunEndEncodedScalar" arrow::RunEndEncodedScalar"(CScalar): + shared_ptr[CScalar] value + + cdef cppclass CExtensionScalar" arrow::ExtensionScalar"(CScalar): + CExtensionScalar(shared_ptr[CScalar] storage, + shared_ptr[CDataType], c_bool is_valid) + shared_ptr[CScalar] value + + shared_ptr[CScalar] MakeScalar[Value](Value value) + + cdef cppclass CConcatenateTablesOptions" arrow::ConcatenateTablesOptions": + c_bool unify_schemas + CField.CMergeOptions field_merge_options + + @staticmethod + CConcatenateTablesOptions Defaults() + + CResult[shared_ptr[CTable]] ConcatenateTables( + const vector[shared_ptr[CTable]]& tables, + CConcatenateTablesOptions options, + CMemoryPool* memory_pool) + + CResult[shared_ptr[CRecordBatch]] ConcatenateRecordBatches( + const vector[shared_ptr[CRecordBatch]]& batches, + CMemoryPool* memory_pool) + + cdef cppclass CDictionaryUnifier" arrow::DictionaryUnifier": + @staticmethod + CResult[shared_ptr[CChunkedArray]] UnifyChunkedArray( + shared_ptr[CChunkedArray] array, CMemoryPool* pool) + + @staticmethod + CResult[shared_ptr[CTable]] UnifyTable( + const CTable& table, CMemoryPool* pool) + + shared_ptr[CScalar] MakeNullScalar(shared_ptr[CDataType] type) + + +cdef extern from "arrow/c/dlpack_abi.h" nogil: + ctypedef enum DLDeviceType: + kDLCPU = 1 + + ctypedef struct DLDevice: + DLDeviceType device_type + int32_t device_id + + ctypedef struct DLManagedTensor: + void (*deleter)(DLManagedTensor*) + + +cdef extern from "arrow/c/dlpack.h" namespace "arrow::dlpack" nogil: + CResult[DLManagedTensor*] ExportArrayToDLPack" arrow::dlpack::ExportArray"( + const shared_ptr[CArray]& arr) + CResult[DLManagedTensor*] ExportTensorToDLPack" arrow::dlpack::ExportTensor"( + const shared_ptr[CTensor]& tensor) + + CResult[DLDevice] ExportDevice(const shared_ptr[CArray]& arr) + CResult[DLDevice] ExportDevice(const shared_ptr[CTensor]& tensor) + + +cdef extern from "arrow/builder.h" namespace "arrow" nogil: + + cdef cppclass CArrayBuilder" arrow::ArrayBuilder": + CArrayBuilder(shared_ptr[CDataType], CMemoryPool* pool) + + int64_t length() + int64_t null_count() + CStatus AppendNull() + CStatus Finish(shared_ptr[CArray]* out) + CStatus Reserve(int64_t additional_capacity) + + cdef cppclass CBooleanBuilder" arrow::BooleanBuilder"(CArrayBuilder): + CBooleanBuilder(CMemoryPool* pool) + CStatus Append(const c_bool val) + CStatus Append(const uint8_t val) + + cdef cppclass CInt8Builder" arrow::Int8Builder"(CArrayBuilder): + CInt8Builder(CMemoryPool* pool) + CStatus Append(const int8_t value) + + cdef cppclass CInt16Builder" arrow::Int16Builder"(CArrayBuilder): + CInt16Builder(CMemoryPool* pool) + CStatus Append(const int16_t value) + + cdef cppclass CInt32Builder" arrow::Int32Builder"(CArrayBuilder): + CInt32Builder(CMemoryPool* pool) + CStatus Append(const int32_t value) + + cdef cppclass CInt64Builder" arrow::Int64Builder"(CArrayBuilder): + CInt64Builder(CMemoryPool* pool) + CStatus Append(const int64_t value) + + cdef cppclass CUInt8Builder" arrow::UInt8Builder"(CArrayBuilder): + CUInt8Builder(CMemoryPool* pool) + CStatus Append(const uint8_t value) + + cdef cppclass CUInt16Builder" arrow::UInt16Builder"(CArrayBuilder): + CUInt16Builder(CMemoryPool* pool) + CStatus Append(const uint16_t value) + + cdef cppclass CUInt32Builder" arrow::UInt32Builder"(CArrayBuilder): + CUInt32Builder(CMemoryPool* pool) + CStatus Append(const uint32_t value) + + cdef cppclass CUInt64Builder" arrow::UInt64Builder"(CArrayBuilder): + CUInt64Builder(CMemoryPool* pool) + CStatus Append(const uint64_t value) + + cdef cppclass CHalfFloatBuilder" arrow::HalfFloatBuilder"(CArrayBuilder): + CHalfFloatBuilder(CMemoryPool* pool) + + cdef cppclass CFloatBuilder" arrow::FloatBuilder"(CArrayBuilder): + CFloatBuilder(CMemoryPool* pool) + CStatus Append(const float value) + + cdef cppclass CDoubleBuilder" arrow::DoubleBuilder"(CArrayBuilder): + CDoubleBuilder(CMemoryPool* pool) + CStatus Append(const double value) + + cdef cppclass CBinaryBuilder" arrow::BinaryBuilder"(CArrayBuilder): + CArrayBuilder(shared_ptr[CDataType], CMemoryPool* pool) + CStatus Append(const char* value, int32_t length) + + cdef cppclass CStringBuilder" arrow::StringBuilder"(CBinaryBuilder): + CStringBuilder(CMemoryPool* pool) + CStatus Append(const c_string& value) + + cdef cppclass CBinaryViewBuilder" arrow::BinaryViewBuilder"(CArrayBuilder): + CBinaryViewBuilder(shared_ptr[CDataType], CMemoryPool* pool) + CStatus Append(const char* value, int32_t length) + + cdef cppclass CStringViewBuilder" arrow::StringViewBuilder"(CBinaryViewBuilder): + CStringViewBuilder(CMemoryPool* pool) + CStatus Append(const c_string& value) + + cdef cppclass CTimestampBuilder "arrow::TimestampBuilder"(CArrayBuilder): + CTimestampBuilder(const shared_ptr[CDataType] typ, CMemoryPool* pool) + CStatus Append(const int64_t value) + + cdef cppclass CDate32Builder "arrow::Date32Builder"(CArrayBuilder): + CDate32Builder(CMemoryPool* pool) + CStatus Append(const int32_t value) + + cdef cppclass CDate64Builder "arrow::Date64Builder"(CArrayBuilder): + CDate64Builder(CMemoryPool* pool) + CStatus Append(const int64_t value) + + +# Use typedef to emulate syntax for std::function +ctypedef void CallbackTransform(object, const shared_ptr[CBuffer]& src, + shared_ptr[CBuffer]* dest) + +ctypedef CResult[shared_ptr[CInputStream]] StreamWrapFunc( + shared_ptr[CInputStream]) + + +cdef extern from "arrow/util/cancel.h" namespace "arrow" nogil: + cdef cppclass CStopToken "arrow::StopToken": + CStatus Poll() + c_bool IsStopRequested() + + cdef cppclass CStopSource "arrow::StopSource": + CStopToken token() + + CResult[CStopSource*] SetSignalStopSource() + void ResetSignalStopSource() + + CStatus RegisterCancellingSignalHandler(vector[int] signals) + void UnregisterCancellingSignalHandler() + + +cdef extern from "arrow/io/api.h" namespace "arrow::io" nogil: + cdef enum FileMode" arrow::io::FileMode::type": + FileMode_READ" arrow::io::FileMode::READ" + FileMode_WRITE" arrow::io::FileMode::WRITE" + FileMode_READWRITE" arrow::io::FileMode::READWRITE" + + cdef enum ObjectType" arrow::io::ObjectType::type": + ObjectType_FILE" arrow::io::ObjectType::FILE" + ObjectType_DIRECTORY" arrow::io::ObjectType::DIRECTORY" + + cdef cppclass CIOContext" arrow::io::IOContext": + CIOContext() + CIOContext(CStopToken) + CIOContext(CMemoryPool*) + CIOContext(CMemoryPool*, CStopToken) + + CIOContext c_default_io_context "arrow::io::default_io_context"() + int GetIOThreadPoolCapacity() + CStatus SetIOThreadPoolCapacity(int threads) + + cdef cppclass FileStatistics: + int64_t size + ObjectType kind + + cdef cppclass FileInterface: + CStatus Close() + CResult[int64_t] Tell() + FileMode mode() + c_bool closed() + + cdef cppclass Readable: + # put overload under a different name to avoid cython bug with multiple + # layers of inheritance + CResult[shared_ptr[CBuffer]] ReadBuffer" Read"(int64_t nbytes) + CResult[int64_t] Read(int64_t nbytes, uint8_t* out) + + cdef cppclass Seekable: + CStatus Seek(int64_t position) + + cdef cppclass Writable: + CStatus WriteBuffer" Write"(shared_ptr[CBuffer] data) + CStatus Write(const uint8_t* data, int64_t nbytes) + CStatus Flush() + + cdef cppclass CCacheOptions "arrow::io::CacheOptions": + int64_t hole_size_limit + int64_t range_size_limit + c_bool lazy + int64_t prefetch_limit + c_bool Equals "operator==" (CCacheOptions other) + + @staticmethod + CCacheOptions MakeFromNetworkMetrics(int64_t time_to_first_byte_millis, + int64_t transfer_bandwidth_mib_per_sec, + double ideal_bandwidth_utilization_frac, + int64_t max_ideal_request_size_mib) + + @staticmethod + CCacheOptions LazyDefaults() + + cdef cppclass COutputStream" arrow::io::OutputStream"(FileInterface, + Writable): + pass + + cdef cppclass CInputStream" arrow::io::InputStream"(FileInterface, + Readable): + CResult[shared_ptr[const CKeyValueMetadata]] ReadMetadata() + + cdef cppclass CRandomAccessFile" arrow::io::RandomAccessFile"(CInputStream, + Seekable): + CResult[int64_t] GetSize() + + @staticmethod + CResult[shared_ptr[CInputStream]] GetStream( + shared_ptr[CRandomAccessFile] file, + int64_t file_offset, + int64_t nbytes) + + CResult[int64_t] ReadAt(int64_t position, int64_t nbytes, + uint8_t* buffer) + CResult[shared_ptr[CBuffer]] ReadAt(int64_t position, int64_t nbytes) + c_bool supports_zero_copy() + + cdef cppclass WritableFile(COutputStream, Seekable): + CStatus WriteAt(int64_t position, const uint8_t* data, + int64_t nbytes) + + cdef cppclass ReadWriteFileInterface(CRandomAccessFile, + WritableFile): + pass + + cdef cppclass CIOFileSystem" arrow::io::FileSystem": + CStatus Stat(const c_string& path, FileStatistics* stat) + + cdef cppclass FileOutputStream(COutputStream): + @staticmethod + CResult[shared_ptr[COutputStream]] Open(const c_string& path) + + @staticmethod + CResult[shared_ptr[COutputStream]] OpenWithAppend" Open"( + const c_string& path, c_bool append) + + int file_descriptor() + + cdef cppclass ReadableFile(CRandomAccessFile): + @staticmethod + CResult[shared_ptr[ReadableFile]] Open(const c_string& path) + + @staticmethod + CResult[shared_ptr[ReadableFile]] Open(const c_string& path, + CMemoryPool* memory_pool) + + int file_descriptor() + + cdef cppclass CMemoryMappedFile \ + " arrow::io::MemoryMappedFile"(ReadWriteFileInterface): + + @staticmethod + CResult[shared_ptr[CMemoryMappedFile]] Create(const c_string& path, + int64_t size) + + @staticmethod + CResult[shared_ptr[CMemoryMappedFile]] Open(const c_string& path, + FileMode mode) + + CStatus Resize(int64_t size) + + int file_descriptor() + + cdef cppclass CCompressedInputStream \ + " arrow::io::CompressedInputStream"(CInputStream): + @staticmethod + CResult[shared_ptr[CCompressedInputStream]] Make( + CCodec* codec, shared_ptr[CInputStream] raw) + + cdef cppclass CCompressedOutputStream \ + " arrow::io::CompressedOutputStream"(COutputStream): + @staticmethod + CResult[shared_ptr[CCompressedOutputStream]] Make( + CCodec* codec, shared_ptr[COutputStream] raw) + + cdef cppclass CBufferedInputStream \ + " arrow::io::BufferedInputStream"(CInputStream): + + @staticmethod + CResult[shared_ptr[CBufferedInputStream]] Create( + int64_t buffer_size, CMemoryPool* pool, + shared_ptr[CInputStream] raw) + + CResult[shared_ptr[CInputStream]] Detach() + + cdef cppclass CBufferedOutputStream \ + " arrow::io::BufferedOutputStream"(COutputStream): + + @staticmethod + CResult[shared_ptr[CBufferedOutputStream]] Create( + int64_t buffer_size, CMemoryPool* pool, + shared_ptr[COutputStream] raw) + + CResult[shared_ptr[COutputStream]] Detach() + + cdef cppclass CTransformInputStreamVTable \ + "arrow::py::TransformInputStreamVTable": + CTransformInputStreamVTable() + function[CallbackTransform] transform + + shared_ptr[CInputStream] MakeTransformInputStream \ + "arrow::py::MakeTransformInputStream"( + shared_ptr[CInputStream] wrapped, CTransformInputStreamVTable vtable, + object method_arg) + + shared_ptr[function[StreamWrapFunc]] MakeStreamTransformFunc \ + "arrow::py::MakeStreamTransformFunc"( + CTransformInputStreamVTable vtable, + object method_arg) + + # ---------------------------------------------------------------------- + # HDFS + + CStatus HaveLibHdfs() + CStatus HaveLibHdfs3() + + cdef enum HdfsDriver" arrow::io::HdfsDriver": + HdfsDriver_LIBHDFS" arrow::io::HdfsDriver::LIBHDFS" + HdfsDriver_LIBHDFS3" arrow::io::HdfsDriver::LIBHDFS3" + + cdef cppclass HdfsConnectionConfig: + c_string host + int port + c_string user + c_string kerb_ticket + unordered_map[c_string, c_string] extra_conf + HdfsDriver driver + + cdef cppclass HdfsPathInfo: + ObjectType kind + c_string name + c_string owner + c_string group + int32_t last_modified_time + int32_t last_access_time + int64_t size + int16_t replication + int64_t block_size + int16_t permissions + + cdef cppclass HdfsReadableFile(CRandomAccessFile): + pass + + cdef cppclass HdfsOutputStream(COutputStream): + pass + + cdef cppclass CIOHadoopFileSystem \ + "arrow::io::HadoopFileSystem"(CIOFileSystem): + @staticmethod + CStatus Connect(const HdfsConnectionConfig* config, + shared_ptr[CIOHadoopFileSystem]* client) + + CStatus MakeDirectory(const c_string& path) + + CStatus Delete(const c_string& path, c_bool recursive) + + CStatus Disconnect() + + c_bool Exists(const c_string& path) + + CStatus Chmod(const c_string& path, int mode) + CStatus Chown(const c_string& path, const char* owner, + const char* group) + + CStatus GetCapacity(int64_t* nbytes) + CStatus GetUsed(int64_t* nbytes) + + CStatus ListDirectory(const c_string& path, + vector[HdfsPathInfo]* listing) + + CStatus GetPathInfo(const c_string& path, HdfsPathInfo* info) + + CStatus Rename(const c_string& src, const c_string& dst) + + CStatus OpenReadable(const c_string& path, + shared_ptr[HdfsReadableFile]* handle) + + CStatus OpenWritable(const c_string& path, c_bool append, + int32_t buffer_size, int16_t replication, + int64_t default_block_size, + shared_ptr[HdfsOutputStream]* handle) + + cdef cppclass CBufferReader \ + " arrow::io::BufferReader"(CRandomAccessFile): + CBufferReader(const shared_ptr[CBuffer]& buffer) + CBufferReader(const uint8_t* data, int64_t nbytes) + + cdef cppclass CBufferOutputStream \ + " arrow::io::BufferOutputStream"(COutputStream): + CBufferOutputStream(const shared_ptr[CResizableBuffer]& buffer) + + cdef cppclass CMockOutputStream \ + " arrow::io::MockOutputStream"(COutputStream): + CMockOutputStream() + int64_t GetExtentBytesWritten() + + cdef cppclass CFixedSizeBufferWriter \ + " arrow::io::FixedSizeBufferWriter"(WritableFile): + CFixedSizeBufferWriter(const shared_ptr[CBuffer]& buffer) + + void set_memcopy_threads(int num_threads) + void set_memcopy_blocksize(int64_t blocksize) + void set_memcopy_threshold(int64_t threshold) + + +cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil: + cdef enum MessageType" arrow::ipc::MessageType": + MessageType_SCHEMA" arrow::ipc::MessageType::SCHEMA" + MessageType_RECORD_BATCH" arrow::ipc::MessageType::RECORD_BATCH" + MessageType_DICTIONARY_BATCH \ + " arrow::ipc::MessageType::DICTIONARY_BATCH" + + # TODO: use "cpdef enum class" to automatically get a Python wrapper? + # See + # https://github.com/cython/cython/commit/2c7c22f51405299a4e247f78edf52957d30cf71d#diff-61c1365c0f761a8137754bb3a73bfbf7 + ctypedef enum CMetadataVersion" arrow::ipc::MetadataVersion": + CMetadataVersion_V1" arrow::ipc::MetadataVersion::V1" + CMetadataVersion_V2" arrow::ipc::MetadataVersion::V2" + CMetadataVersion_V3" arrow::ipc::MetadataVersion::V3" + CMetadataVersion_V4" arrow::ipc::MetadataVersion::V4" + CMetadataVersion_V5" arrow::ipc::MetadataVersion::V5" + + cdef cppclass CIpcWriteOptions" arrow::ipc::IpcWriteOptions": + c_bool allow_64bit + int max_recursion_depth + int32_t alignment + c_bool write_legacy_ipc_format + CMemoryPool* memory_pool + CMetadataVersion metadata_version + shared_ptr[CCodec] codec + c_bool use_threads + c_bool emit_dictionary_deltas + c_bool unify_dictionaries + + CIpcWriteOptions() + CIpcWriteOptions(CIpcWriteOptions) + + @staticmethod + CIpcWriteOptions Defaults() + + ctypedef enum CAlignment" arrow::ipc::Alignment": + CAlignment_Any" arrow::ipc::Alignment::kAnyAlignment" + CAlignment_DataTypeSpecific" arrow::ipc::Alignment::kDataTypeSpecificAlignment" + CAlignment_64Byte" arrow::ipc::Alignment::k64ByteAlignment" + + cdef cppclass CIpcReadOptions" arrow::ipc::IpcReadOptions": + int max_recursion_depth + CMemoryPool* memory_pool + vector[int] included_fields + c_bool use_threads + c_bool ensure_native_endian + CAlignment ensure_alignment + + @staticmethod + CIpcReadOptions Defaults() + + cdef cppclass CIpcWriteStats" arrow::ipc::WriteStats": + int64_t num_messages + int64_t num_record_batches + int64_t num_dictionary_batches + int64_t num_dictionary_deltas + int64_t num_replaced_dictionaries + + cdef cppclass CIpcReadStats" arrow::ipc::ReadStats": + int64_t num_messages + int64_t num_record_batches + int64_t num_dictionary_batches + int64_t num_dictionary_deltas + int64_t num_replaced_dictionaries + + cdef cppclass CDictionaryMemo" arrow::ipc::DictionaryMemo": + pass + + cdef cppclass CIpcPayload" arrow::ipc::IpcPayload": + MessageType type + shared_ptr[CBuffer] metadata + vector[shared_ptr[CBuffer]] body_buffers + int64_t body_length + + cdef cppclass CMessage" arrow::ipc::Message": + CResult[unique_ptr[CMessage]] Open(shared_ptr[CBuffer] metadata, + shared_ptr[CBuffer] body) + + shared_ptr[CBuffer] body() + + c_bool Equals(const CMessage& other) + + shared_ptr[CBuffer] metadata() + CMetadataVersion metadata_version() + MessageType type() + + CStatus SerializeTo(COutputStream* stream, + const CIpcWriteOptions& options, + int64_t* output_length) + + c_string FormatMessageType(MessageType type) + + cdef cppclass CMessageReader" arrow::ipc::MessageReader": + @staticmethod + unique_ptr[CMessageReader] Open(const shared_ptr[CInputStream]& stream) + + CResult[unique_ptr[CMessage]] ReadNextMessage() + + cdef cppclass CRecordBatchWriter" arrow::ipc::RecordBatchWriter": + CStatus Close() + CStatus WriteRecordBatch(const CRecordBatch& batch) + CStatus WriteRecordBatch( + const CRecordBatch& batch, + const shared_ptr[const CKeyValueMetadata]& metadata) + CStatus WriteTable(const CTable& table, int64_t max_chunksize) + + CIpcWriteStats stats() + + cdef cppclass CRecordBatchStreamReader \ + " arrow::ipc::RecordBatchStreamReader"(CRecordBatchReader): + @staticmethod + CResult[shared_ptr[CRecordBatchReader]] Open( + const shared_ptr[CInputStream], const CIpcReadOptions&) + + @staticmethod + CResult[shared_ptr[CRecordBatchReader]] Open2" Open"( + unique_ptr[CMessageReader] message_reader, + const CIpcReadOptions& options) + + CIpcReadStats stats() + + cdef cppclass CRecordBatchFileReader \ + " arrow::ipc::RecordBatchFileReader": + @staticmethod + CResult[shared_ptr[CRecordBatchFileReader]] Open( + CRandomAccessFile* file, + const CIpcReadOptions& options) + + @staticmethod + CResult[shared_ptr[CRecordBatchFileReader]] Open2" Open"( + CRandomAccessFile* file, int64_t footer_offset, + const CIpcReadOptions& options) + + shared_ptr[CSchema] schema() + + int num_record_batches() + + CResult[shared_ptr[CRecordBatch]] ReadRecordBatch(int i) + + CResult[CRecordBatchWithMetadata] ReadRecordBatchWithCustomMetadata(int i) + + CIpcReadStats stats() + + shared_ptr[const CKeyValueMetadata] metadata() + + CResult[shared_ptr[CRecordBatchWriter]] MakeStreamWriter( + shared_ptr[COutputStream] sink, const shared_ptr[CSchema]& schema, + CIpcWriteOptions& options) + + CResult[shared_ptr[CRecordBatchWriter]] MakeFileWriter( + shared_ptr[COutputStream] sink, const shared_ptr[CSchema]& schema, + CIpcWriteOptions& options, shared_ptr[const CKeyValueMetadata] metadata) + + CResult[unique_ptr[CMessage]] ReadMessage(CInputStream* stream, + CMemoryPool* pool) + + CStatus GetRecordBatchSize(const CRecordBatch& batch, int64_t* size) + CStatus GetTensorSize(const CTensor& tensor, int64_t* size) + + CStatus WriteTensor(const CTensor& tensor, COutputStream* dst, + int32_t* metadata_length, + int64_t* body_length) + + CResult[shared_ptr[CTensor]] ReadTensor(CInputStream* stream) + + CResult[shared_ptr[CRecordBatch]] ReadRecordBatch( + const CMessage& message, const shared_ptr[CSchema]& schema, + CDictionaryMemo* dictionary_memo, + const CIpcReadOptions& options) + + CResult[shared_ptr[CBuffer]] SerializeSchema( + const CSchema& schema, CMemoryPool* pool) + + CResult[shared_ptr[CBuffer]] SerializeRecordBatch( + const CRecordBatch& schema, const CIpcWriteOptions& options) + + CResult[shared_ptr[CSchema]] ReadSchema(const CMessage& message, + CDictionaryMemo* dictionary_memo) + + CResult[shared_ptr[CSchema]] ReadSchema(CInputStream* stream, + CDictionaryMemo* dictionary_memo) + + CResult[shared_ptr[CRecordBatch]] ReadRecordBatch( + const shared_ptr[CSchema]& schema, + CDictionaryMemo* dictionary_memo, + const CIpcReadOptions& options, + CInputStream* stream) + + CStatus AlignStream(CInputStream* stream, int64_t alignment) + CStatus AlignStream(COutputStream* stream, int64_t alignment) + + cdef CStatus GetRecordBatchPayload \ + " arrow::ipc::GetRecordBatchPayload"( + const CRecordBatch& batch, + const CIpcWriteOptions& options, + CIpcPayload* out) + + +cdef extern from "arrow/util/value_parsing.h" namespace "arrow" nogil: + cdef cppclass CTimestampParser" arrow::TimestampParser": + const char* kind() const + const char* format() const + + @staticmethod + shared_ptr[CTimestampParser] MakeStrptime(c_string format) + + @staticmethod + shared_ptr[CTimestampParser] MakeISO8601() + + +cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil: + + cdef cppclass CCSVInvalidRow" arrow::csv::InvalidRow": + int32_t expected_columns + int32_t actual_columns + int64_t number + c_string text + + ctypedef enum CInvalidRowResult" arrow::csv::InvalidRowResult": + CInvalidRowResult_Error" arrow::csv::InvalidRowResult::Error" + CInvalidRowResult_Skip" arrow::csv::InvalidRowResult::Skip" + + ctypedef CInvalidRowResult CInvalidRowHandler(const CCSVInvalidRow&) + + +cdef extern from "arrow/csv/api.h" namespace "arrow::csv" nogil: + + ctypedef enum CQuotingStyle "arrow::csv::QuotingStyle": + CQuotingStyle_Needed "arrow::csv::QuotingStyle::Needed" + CQuotingStyle_AllValid "arrow::csv::QuotingStyle::AllValid" + CQuotingStyle_None "arrow::csv::QuotingStyle::None" + + cdef cppclass CCSVParseOptions" arrow::csv::ParseOptions": + unsigned char delimiter + c_bool quoting + unsigned char quote_char + c_bool double_quote + c_bool escaping + unsigned char escape_char + c_bool newlines_in_values + c_bool ignore_empty_lines + function[CInvalidRowHandler] invalid_row_handler + + CCSVParseOptions() + CCSVParseOptions(CCSVParseOptions) + + @staticmethod + CCSVParseOptions Defaults() + + CStatus Validate() + + cdef cppclass CCSVConvertOptions" arrow::csv::ConvertOptions": + c_bool check_utf8 + unordered_map[c_string, shared_ptr[CDataType]] column_types + vector[c_string] null_values + vector[c_string] true_values + vector[c_string] false_values + c_bool strings_can_be_null + c_bool quoted_strings_can_be_null + vector[shared_ptr[CTimestampParser]] timestamp_parsers + + c_bool auto_dict_encode + int32_t auto_dict_max_cardinality + unsigned char decimal_point + + vector[c_string] include_columns + c_bool include_missing_columns + + CCSVConvertOptions() + CCSVConvertOptions(CCSVConvertOptions) + + @staticmethod + CCSVConvertOptions Defaults() + + CStatus Validate() + + cdef cppclass CCSVReadOptions" arrow::csv::ReadOptions": + c_bool use_threads + int32_t block_size + int32_t skip_rows + int32_t skip_rows_after_names + vector[c_string] column_names + c_bool autogenerate_column_names + + CCSVReadOptions() + CCSVReadOptions(CCSVReadOptions) + + @staticmethod + CCSVReadOptions Defaults() + + CStatus Validate() + + cdef cppclass CCSVWriteOptions" arrow::csv::WriteOptions": + c_bool include_header + int32_t batch_size + unsigned char delimiter + CQuotingStyle quoting_style + CQuotingStyle quoting_header + CIOContext io_context + + CCSVWriteOptions() + CCSVWriteOptions(CCSVWriteOptions) + + @staticmethod + CCSVWriteOptions Defaults() + + CStatus Validate() + + cdef cppclass CCSVReader" arrow::csv::TableReader": + @staticmethod + CResult[shared_ptr[CCSVReader]] Make( + CIOContext, shared_ptr[CInputStream], + CCSVReadOptions, CCSVParseOptions, CCSVConvertOptions) + + CResult[shared_ptr[CTable]] Read() + + cdef cppclass CCSVStreamingReader" arrow::csv::StreamingReader"( + CRecordBatchReader): + @staticmethod + CResult[shared_ptr[CCSVStreamingReader]] Make( + CIOContext, shared_ptr[CInputStream], + CCSVReadOptions, CCSVParseOptions, CCSVConvertOptions) + + cdef CStatus WriteCSV(CTable&, CCSVWriteOptions& options, COutputStream*) + cdef CStatus WriteCSV( + CRecordBatch&, CCSVWriteOptions& options, COutputStream*) + cdef CResult[shared_ptr[CRecordBatchWriter]] MakeCSVWriter( + shared_ptr[COutputStream], shared_ptr[CSchema], + CCSVWriteOptions& options) + + +cdef extern from "arrow/json/options.h" nogil: + + ctypedef enum CUnexpectedFieldBehavior \ + "arrow::json::UnexpectedFieldBehavior": + CUnexpectedFieldBehavior_Ignore \ + "arrow::json::UnexpectedFieldBehavior::Ignore" + CUnexpectedFieldBehavior_Error \ + "arrow::json::UnexpectedFieldBehavior::Error" + CUnexpectedFieldBehavior_InferType \ + "arrow::json::UnexpectedFieldBehavior::InferType" + + cdef cppclass CJSONReadOptions" arrow::json::ReadOptions": + c_bool use_threads + int32_t block_size + + @staticmethod + CJSONReadOptions Defaults() + + cdef cppclass CJSONParseOptions" arrow::json::ParseOptions": + shared_ptr[CSchema] explicit_schema + c_bool newlines_in_values + CUnexpectedFieldBehavior unexpected_field_behavior + + @staticmethod + CJSONParseOptions Defaults() + + +cdef extern from "arrow/json/reader.h" namespace "arrow::json" nogil: + + cdef cppclass CJSONReader" arrow::json::TableReader": + @staticmethod + CResult[shared_ptr[CJSONReader]] Make( + CMemoryPool*, shared_ptr[CInputStream], + CJSONReadOptions, CJSONParseOptions) + + CResult[shared_ptr[CTable]] Read() + + cdef cppclass CJSONStreamingReader" arrow::json::StreamingReader"( + CRecordBatchReader): + @staticmethod + CResult[shared_ptr[CJSONStreamingReader]] Make( + shared_ptr[CInputStream], + CJSONReadOptions, CJSONParseOptions, CIOContext) + + +cdef extern from "arrow/util/thread_pool.h" namespace "arrow::internal" nogil: + + cdef cppclass CExecutor "arrow::internal::Executor": + pass + + cdef cppclass CThreadPool "arrow::internal::ThreadPool"(CExecutor): + @staticmethod + CResult[shared_ptr[CThreadPool]] Make(int threads) + + CThreadPool* GetCpuThreadPool() + + +cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil: + + CStatus InitializeCompute " arrow::compute::Initialize"() + + cdef cppclass CExecBatch "arrow::compute::ExecBatch": + vector[CDatum] values + int64_t length + + cdef cppclass CExecContext" arrow::compute::ExecContext": + CExecContext() + CExecContext(CMemoryPool* pool) + CExecContext(CMemoryPool* pool, CExecutor* exc) + + CMemoryPool* memory_pool() const + CExecutor* executor() + + cdef cppclass CKernelSignature" arrow::compute::KernelSignature": + c_string ToString() const + + cdef cppclass CKernel" arrow::compute::Kernel": + shared_ptr[CKernelSignature] signature + + cdef cppclass CArrayKernel" arrow::compute::ArrayKernel"(CKernel): + pass + + cdef cppclass CScalarKernel" arrow::compute::ScalarKernel"(CArrayKernel): + pass + + cdef cppclass CVectorKernel" arrow::compute::VectorKernel"(CArrayKernel): + pass + + cdef cppclass CScalarAggregateKernel \ + " arrow::compute::ScalarAggregateKernel"(CKernel): + pass + + cdef cppclass CHashAggregateKernel \ + " arrow::compute::HashAggregateKernel"(CKernel): + pass + + cdef cppclass CArity" arrow::compute::Arity": + int num_args + c_bool is_varargs + + CArity() + + CArity(int num_args, c_bool is_varargs) + + cdef enum FunctionKind" arrow::compute::Function::Kind": + FunctionKind_SCALAR" arrow::compute::Function::SCALAR" + FunctionKind_VECTOR" arrow::compute::Function::VECTOR" + FunctionKind_SCALAR_AGGREGATE \ + " arrow::compute::Function::SCALAR_AGGREGATE" + FunctionKind_HASH_AGGREGATE \ + " arrow::compute::Function::HASH_AGGREGATE" + FunctionKind_META \ + " arrow::compute::Function::META" + + cdef cppclass CFunctionDoc" arrow::compute::FunctionDoc": + c_string summary + c_string description + vector[c_string] arg_names + c_string options_class + c_bool options_required + + cdef cppclass CFunctionOptionsType" arrow::compute::FunctionOptionsType": + const char* type_name() const + + cdef cppclass CFunctionOptions" arrow::compute::FunctionOptions": + const CFunctionOptionsType* options_type() const + const char* type_name() const + c_bool Equals(const CFunctionOptions& other) const + c_string ToString() const + unique_ptr[CFunctionOptions] Copy() const + CResult[shared_ptr[CBuffer]] Serialize() const + + @staticmethod + CResult[unique_ptr[CFunctionOptions]] Deserialize( + const c_string& type_name, const CBuffer& buffer) + + cdef cppclass CFunction" arrow::compute::Function": + const c_string& name() const + FunctionKind kind() const + const CArity& arity() const + const CFunctionDoc& doc() const + int num_kernels() const + CResult[CDatum] Execute(const vector[CDatum]& args, + const CFunctionOptions* options, + CExecContext* ctx) const + CResult[CDatum] Execute(const CExecBatch& args, + const CFunctionOptions* options, + CExecContext* ctx) const + + cdef cppclass CScalarFunction" arrow::compute::ScalarFunction"(CFunction): + vector[const CScalarKernel*] kernels() const + + cdef cppclass CVectorFunction" arrow::compute::VectorFunction"(CFunction): + vector[const CVectorKernel*] kernels() const + + cdef cppclass CScalarAggregateFunction \ + " arrow::compute::ScalarAggregateFunction"(CFunction): + vector[const CScalarAggregateKernel*] kernels() const + + cdef cppclass CHashAggregateFunction \ + " arrow::compute::HashAggregateFunction"(CFunction): + vector[const CHashAggregateKernel*] kernels() const + + cdef cppclass CMetaFunction" arrow::compute::MetaFunction"(CFunction): + pass + + cdef cppclass CFunctionRegistry" arrow::compute::FunctionRegistry": + CResult[shared_ptr[CFunction]] GetFunction( + const c_string& name) const + vector[c_string] GetFunctionNames() const + int num_functions() const + + CFunctionRegistry* GetFunctionRegistry() + + cdef cppclass CElementWiseAggregateOptions \ + "arrow::compute::ElementWiseAggregateOptions"(CFunctionOptions): + CElementWiseAggregateOptions(c_bool skip_nulls) + c_bool skip_nulls + + ctypedef enum CRoundMode \ + "arrow::compute::RoundMode": + CRoundMode_DOWN \ + "arrow::compute::RoundMode::DOWN" + CRoundMode_UP \ + "arrow::compute::RoundMode::UP" + CRoundMode_TOWARDS_ZERO \ + "arrow::compute::RoundMode::TOWARDS_ZERO" + CRoundMode_TOWARDS_INFINITY \ + "arrow::compute::RoundMode::TOWARDS_INFINITY" + CRoundMode_HALF_DOWN \ + "arrow::compute::RoundMode::HALF_DOWN" + CRoundMode_HALF_UP \ + "arrow::compute::RoundMode::HALF_UP" + CRoundMode_HALF_TOWARDS_ZERO \ + "arrow::compute::RoundMode::HALF_TOWARDS_ZERO" + CRoundMode_HALF_TOWARDS_INFINITY \ + "arrow::compute::RoundMode::HALF_TOWARDS_INFINITY" + CRoundMode_HALF_TO_EVEN \ + "arrow::compute::RoundMode::HALF_TO_EVEN" + CRoundMode_HALF_TO_ODD \ + "arrow::compute::RoundMode::HALF_TO_ODD" + + cdef cppclass CRoundOptions \ + "arrow::compute::RoundOptions"(CFunctionOptions): + CRoundOptions(int64_t ndigits, CRoundMode round_mode) + int64_t ndigits + CRoundMode round_mode + + cdef cppclass CRoundBinaryOptions \ + "arrow::compute::RoundBinaryOptions"(CFunctionOptions): + CRoundBinaryOptions(CRoundMode round_mode) + CRoundMode round_mode + + ctypedef enum CCalendarUnit \ + "arrow::compute::CalendarUnit": + CCalendarUnit_NANOSECOND \ + "arrow::compute::CalendarUnit::NANOSECOND" + CCalendarUnit_MICROSECOND \ + "arrow::compute::CalendarUnit::MICROSECOND" + CCalendarUnit_MILLISECOND \ + "arrow::compute::CalendarUnit::MILLISECOND" + CCalendarUnit_SECOND \ + "arrow::compute::CalendarUnit::SECOND" + CCalendarUnit_MINUTE \ + "arrow::compute::CalendarUnit::MINUTE" + CCalendarUnit_HOUR \ + "arrow::compute::CalendarUnit::HOUR" + CCalendarUnit_DAY \ + "arrow::compute::CalendarUnit::DAY" + CCalendarUnit_WEEK \ + "arrow::compute::CalendarUnit::WEEK" + CCalendarUnit_MONTH \ + "arrow::compute::CalendarUnit::MONTH" + CCalendarUnit_QUARTER \ + "arrow::compute::CalendarUnit::QUARTER" + CCalendarUnit_YEAR \ + "arrow::compute::CalendarUnit::YEAR" + + cdef cppclass CRoundTemporalOptions \ + "arrow::compute::RoundTemporalOptions"(CFunctionOptions): + CRoundTemporalOptions(int multiple, CCalendarUnit unit, + c_bool week_starts_monday, + c_bool ceil_is_strictly_greater, + c_bool calendar_based_origin) + int multiple + CCalendarUnit unit + c_bool week_starts_monday + c_bool ceil_is_strictly_greater + c_bool calendar_based_origin + + cdef cppclass CRoundToMultipleOptions \ + "arrow::compute::RoundToMultipleOptions"(CFunctionOptions): + CRoundToMultipleOptions(shared_ptr[CScalar] multiple, CRoundMode round_mode) + shared_ptr[CScalar] multiple + CRoundMode round_mode + + cdef enum CJoinNullHandlingBehavior \ + "arrow::compute::JoinOptions::NullHandlingBehavior": + CJoinNullHandlingBehavior_EMIT_NULL \ + "arrow::compute::JoinOptions::EMIT_NULL" + CJoinNullHandlingBehavior_SKIP \ + "arrow::compute::JoinOptions::SKIP" + CJoinNullHandlingBehavior_REPLACE \ + "arrow::compute::JoinOptions::REPLACE" + + cdef cppclass CJoinOptions \ + "arrow::compute::JoinOptions"(CFunctionOptions): + CJoinOptions(CJoinNullHandlingBehavior null_handling, + c_string null_replacement) + CJoinNullHandlingBehavior null_handling + c_string null_replacement + + cdef cppclass CMatchSubstringOptions \ + "arrow::compute::MatchSubstringOptions"(CFunctionOptions): + CMatchSubstringOptions(c_string pattern, c_bool ignore_case) + c_string pattern + c_bool ignore_case + + cdef cppclass CTrimOptions \ + "arrow::compute::TrimOptions"(CFunctionOptions): + CTrimOptions(c_string characters) + c_string characters + + cdef cppclass CPadOptions \ + "arrow::compute::PadOptions"(CFunctionOptions): + CPadOptions(int64_t width, c_string padding, c_bool lean_left_on_odd_padding) + int64_t width + c_string padding + c_bool lean_left_on_odd_padding + + cdef cppclass CZeroFillOptions \ + "arrow::compute::ZeroFillOptions"(CFunctionOptions): + CZeroFillOptions(int64_t width, c_string padding) + int64_t width + c_string padding + + cdef cppclass CSliceOptions \ + "arrow::compute::SliceOptions"(CFunctionOptions): + CSliceOptions(int64_t start, int64_t stop, int64_t step) + int64_t start + int64_t stop + int64_t step + + cdef cppclass CListSliceOptions \ + "arrow::compute::ListSliceOptions"(CFunctionOptions): + CListSliceOptions(int64_t start, optional[int64_t] stop, + int64_t step, + optional[c_bool] return_fixed_size_list) + int64_t start + optional[int64_t] stop + int64_t step + optional[c_bool] return_fixed_size_list + + cdef cppclass CSplitOptions \ + "arrow::compute::SplitOptions"(CFunctionOptions): + CSplitOptions(int64_t max_splits, c_bool reverse) + int64_t max_splits + c_bool reverse + + cdef cppclass CSplitPatternOptions \ + "arrow::compute::SplitPatternOptions"(CFunctionOptions): + CSplitPatternOptions(c_string pattern, int64_t max_splits, + c_bool reverse) + int64_t max_splits + c_bool reverse + c_string pattern + + cdef cppclass CReplaceSliceOptions \ + "arrow::compute::ReplaceSliceOptions"(CFunctionOptions): + CReplaceSliceOptions(int64_t start, int64_t stop, c_string replacement) + int64_t start + int64_t stop + c_string replacement + + cdef cppclass CReplaceSubstringOptions \ + "arrow::compute::ReplaceSubstringOptions"(CFunctionOptions): + CReplaceSubstringOptions(c_string pattern, c_string replacement, + int64_t max_replacements) + c_string pattern + c_string replacement + int64_t max_replacements + + cdef cppclass CExtractRegexOptions \ + "arrow::compute::ExtractRegexOptions"(CFunctionOptions): + CExtractRegexOptions(c_string pattern) + c_string pattern + + cdef cppclass CExtractRegexSpanOptions \ + "arrow::compute::ExtractRegexSpanOptions"(CFunctionOptions): + CExtractRegexSpanOptions(c_string pattern) + c_string pattern + + cdef cppclass CCastOptions" arrow::compute::CastOptions"(CFunctionOptions): + CCastOptions() + CCastOptions(c_bool safe) + CCastOptions(CCastOptions options) + + @staticmethod + CCastOptions Safe() + + @staticmethod + CCastOptions Unsafe() + shared_ptr[CDataType] to_type + c_bool allow_int_overflow + c_bool allow_time_truncate + c_bool allow_time_overflow + c_bool allow_decimal_truncate + c_bool allow_float_truncate + c_bool allow_invalid_utf8 + + cdef enum CFilterNullSelectionBehavior \ + "arrow::compute::FilterOptions::NullSelectionBehavior": + CFilterNullSelectionBehavior_DROP \ + "arrow::compute::FilterOptions::DROP" + CFilterNullSelectionBehavior_EMIT_NULL \ + "arrow::compute::FilterOptions::EMIT_NULL" + + cdef cppclass CFilterOptions \ + " arrow::compute::FilterOptions"(CFunctionOptions): + CFilterOptions() + CFilterOptions(CFilterNullSelectionBehavior null_selection_behavior) + CFilterNullSelectionBehavior null_selection_behavior + + cdef enum CDictionaryEncodeNullEncodingBehavior \ + "arrow::compute::DictionaryEncodeOptions::NullEncodingBehavior": + CDictionaryEncodeNullEncodingBehavior_ENCODE \ + "arrow::compute::DictionaryEncodeOptions::ENCODE" + CDictionaryEncodeNullEncodingBehavior_MASK \ + "arrow::compute::DictionaryEncodeOptions::MASK" + + cdef cppclass CDictionaryEncodeOptions \ + "arrow::compute::DictionaryEncodeOptions"(CFunctionOptions): + CDictionaryEncodeOptions( + CDictionaryEncodeNullEncodingBehavior null_encoding) + CDictionaryEncodeNullEncodingBehavior null_encoding + + cdef cppclass CRunEndEncodeOptions \ + "arrow::compute::RunEndEncodeOptions"(CFunctionOptions): + CRunEndEncodeOptions() + CRunEndEncodeOptions(shared_ptr[CDataType] run_end_type) + shared_ptr[CDataType] run_end_type + + cdef cppclass CTakeOptions \ + " arrow::compute::TakeOptions"(CFunctionOptions): + CTakeOptions(c_bool boundscheck) + c_bool boundscheck + + cdef cppclass CInversePermutationOptions \ + "arrow::compute::InversePermutationOptions"(CFunctionOptions): + CInversePermutationOptions(int64_t max_index, optional[shared_ptr[CDataType]] output_type) + int64_t max_index + optional[shared_ptr[CDataType]] output_type + + cdef cppclass CScatterOptions \ + "arrow::compute::ScatterOptions"(CFunctionOptions): + CScatterOptions(int64_t max_index) + int64_t max_index + + cdef cppclass CStrptimeOptions \ + "arrow::compute::StrptimeOptions"(CFunctionOptions): + CStrptimeOptions(c_string format, TimeUnit unit, c_bool raise_error) + c_string format + TimeUnit unit + c_bool raise_error + + cdef cppclass CStrftimeOptions \ + "arrow::compute::StrftimeOptions"(CFunctionOptions): + CStrftimeOptions(c_string format, c_string locale) + c_string format + c_string locale + + cdef cppclass CDayOfWeekOptions \ + "arrow::compute::DayOfWeekOptions"(CFunctionOptions): + CDayOfWeekOptions(c_bool count_from_zero, uint32_t week_start) + c_bool count_from_zero + uint32_t week_start + + cdef enum CAssumeTimezoneAmbiguous \ + "arrow::compute::AssumeTimezoneOptions::Ambiguous": + CAssumeTimezoneAmbiguous_AMBIGUOUS_RAISE \ + "arrow::compute::AssumeTimezoneOptions::AMBIGUOUS_RAISE" + CAssumeTimezoneAmbiguous_AMBIGUOUS_EARLIEST \ + "arrow::compute::AssumeTimezoneOptions::AMBIGUOUS_EARLIEST" + CAssumeTimezoneAmbiguous_AMBIGUOUS_LATEST \ + "arrow::compute::AssumeTimezoneOptions::AMBIGUOUS_LATEST" + + cdef enum CAssumeTimezoneNonexistent \ + "arrow::compute::AssumeTimezoneOptions::Nonexistent": + CAssumeTimezoneNonexistent_NONEXISTENT_RAISE \ + "arrow::compute::AssumeTimezoneOptions::NONEXISTENT_RAISE" + CAssumeTimezoneNonexistent_NONEXISTENT_EARLIEST \ + "arrow::compute::AssumeTimezoneOptions::NONEXISTENT_EARLIEST" + CAssumeTimezoneNonexistent_NONEXISTENT_LATEST \ + "arrow::compute::AssumeTimezoneOptions::NONEXISTENT_LATEST" + + cdef cppclass CAssumeTimezoneOptions \ + "arrow::compute::AssumeTimezoneOptions"(CFunctionOptions): + CAssumeTimezoneOptions(c_string timezone, + CAssumeTimezoneAmbiguous ambiguous, + CAssumeTimezoneNonexistent nonexistent) + c_string timezone + CAssumeTimezoneAmbiguous ambiguous + CAssumeTimezoneNonexistent nonexistent + + cdef cppclass CWeekOptions \ + "arrow::compute::WeekOptions"(CFunctionOptions): + CWeekOptions(c_bool week_starts_monday, c_bool count_from_zero, + c_bool first_week_is_fully_in_year) + c_bool week_starts_monday + c_bool count_from_zero + c_bool first_week_is_fully_in_year + + cdef cppclass CNullOptions \ + "arrow::compute::NullOptions"(CFunctionOptions): + CNullOptions(c_bool nan_is_null) + c_bool nan_is_null + + cdef cppclass CVarianceOptions \ + "arrow::compute::VarianceOptions"(CFunctionOptions): + CVarianceOptions(int ddof, c_bool skip_nulls, uint32_t min_count) + int ddof + c_bool skip_nulls + uint32_t min_count + + cdef cppclass CSkewOptions \ + "arrow::compute::SkewOptions"(CFunctionOptions): + CSkewOptions(c_bool skip_nulls, c_bool biased, uint32_t min_count) + c_bool skip_nulls + c_bool biased + uint32_t min_count + + cdef cppclass CScalarAggregateOptions \ + "arrow::compute::ScalarAggregateOptions"(CFunctionOptions): + CScalarAggregateOptions(c_bool skip_nulls, uint32_t min_count) + c_bool skip_nulls + uint32_t min_count + + cdef enum CCountMode "arrow::compute::CountOptions::CountMode": + CCountMode_ONLY_VALID "arrow::compute::CountOptions::ONLY_VALID" + CCountMode_ONLY_NULL "arrow::compute::CountOptions::ONLY_NULL" + CCountMode_ALL "arrow::compute::CountOptions::ALL" + + cdef cppclass CCountOptions \ + "arrow::compute::CountOptions"(CFunctionOptions): + CCountOptions(CCountMode mode) + CCountMode mode + + cdef cppclass CModeOptions \ + "arrow::compute::ModeOptions"(CFunctionOptions): + CModeOptions(int64_t n, c_bool skip_nulls, uint32_t min_count) + int64_t n + c_bool skip_nulls + uint32_t min_count + + cdef cppclass CIndexOptions \ + "arrow::compute::IndexOptions"(CFunctionOptions): + CIndexOptions(shared_ptr[CScalar] value) + shared_ptr[CScalar] value + + cdef cppclass CAggregate "arrow::compute::Aggregate": + c_string function + shared_ptr[CFunctionOptions] options + vector[CFieldRef] target + c_string name + + cdef enum CMapLookupOccurrence \ + "arrow::compute::MapLookupOptions::Occurrence": + CMapLookupOccurrence_ALL "arrow::compute::MapLookupOptions::ALL" + CMapLookupOccurrence_FIRST "arrow::compute::MapLookupOptions::FIRST" + CMapLookupOccurrence_LAST "arrow::compute::MapLookupOptions::LAST" + + cdef cppclass CMapLookupOptions \ + "arrow::compute::MapLookupOptions"(CFunctionOptions): + CMapLookupOptions(shared_ptr[CScalar] query_key, + CMapLookupOccurrence occurrence) + CMapLookupOccurrence occurrence + shared_ptr[CScalar] query_key + + cdef cppclass CMakeStructOptions \ + "arrow::compute::MakeStructOptions"(CFunctionOptions): + CMakeStructOptions(vector[c_string] n, + vector[c_bool] r, + vector[shared_ptr[const CKeyValueMetadata]] m) + CMakeStructOptions(vector[c_string] n) + vector[c_string] field_names + vector[c_bool] field_nullability + vector[shared_ptr[const CKeyValueMetadata]] field_metadata + + cdef cppclass CStructFieldOptions \ + "arrow::compute::StructFieldOptions"(CFunctionOptions): + CStructFieldOptions(vector[int] indices) + CStructFieldOptions(CFieldRef field_ref) + vector[int] indices + CFieldRef field_ref + + ctypedef enum CSortOrder" arrow::compute::SortOrder": + CSortOrder_Ascending \ + "arrow::compute::SortOrder::Ascending" + CSortOrder_Descending \ + "arrow::compute::SortOrder::Descending" + + ctypedef enum CNullPlacement" arrow::compute::NullPlacement": + CNullPlacement_AtStart \ + "arrow::compute::NullPlacement::AtStart" + CNullPlacement_AtEnd \ + "arrow::compute::NullPlacement::AtEnd" + + cdef cppclass CPartitionNthOptions \ + "arrow::compute::PartitionNthOptions"(CFunctionOptions): + CPartitionNthOptions(int64_t pivot, CNullPlacement) + int64_t pivot + CNullPlacement null_placement + + cdef cppclass CWinsorizeOptions \ + "arrow::compute::WinsorizeOptions"(CFunctionOptions): + CWinsorizeOptions(double lower_limit, double upper_limit) + + cdef cppclass CCumulativeOptions \ + "arrow::compute::CumulativeOptions"(CFunctionOptions): + CCumulativeOptions(c_bool skip_nulls) + CCumulativeOptions(shared_ptr[CScalar] start, c_bool skip_nulls) + optional[shared_ptr[CScalar]] start + c_bool skip_nulls + + cdef cppclass CPairwiseOptions \ + "arrow::compute::PairwiseOptions"(CFunctionOptions): + CPairwiseOptions(int64_t period) + int64_t period + + cdef cppclass CListFlattenOptions\ + "arrow::compute::ListFlattenOptions"(CFunctionOptions): + CListFlattenOptions(c_bool recursive) + c_bool recursive + + cdef cppclass CArraySortOptions \ + "arrow::compute::ArraySortOptions"(CFunctionOptions): + CArraySortOptions(CSortOrder, CNullPlacement) + CSortOrder order + CNullPlacement null_placement + + cdef cppclass CSortKey" arrow::compute::SortKey": + CSortKey(CFieldRef target, CSortOrder order) + CFieldRef target + CSortOrder order + + cdef cppclass COrdering" arrow::compute::Ordering": + COrdering(vector[CSortKey] sort_keys, CNullPlacement null_placement) + + cdef cppclass CSortOptions \ + "arrow::compute::SortOptions"(CFunctionOptions): + CSortOptions(vector[CSortKey] sort_keys, CNullPlacement) + vector[CSortKey] sort_keys + CNullPlacement null_placement + + cdef cppclass CSelectKOptions \ + "arrow::compute::SelectKOptions"(CFunctionOptions): + CSelectKOptions(int64_t k, vector[CSortKey] sort_keys) + int64_t k + vector[CSortKey] sort_keys + + cdef enum CQuantileInterp \ + "arrow::compute::QuantileOptions::Interpolation": + CQuantileInterp_LINEAR "arrow::compute::QuantileOptions::LINEAR" + CQuantileInterp_LOWER "arrow::compute::QuantileOptions::LOWER" + CQuantileInterp_HIGHER "arrow::compute::QuantileOptions::HIGHER" + CQuantileInterp_NEAREST "arrow::compute::QuantileOptions::NEAREST" + CQuantileInterp_MIDPOINT "arrow::compute::QuantileOptions::MIDPOINT" + + cdef cppclass CQuantileOptions \ + "arrow::compute::QuantileOptions"(CFunctionOptions): + CQuantileOptions(vector[double] q, CQuantileInterp interpolation, + c_bool skip_nulls, uint32_t min_count) + vector[double] q + CQuantileInterp interpolation + c_bool skip_nulls + uint32_t min_count + + cdef cppclass CTDigestOptions \ + "arrow::compute::TDigestOptions"(CFunctionOptions): + CTDigestOptions(vector[double] q, + uint32_t delta, uint32_t buffer_size, + c_bool skip_nulls, uint32_t min_count) + vector[double] q + uint32_t delta + uint32_t buffer_size + c_bool skip_nulls + uint32_t min_count + + cdef enum CUtf8NormalizeForm \ + "arrow::compute::Utf8NormalizeOptions::Form": + CUtf8NormalizeForm_NFC "arrow::compute::Utf8NormalizeOptions::NFC" + CUtf8NormalizeForm_NFKC "arrow::compute::Utf8NormalizeOptions::NFKC" + CUtf8NormalizeForm_NFD "arrow::compute::Utf8NormalizeOptions::NFD" + CUtf8NormalizeForm_NFKD "arrow::compute::Utf8NormalizeOptions::NFKD" + + cdef cppclass CUtf8NormalizeOptions \ + "arrow::compute::Utf8NormalizeOptions"(CFunctionOptions): + CUtf8NormalizeOptions(CUtf8NormalizeForm form) + CUtf8NormalizeForm form + + cdef cppclass CSetLookupOptions \ + "arrow::compute::SetLookupOptions"(CFunctionOptions): + CSetLookupOptions(CDatum value_set, c_bool skip_nulls) + CDatum value_set + c_bool skip_nulls + + cdef cppclass CRandomOptions \ + "arrow::compute::RandomOptions"(CFunctionOptions): + CRandomOptions(CRandomOptions) + + @staticmethod + CRandomOptions FromSystemRandom() + + @staticmethod + CRandomOptions FromSeed(uint64_t seed) + + cdef enum CRankOptionsTiebreaker \ + "arrow::compute::RankOptions::Tiebreaker": + CRankOptionsTiebreaker_Min "arrow::compute::RankOptions::Min" + CRankOptionsTiebreaker_Max "arrow::compute::RankOptions::Max" + CRankOptionsTiebreaker_First "arrow::compute::RankOptions::First" + CRankOptionsTiebreaker_Dense "arrow::compute::RankOptions::Dense" + + cdef cppclass CRankOptions \ + "arrow::compute::RankOptions"(CFunctionOptions): + CRankOptions(vector[CSortKey] sort_keys, CNullPlacement, + CRankOptionsTiebreaker tiebreaker) + vector[CSortKey] sort_keys + CNullPlacement null_placement + CRankOptionsTiebreaker tiebreaker + + cdef cppclass CRankQuantileOptions \ + "arrow::compute::RankQuantileOptions"(CFunctionOptions): + CRankQuantileOptions(vector[CSortKey] sort_keys, CNullPlacement) + vector[CSortKey] sort_keys + CNullPlacement null_placement + + cdef enum PivotWiderUnexpectedKeyBehavior \ + "arrow::compute::PivotWiderOptions::UnexpectedKeyBehavior": + PivotWiderUnexpectedKeyBehavior_Ignore "arrow::compute::PivotWiderOptions::kIgnore" + PivotWiderUnexpectedKeyBehavior_Raise "arrow::compute::PivotWiderOptions::kRaise" + + cdef cppclass CPivotWiderOptions \ + "arrow::compute::PivotWiderOptions"(CFunctionOptions): + CPivotWiderOptions(vector[c_string] key_names, + PivotWiderUnexpectedKeyBehavior) + + cdef enum DatumType" arrow::Datum::type": + DatumType_NONE" arrow::Datum::NONE" + DatumType_SCALAR" arrow::Datum::SCALAR" + DatumType_ARRAY" arrow::Datum::ARRAY" + DatumType_CHUNKED_ARRAY" arrow::Datum::CHUNKED_ARRAY" + DatumType_RECORD_BATCH" arrow::Datum::RECORD_BATCH" + DatumType_TABLE" arrow::Datum::TABLE" + DatumType_COLLECTION" arrow::Datum::COLLECTION" + + cdef cppclass CDatum" arrow::Datum": + CDatum() + CDatum(const shared_ptr[CArray]& value) + CDatum(const shared_ptr[CChunkedArray]& value) + CDatum(const shared_ptr[CScalar]& value) + CDatum(const shared_ptr[CRecordBatch]& value) + CDatum(const shared_ptr[CTable]& value) + + DatumType kind() const + c_string ToString() const + + const shared_ptr[CArrayData]& array() const + const shared_ptr[CChunkedArray]& chunked_array() const + const shared_ptr[CRecordBatch]& record_batch() const + const shared_ptr[CTable]& table() const + const shared_ptr[CScalar]& scalar() const + + cdef c_string ToString(DatumType kind) + + +cdef extern from * namespace "arrow::compute": + # inlined from compute/function_internal.h to avoid exposing + # implementation details + """ + #include "arrow/compute/function.h" + namespace arrow { + namespace compute { + namespace internal { + Result> DeserializeFunctionOptions( + const Buffer& buffer); + } // namespace internal + } // namespace compute + } // namespace arrow + """ + CResult[unique_ptr[CFunctionOptions]] DeserializeFunctionOptions \ + " arrow::compute::internal::DeserializeFunctionOptions"( + const CBuffer& buffer) + + +cdef extern from * namespace "arrow::compute": + # inlined from expression_internal.h to avoid + # proliferation of #include + """ + #include + + #include "arrow/type.h" + #include "arrow/datum.h" + + namespace arrow { + namespace compute { + struct KnownFieldValues { + std::unordered_map map; + }; + } // namespace compute + } // namespace arrow + """ + cdef struct CKnownFieldValues "arrow::compute::KnownFieldValues": + unordered_map[CFieldRef, CDatum, CFieldRefHash] map + +cdef extern from "arrow/compute/expression.h" \ + namespace "arrow::compute" nogil: + + cdef cppclass CExpression "arrow::compute::Expression": + c_bool Equals(const CExpression& other) const + c_string ToString() const + CResult[CExpression] Bind(const CSchema&) + const CFieldRef* field_ref() const + + cdef CExpression CMakeScalarExpression \ + "arrow::compute::literal"(shared_ptr[CScalar] value) + + cdef CExpression CMakeFieldExpression \ + "arrow::compute::field_ref"(CFieldRef) + + cdef CExpression CMakeFieldExpressionByIndex \ + "arrow::compute::field_ref"(int idx) + + cdef CExpression CMakeCallExpression \ + "arrow::compute::call"(c_string function, + vector[CExpression] arguments, + shared_ptr[CFunctionOptions] options) + + cdef CResult[shared_ptr[CBuffer]] CSerializeExpression \ + "arrow::compute::Serialize"(const CExpression&) + + cdef CResult[CExpression] CDeserializeExpression \ + "arrow::compute::Deserialize"(shared_ptr[CBuffer]) + + cdef CResult[CKnownFieldValues] \ + CExtractKnownFieldValues "arrow::compute::ExtractKnownFieldValues"( + const CExpression& partition_expression) + + +cdef extern from "arrow/extension_type.h" namespace "arrow": + cdef cppclass CExtensionTypeRegistry" arrow::ExtensionTypeRegistry": + @staticmethod + shared_ptr[CExtensionTypeRegistry] GetGlobalRegistry() + + cdef cppclass CExtensionType" arrow::ExtensionType"(CDataType): + c_string extension_name() + shared_ptr[CDataType] storage_type() + int byte_width() + int bit_width() + + @staticmethod + shared_ptr[CArray] WrapArray(shared_ptr[CDataType] ext_type, + shared_ptr[CArray] storage) + + @staticmethod + shared_ptr[CChunkedArray] WrapArray(shared_ptr[CDataType] ext_type, + shared_ptr[CChunkedArray] storage) + + cdef cppclass CExtensionArray" arrow::ExtensionArray"(CArray): + CExtensionArray(shared_ptr[CDataType], shared_ptr[CArray] storage) + + shared_ptr[CArray] storage() + + +cdef extern from "arrow/extension/json.h" namespace "arrow::extension" nogil: + cdef cppclass CJsonType" arrow::extension::JsonExtensionType"(CExtensionType): + + @staticmethod + CResult[shared_ptr[CDataType]] Make(shared_ptr[CDataType]& storage_type) + + +cdef extern from "arrow/extension/uuid.h" namespace "arrow::extension" nogil: + cdef cppclass CUuidType" arrow::extension::UuidType"(CExtensionType): + + @staticmethod + CResult[shared_ptr[CDataType]] Make() + + cdef cppclass CUuidArray" arrow::extension::UuidArray"(CExtensionArray): + pass + + +cdef extern from "arrow/extension/fixed_shape_tensor.h" namespace "arrow::extension" nogil: + cdef cppclass CFixedShapeTensorType \ + " arrow::extension::FixedShapeTensorType"(CExtensionType): + + CResult[shared_ptr[CTensor]] MakeTensor(const shared_ptr[CExtensionScalar]& scalar) const + + @staticmethod + CResult[shared_ptr[CDataType]] Make(const shared_ptr[CDataType]& value_type, + const vector[int64_t]& shape, + const vector[int64_t]& permutation, + const vector[c_string]& dim_names) + + const shared_ptr[CDataType] value_type() + const vector[int64_t] shape() + const vector[int64_t] permutation() + const vector[c_string] dim_names() + + cdef cppclass CFixedShapeTensorArray \ + " arrow::extension::FixedShapeTensorArray"(CExtensionArray): + const CResult[shared_ptr[CTensor]] ToTensor() const + + +cdef extern from "arrow/extension/opaque.h" namespace "arrow::extension" nogil: + cdef cppclass COpaqueType \ + " arrow::extension::OpaqueType"(CExtensionType): + + c_string type_name() + c_string vendor_name() + + cdef cppclass COpaqueArray \ + " arrow::extension::OpaqueArray"(CExtensionArray): + pass + + +cdef extern from "arrow/extension/bool8.h" namespace "arrow::extension" nogil: + cdef cppclass CBool8Type" arrow::extension::Bool8Type"(CExtensionType): + + @staticmethod + CResult[shared_ptr[CDataType]] Make() + + cdef cppclass CBool8Array" arrow::extension::Bool8Array"(CExtensionArray): + pass + +cdef extern from "arrow/util/compression.h" namespace "arrow" nogil: + cdef enum CCompressionType" arrow::Compression::type": + CCompressionType_UNCOMPRESSED" arrow::Compression::UNCOMPRESSED" + CCompressionType_SNAPPY" arrow::Compression::SNAPPY" + CCompressionType_GZIP" arrow::Compression::GZIP" + CCompressionType_BROTLI" arrow::Compression::BROTLI" + CCompressionType_ZSTD" arrow::Compression::ZSTD" + CCompressionType_LZ4" arrow::Compression::LZ4" + CCompressionType_LZ4_FRAME" arrow::Compression::LZ4_FRAME" + CCompressionType_BZ2" arrow::Compression::BZ2" + + cdef cppclass CCodec" arrow::util::Codec": + @staticmethod + CResult[unique_ptr[CCodec]] Create(CCompressionType codec) + + @staticmethod + CResult[unique_ptr[CCodec]] CreateWithLevel" Create"( + CCompressionType codec, + int compression_level) + + @staticmethod + c_bool SupportsCompressionLevel(CCompressionType codec) + + @staticmethod + CResult[int] MinimumCompressionLevel(CCompressionType codec) + + @staticmethod + CResult[int] MaximumCompressionLevel(CCompressionType codec) + + @staticmethod + CResult[int] DefaultCompressionLevel(CCompressionType codec) + + @staticmethod + c_bool IsAvailable(CCompressionType codec) + + CResult[int64_t] Decompress(int64_t input_len, const uint8_t* input, + int64_t output_len, + uint8_t* output_buffer) + CResult[int64_t] Compress(int64_t input_len, const uint8_t* input, + int64_t output_buffer_len, + uint8_t* output_buffer) + c_string name() const + int compression_level() const + int64_t MaxCompressedLen(int64_t input_len, const uint8_t* input) + + +cdef extern from "arrow/util/io_util.h" namespace "arrow::internal" nogil: + int ErrnoFromStatus(CStatus status) + int WinErrorFromStatus(CStatus status) + int SignalFromStatus(CStatus status) + + CStatus SendSignal(int signum) + CStatus SendSignalToThread(int signum, uint64_t thread_id) + + +cdef extern from "arrow/util/iterator.h" namespace "arrow" nogil: + cdef cppclass CIterator" arrow::Iterator"[T]: + CResult[T] Next() + CStatus Visit[Visitor](Visitor&& visitor) + cppclass RangeIterator: + CResult[T] operator*() + RangeIterator& operator++() + c_bool operator!=(RangeIterator) const + RangeIterator begin() + RangeIterator end() + CIterator[T] MakeVectorIterator[T](vector[T] v) + + +cdef extern from "arrow/util/secure_string.h" namespace "arrow" nogil: + cdef cppclass CSecureString" arrow::util::SecureString": + CSecureString() + CSecureString(c_string s) + CSecureString(const CSecureString& s) + CSecureString(size_t n, char c) + cpp_string_view as_view() + + +cdef extern from "arrow/util/thread_pool.h" namespace "arrow" nogil: + int GetCpuThreadPoolCapacity() + CStatus SetCpuThreadPoolCapacity(int threads) + +cdef extern from "arrow/array/concatenate.h" namespace "arrow" nogil: + CResult[shared_ptr[CArray]] Concatenate( + const vector[shared_ptr[CArray]]& arrays, + CMemoryPool* pool) + +cdef extern from "arrow/c/abi.h": + cdef struct ArrowSchema: + void (*release)(ArrowSchema*) noexcept nogil + + cdef struct ArrowArray: + void (*release)(ArrowArray*) noexcept nogil + + cdef struct ArrowArrayStream: + void (*release)(ArrowArrayStream*) noexcept nogil + + ctypedef int32_t ArrowDeviceType + cdef ArrowDeviceType ARROW_DEVICE_CUDA + + cdef struct ArrowDeviceArray: + ArrowArray array + int64_t device_id + int32_t device_type + +cdef extern from "arrow/c/bridge.h" namespace "arrow" nogil: + CStatus ExportType(CDataType&, ArrowSchema* out) + CResult[shared_ptr[CDataType]] ImportType(ArrowSchema*) + + CStatus ExportField(CField&, ArrowSchema* out) + CResult[shared_ptr[CField]] ImportField(ArrowSchema*) + + CStatus ExportSchema(CSchema&, ArrowSchema* out) + CResult[shared_ptr[CSchema]] ImportSchema(ArrowSchema*) + + CStatus ExportArray(CArray&, ArrowArray* out) + CStatus ExportArray(CArray&, ArrowArray* out, ArrowSchema* out_schema) + CResult[shared_ptr[CArray]] ImportArray(ArrowArray*, + shared_ptr[CDataType]) + CResult[shared_ptr[CArray]] ImportArray(ArrowArray*, ArrowSchema*) + + CStatus ExportRecordBatch(CRecordBatch&, ArrowArray* out) + CStatus ExportRecordBatch(CRecordBatch&, ArrowArray* out, + ArrowSchema* out_schema) + CResult[shared_ptr[CRecordBatch]] ImportRecordBatch(ArrowArray*, + shared_ptr[CSchema]) + CResult[shared_ptr[CRecordBatch]] ImportRecordBatch(ArrowArray*, + ArrowSchema*) + + CStatus ExportRecordBatchReader(shared_ptr[CRecordBatchReader], + ArrowArrayStream*) + CResult[shared_ptr[CRecordBatchReader]] ImportRecordBatchReader( + ArrowArrayStream*) + + CStatus ExportChunkedArray(shared_ptr[CChunkedArray], ArrowArrayStream*) + CResult[shared_ptr[CChunkedArray]] ImportChunkedArray(ArrowArrayStream*) + + CStatus ExportDeviceArray(const CArray&, shared_ptr[CSyncEvent], + ArrowDeviceArray* out, ArrowSchema*) + CResult[shared_ptr[CArray]] ImportDeviceArray( + ArrowDeviceArray*, shared_ptr[CDataType]) + CResult[shared_ptr[CArray]] ImportDeviceArray( + ArrowDeviceArray*, ArrowSchema*) + + CStatus ExportDeviceRecordBatch(const CRecordBatch&, shared_ptr[CSyncEvent], + ArrowDeviceArray* out, ArrowSchema*) + CResult[shared_ptr[CRecordBatch]] ImportDeviceRecordBatch( + ArrowDeviceArray*, shared_ptr[CSchema]) + CResult[shared_ptr[CRecordBatch]] ImportDeviceRecordBatch( + ArrowDeviceArray*, ArrowSchema*) + + +cdef extern from "arrow/util/byte_size.h" namespace "arrow::util" nogil: + CResult[int64_t] ReferencedBufferSize(const CArray& array_data) + CResult[int64_t] ReferencedBufferSize(const CRecordBatch& record_batch) + CResult[int64_t] ReferencedBufferSize(const CChunkedArray& chunked_array) + CResult[int64_t] ReferencedBufferSize(const CTable& table) + int64_t TotalBufferSize(const CArray& array) + int64_t TotalBufferSize(const CChunkedArray& array) + int64_t TotalBufferSize(const CRecordBatch& record_batch) + int64_t TotalBufferSize(const CTable& table) + +ctypedef PyObject* CallbackUdf(object user_function, const CUdfContext& context, object inputs) + + +cdef extern from "arrow/api.h" namespace "arrow" nogil: + + cdef cppclass CRecordBatchIterator "arrow::RecordBatchIterator"( + CIterator[shared_ptr[CRecordBatch]]): + pass + + +cdef extern from "arrow/python/udf.h" namespace "arrow::py" nogil: + cdef cppclass CUdfContext" arrow::py::UdfContext": + CMemoryPool *pool + int64_t batch_length + + cdef cppclass CUdfOptions" arrow::py::UdfOptions": + c_string func_name + CArity arity + CFunctionDoc func_doc + vector[shared_ptr[CDataType]] input_types + shared_ptr[CDataType] output_type + + CStatus RegisterScalarFunction(PyObject* function, + function[CallbackUdf] wrapper, const CUdfOptions& options, + CFunctionRegistry* registry) + + CStatus RegisterTabularFunction(PyObject* function, + function[CallbackUdf] wrapper, const CUdfOptions& options, + CFunctionRegistry* registry) + + CStatus RegisterAggregateFunction(PyObject* function, + function[CallbackUdf] wrapper, const CUdfOptions& options, + CFunctionRegistry* registry) + + CStatus RegisterVectorFunction(PyObject* function, + function[CallbackUdf] wrapper, const CUdfOptions& options, + CFunctionRegistry* registry) + + CResult[shared_ptr[CRecordBatchReader]] CallTabularFunction( + const c_string& func_name, const vector[CDatum]& args, CFunctionRegistry* registry) + +cdef extern from "arrow/compute/cast.h" namespace "arrow::compute": + CResult[CDatum] Cast(const CDatum& value, const CCastOptions& options) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_acero.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_acero.pxd new file mode 100644 index 0000000000000000000000000000000000000000..dc9babee190e110583d9d2cc24b71567203290af --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_acero.pxd @@ -0,0 +1,118 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * + + +cdef extern from "arrow/acero/options.h" namespace "arrow::acero" nogil: + cdef enum CJoinType "arrow::acero::JoinType": + CJoinType_LEFT_SEMI "arrow::acero::JoinType::LEFT_SEMI" + CJoinType_RIGHT_SEMI "arrow::acero::JoinType::RIGHT_SEMI" + CJoinType_LEFT_ANTI "arrow::acero::JoinType::LEFT_ANTI" + CJoinType_RIGHT_ANTI "arrow::acero::JoinType::RIGHT_ANTI" + CJoinType_INNER "arrow::acero::JoinType::INNER" + CJoinType_LEFT_OUTER "arrow::acero::JoinType::LEFT_OUTER" + CJoinType_RIGHT_OUTER "arrow::acero::JoinType::RIGHT_OUTER" + CJoinType_FULL_OUTER "arrow::acero::JoinType::FULL_OUTER" + + cdef cppclass CExecNodeOptions "arrow::acero::ExecNodeOptions": + pass + + cdef cppclass CSourceNodeOptions "arrow::acero::SourceNodeOptions"(CExecNodeOptions): + pass + + cdef cppclass CTableSourceNodeOptions "arrow::acero::TableSourceNodeOptions"(CExecNodeOptions): + CTableSourceNodeOptions(shared_ptr[CTable] table) + CTableSourceNodeOptions(shared_ptr[CTable] table, int64_t max_batch_size) + + cdef cppclass CSinkNodeOptions "arrow::acero::SinkNodeOptions"(CExecNodeOptions): + pass + + cdef cppclass CFilterNodeOptions "arrow::acero::FilterNodeOptions"(CExecNodeOptions): + CFilterNodeOptions(CExpression) + + cdef cppclass CProjectNodeOptions "arrow::acero::ProjectNodeOptions"(CExecNodeOptions): + CProjectNodeOptions(vector[CExpression] expressions) + CProjectNodeOptions(vector[CExpression] expressions, + vector[c_string] names) + + cdef cppclass CAggregateNodeOptions "arrow::acero::AggregateNodeOptions"(CExecNodeOptions): + CAggregateNodeOptions(vector[CAggregate] aggregates, vector[CFieldRef] names) + + cdef cppclass COrderByNodeOptions "arrow::acero::OrderByNodeOptions"(CExecNodeOptions): + COrderByNodeOptions(COrdering ordering) + + cdef cppclass CHashJoinNodeOptions "arrow::acero::HashJoinNodeOptions"(CExecNodeOptions): + CHashJoinNodeOptions(CJoinType, vector[CFieldRef] in_left_keys, + vector[CFieldRef] in_right_keys) + CHashJoinNodeOptions(CJoinType, vector[CFieldRef] in_left_keys, + vector[CFieldRef] in_right_keys, + CExpression filter, + c_string output_suffix_for_left, + c_string output_suffix_for_right) + CHashJoinNodeOptions(CJoinType join_type, + vector[CFieldRef] left_keys, + vector[CFieldRef] right_keys, + vector[CFieldRef] left_output, + vector[CFieldRef] right_output, + CExpression filter, + c_string output_suffix_for_left, + c_string output_suffix_for_right) + + cdef struct CAsofJoinKeys "arrow::acero::AsofJoinNodeOptions::Keys": + CFieldRef on_key + vector[CFieldRef] by_key + + cdef cppclass CAsofJoinNodeOptions "arrow::acero::AsofJoinNodeOptions"(CExecNodeOptions): + CAsofJoinNodeOptions(vector[CAsofJoinKeys] keys, int64_t tolerance) + + +cdef extern from "arrow/acero/exec_plan.h" namespace "arrow::acero" nogil: + cdef cppclass CDeclaration "arrow::acero::Declaration": + cppclass Input: + Input(CExecNode*) + Input(CDeclaration) + + c_string label + vector[Input] inputs + + CDeclaration() + CDeclaration(c_string factory_name, CExecNodeOptions options) + CDeclaration(c_string factory_name, vector[Input] inputs, shared_ptr[CExecNodeOptions] options) + + @staticmethod + CDeclaration Sequence(vector[CDeclaration] decls) + + cdef cppclass CExecNode "arrow::acero::ExecNode": + const vector[CExecNode*]& inputs() const + const shared_ptr[CSchema]& output_schema() const + + CResult[shared_ptr[CTable]] DeclarationToTable( + CDeclaration declaration, c_bool use_threads + ) + CResult[shared_ptr[CTable]] DeclarationToTable( + CDeclaration declaration, c_bool use_threads, + CMemoryPool* memory_pool, CFunctionRegistry* function_registry + ) + CResult[unique_ptr[CRecordBatchReader]] DeclarationToReader( + CDeclaration declaration, c_bool use_threads + ) + + CResult[c_string] DeclarationToString(const CDeclaration& declaration) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_cuda.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_cuda.pxd new file mode 100644 index 0000000000000000000000000000000000000000..39ca53908a68b14966583d3568703da2c0ab408c --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_cuda.pxd @@ -0,0 +1,109 @@ +# 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 pyarrow.includes.libarrow cimport * + +cdef extern from "arrow/gpu/cuda_api.h" namespace "arrow::cuda" nogil: + + cdef cppclass CCudaDeviceManager" arrow::cuda::CudaDeviceManager": + @staticmethod + CResult[CCudaDeviceManager*] Instance() + CResult[shared_ptr[CCudaContext]] GetContext(int gpu_number) + CResult[shared_ptr[CCudaContext]] GetSharedContext(int gpu_number, + void* handle) + CStatus AllocateHost(int device_number, int64_t nbytes, + shared_ptr[CCudaHostBuffer]* buffer) + int num_devices() const + + cdef cppclass CCudaContext" arrow::cuda::CudaContext": + CResult[shared_ptr[CCudaBuffer]] Allocate(int64_t nbytes) + CResult[shared_ptr[CCudaBuffer]] View(uint8_t* data, int64_t nbytes) + CResult[shared_ptr[CCudaBuffer]] OpenIpcBuffer( + const CCudaIpcMemHandle& ipc_handle) + CStatus Synchronize() + int64_t bytes_allocated() const + const void* handle() const + int device_number() const + CResult[uintptr_t] GetDeviceAddress(uintptr_t addr) + shared_ptr[CDevice] device() const + shared_ptr[CMemoryManager] memory_manager() const + + cdef cppclass CCudaIpcMemHandle" arrow::cuda::CudaIpcMemHandle": + @staticmethod + CResult[shared_ptr[CCudaIpcMemHandle]] FromBuffer( + const void* opaque_handle) + CResult[shared_ptr[CBuffer]] Serialize(CMemoryPool* pool) const + + cdef cppclass CCudaBuffer" arrow::cuda::CudaBuffer"(CBuffer): + CCudaBuffer(uint8_t* data, int64_t size, + const shared_ptr[CCudaContext]& context, + c_bool own_data=false, c_bool is_ipc=false) + CCudaBuffer(const shared_ptr[CCudaBuffer]& parent, + const int64_t offset, const int64_t size) + + @staticmethod + CResult[shared_ptr[CCudaBuffer]] FromBuffer(shared_ptr[CBuffer] buf) + + CStatus CopyToHost(const int64_t position, const int64_t nbytes, + void* out) const + CStatus CopyFromHost(const int64_t position, const void* data, + int64_t nbytes) + CStatus CopyFromDevice(const int64_t position, const void* data, + int64_t nbytes) + CStatus CopyFromAnotherDevice(const shared_ptr[CCudaContext]& src_ctx, + const int64_t position, const void* data, + int64_t nbytes) + CResult[shared_ptr[CCudaIpcMemHandle]] ExportForIpc() + shared_ptr[CCudaContext] context() const + + cdef cppclass \ + CCudaHostBuffer" arrow::cuda::CudaHostBuffer"(CMutableBuffer): + pass + + cdef cppclass \ + CCudaBufferReader" arrow::cuda::CudaBufferReader"(CBufferReader): + CCudaBufferReader(const shared_ptr[CBuffer]& buffer) + CResult[int64_t] Read(int64_t nbytes, void* buffer) + CResult[shared_ptr[CBuffer]] Read(int64_t nbytes) + + cdef cppclass \ + CCudaBufferWriter" arrow::cuda::CudaBufferWriter"(WritableFile): + CCudaBufferWriter(const shared_ptr[CCudaBuffer]& buffer) + CStatus Close() + CStatus Write(const void* data, int64_t nbytes) + CStatus WriteAt(int64_t position, const void* data, int64_t nbytes) + CStatus SetBufferSize(const int64_t buffer_size) + int64_t buffer_size() + int64_t num_bytes_buffered() const + + CResult[shared_ptr[CCudaHostBuffer]] AllocateCudaHostBuffer( + int device_number, const int64_t size) + + # Cuda prefix is added to avoid picking up arrow::cuda functions + # from arrow namespace. + CResult[shared_ptr[CCudaBuffer]] \ + CudaSerializeRecordBatch" arrow::cuda::SerializeRecordBatch"\ + (const CRecordBatch& batch, + CCudaContext* ctx) + CResult[shared_ptr[CRecordBatch]] \ + CudaReadRecordBatch" arrow::cuda::ReadRecordBatch"\ + (const shared_ptr[CSchema]& schema, + CDictionaryMemo* dictionary_memo, + const shared_ptr[CCudaBuffer]& buffer, + CMemoryPool* pool) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset.pxd new file mode 100644 index 0000000000000000000000000000000000000000..12cf82733935870b592e882d27412cbbdf293784 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset.pxd @@ -0,0 +1,423 @@ +# 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.unordered_map cimport unordered_map +from libcpp cimport bool as c_bool + +from pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * +from pyarrow.includes.libarrow_acero cimport * +from pyarrow.includes.libarrow_fs cimport * + + +cdef extern from "arrow/dataset/plan.h" namespace "arrow::dataset::internal" nogil: + + cdef void Initialize() + + +ctypedef CStatus cb_writer_finish_internal(CFileWriter*) +ctypedef void cb_writer_finish(dict, CFileWriter*) + +cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil: + + cdef enum ExistingDataBehavior" arrow::dataset::ExistingDataBehavior": + ExistingDataBehavior_DELETE_MATCHING" \ + arrow::dataset::ExistingDataBehavior::kDeleteMatchingPartitions" + ExistingDataBehavior_OVERWRITE_OR_IGNORE" \ + arrow::dataset::ExistingDataBehavior::kOverwriteOrIgnore" + ExistingDataBehavior_ERROR" \ + arrow::dataset::ExistingDataBehavior::kError" + + cdef cppclass CScanOptions "arrow::dataset::ScanOptions": + shared_ptr[CSchema] dataset_schema + shared_ptr[CSchema] projected_schema + c_bool use_threads + c_bool cache_metadata + CExpression filter + + cdef cppclass CScanNodeOptions "arrow::dataset::ScanNodeOptions"(CExecNodeOptions): + CScanNodeOptions(shared_ptr[CDataset] dataset, shared_ptr[CScanOptions] scan_options, bint require_sequenced_output, bint implicit_ordering) + + shared_ptr[CScanOptions] scan_options + + cdef cppclass CFragmentScanOptions "arrow::dataset::FragmentScanOptions": + c_string type_name() const + + ctypedef CIterator[shared_ptr[CScanTask]] CScanTaskIterator \ + "arrow::dataset::ScanTaskIterator" + + cdef cppclass CScanTask" arrow::dataset::ScanTask": + CResult[CRecordBatchIterator] Execute() + + cdef cppclass CFragment "arrow::dataset::Fragment": + CResult[shared_ptr[CSchema]] ReadPhysicalSchema() + CResult[CScanTaskIterator] Scan(shared_ptr[CScanOptions] options) + c_bool splittable() const + c_string type_name() const + const CExpression& partition_expression() const + + ctypedef vector[shared_ptr[CFragment]] CFragmentVector \ + "arrow::dataset::FragmentVector" + + ctypedef CIterator[shared_ptr[CFragment]] CFragmentIterator \ + "arrow::dataset::FragmentIterator" + + cdef cppclass CInMemoryFragment "arrow::dataset::InMemoryFragment"( + CFragment): + CInMemoryFragment(vector[shared_ptr[CRecordBatch]] record_batches, + CExpression partition_expression) + + cdef cppclass CTaggedRecordBatch "arrow::dataset::TaggedRecordBatch": + shared_ptr[CRecordBatch] record_batch + shared_ptr[CFragment] fragment + + ctypedef CIterator[CTaggedRecordBatch] CTaggedRecordBatchIterator \ + "arrow::dataset::TaggedRecordBatchIterator" + + cdef cppclass CScanner "arrow::dataset::Scanner": + CScanner(shared_ptr[CDataset], shared_ptr[CScanOptions]) + CScanner(shared_ptr[CFragment], shared_ptr[CScanOptions]) + CResult[CScanTaskIterator] Scan() + CResult[CTaggedRecordBatchIterator] ScanBatches() + CResult[shared_ptr[CTable]] ToTable() + CResult[shared_ptr[CTable]] TakeRows(const CArray& indices) + CResult[shared_ptr[CTable]] Head(int64_t num_rows) + CResult[int64_t] CountRows() + CResult[CFragmentIterator] GetFragments() + CResult[shared_ptr[CRecordBatchReader]] ToRecordBatchReader() + const shared_ptr[CScanOptions]& options() + + cdef cppclass CScannerBuilder "arrow::dataset::ScannerBuilder": + CScannerBuilder(shared_ptr[CDataset], + shared_ptr[CScanOptions] scan_options) + CScannerBuilder(shared_ptr[CSchema], shared_ptr[CFragment], + shared_ptr[CScanOptions] scan_options) + + @staticmethod + shared_ptr[CScannerBuilder] FromRecordBatchReader( + shared_ptr[CRecordBatchReader] reader) + CStatus ProjectColumns "Project"(const vector[c_string]& columns) + CStatus Project(vector[CExpression]& exprs, vector[c_string]& columns) + CStatus Filter(CExpression filter) + CStatus UseThreads(c_bool use_threads) + CStatus CacheMetadata(c_bool cache_metadata) + CStatus Pool(CMemoryPool* pool) + CStatus BatchSize(int64_t batch_size) + CStatus BatchReadahead(int32_t batch_readahead) + CStatus FragmentReadahead(int32_t fragment_readahead) + CStatus FragmentScanOptions( + shared_ptr[CFragmentScanOptions] fragment_scan_options) + CResult[shared_ptr[CScanOptions]] GetScanOptions() + CResult[shared_ptr[CScanner]] Finish() + shared_ptr[CSchema] schema() const + + ctypedef vector[shared_ptr[CDataset]] CDatasetVector \ + "arrow::dataset::DatasetVector" + + cdef cppclass CDataset "arrow::dataset::Dataset": + const shared_ptr[CSchema] & schema() + CResult[CFragmentIterator] GetFragments() + CResult[CFragmentIterator] GetFragments(CExpression predicate) + const CExpression & partition_expression() + c_string type_name() + + CResult[shared_ptr[CDataset]] ReplaceSchema(shared_ptr[CSchema]) + + CResult[shared_ptr[CScannerBuilder]] NewScan() + + cdef cppclass CInMemoryDataset "arrow::dataset::InMemoryDataset"( + CDataset): + CInMemoryDataset(shared_ptr[CRecordBatchReader]) + CInMemoryDataset(shared_ptr[CTable]) + + cdef cppclass CUnionDataset "arrow::dataset::UnionDataset"( + CDataset): + @staticmethod + CResult[shared_ptr[CUnionDataset]] Make(shared_ptr[CSchema] schema, + CDatasetVector children) + + const CDatasetVector& children() const + + cdef cppclass CInspectOptions "arrow::dataset::InspectOptions": + int fragments + CField.CMergeOptions field_merge_options + + cdef cppclass CFinishOptions "arrow::dataset::FinishOptions": + shared_ptr[CSchema] schema + CInspectOptions inspect_options + c_bool validate_fragments + + cdef cppclass CDatasetFactory "arrow::dataset::DatasetFactory": + CResult[vector[shared_ptr[CSchema]]] InspectSchemas(CInspectOptions) + CResult[shared_ptr[CSchema]] Inspect(CInspectOptions) + CResult[shared_ptr[CDataset]] FinishWithSchema "Finish"( + const shared_ptr[CSchema]& schema) + CResult[shared_ptr[CDataset]] Finish() + const CExpression& root_partition() + CStatus SetRootPartition(CExpression partition) + + cdef cppclass CUnionDatasetFactory "arrow::dataset::UnionDatasetFactory": + @staticmethod + CResult[shared_ptr[CDatasetFactory]] Make( + vector[shared_ptr[CDatasetFactory]] factories) + + cdef cppclass CFileSource "arrow::dataset::FileSource": + const c_string& path() const + const shared_ptr[CFileSystem]& filesystem() const + const shared_ptr[CBuffer]& buffer() const + const int64_t size() const + CResult[shared_ptr[CRandomAccessFile]] Open() const + # HACK: Cython can't handle all the overloads so don't declare them. + # This means invalid construction of CFileSource won't be caught in + # the C++ generation phase (though it will still be caught when + # the generated C++ is compiled). + CFileSource(...) + + cdef cppclass CFileWriteOptions \ + "arrow::dataset::FileWriteOptions": + const shared_ptr[CFileFormat]& format() const + c_string type_name() const + + cdef cppclass CFileWriter \ + "arrow::dataset::FileWriter": + const shared_ptr[CFileFormat]& format() const + const shared_ptr[CSchema]& schema() const + const shared_ptr[CFileWriteOptions]& options() const + const CFileLocator& destination() const + CResult[int64_t] GetBytesWritten() + + cdef cppclass CFileFormat "arrow::dataset::FileFormat": + shared_ptr[CFragmentScanOptions] default_fragment_scan_options + c_string type_name() const + CResult[shared_ptr[CSchema]] Inspect(const CFileSource&) const + CResult[shared_ptr[CFileFragment]] MakeFragment( + CFileSource source, + CExpression partition_expression, + shared_ptr[CSchema] physical_schema) + shared_ptr[CFileWriteOptions] DefaultWriteOptions() + + cdef cppclass CFileFragment "arrow::dataset::FileFragment"( + CFragment): + const CFileSource& source() const + const shared_ptr[CFileFormat]& format() const + + cdef cppclass CFileSystemDatasetWriteOptions \ + "arrow::dataset::FileSystemDatasetWriteOptions": + shared_ptr[CFileWriteOptions] file_write_options + shared_ptr[CFileSystem] filesystem + c_string base_dir + shared_ptr[CPartitioning] partitioning + c_bool preserve_order + int max_partitions + c_string basename_template + function[cb_writer_finish_internal] writer_pre_finish + function[cb_writer_finish_internal] writer_post_finish + ExistingDataBehavior existing_data_behavior + c_bool create_dir + uint32_t max_open_files + uint64_t max_rows_per_file + uint64_t min_rows_per_group + uint64_t max_rows_per_group + + cdef cppclass CFileSystemDataset \ + "arrow::dataset::FileSystemDataset"(CDataset): + @staticmethod + CResult[shared_ptr[CDataset]] Make( + shared_ptr[CSchema] schema, + CExpression source_partition, + shared_ptr[CFileFormat] format, + shared_ptr[CFileSystem] filesystem, + vector[shared_ptr[CFileFragment]] fragments) + + @staticmethod + CStatus Write( + const CFileSystemDatasetWriteOptions& write_options, + shared_ptr[CScanner] scanner) + + c_string type() + vector[c_string] files() + const shared_ptr[CFileFormat]& format() const + const shared_ptr[CFileSystem]& filesystem() const + const shared_ptr[CPartitioning]& partitioning() const + + cdef cppclass CIpcFileWriteOptions \ + "arrow::dataset::IpcFileWriteOptions"(CFileWriteOptions): + shared_ptr[CIpcWriteOptions] options + + cdef cppclass CIpcFileFormat "arrow::dataset::IpcFileFormat"( + CFileFormat): + pass + + cdef cppclass COrcFileFormat "arrow::dataset::OrcFileFormat"( + CFileFormat): + pass + + cdef cppclass CCsvFileWriteOptions \ + "arrow::dataset::CsvFileWriteOptions"(CFileWriteOptions): + shared_ptr[CCSVWriteOptions] write_options + CMemoryPool* pool + + cdef cppclass CCsvFileFormat "arrow::dataset::CsvFileFormat"( + CFileFormat): + CCSVParseOptions parse_options + + cdef cppclass CCsvFragmentScanOptions \ + "arrow::dataset::CsvFragmentScanOptions"(CFragmentScanOptions): + CCSVConvertOptions convert_options + CCSVReadOptions read_options + function[StreamWrapFunc] stream_transform_func + + cdef cppclass CJsonFileFormat "arrow::dataset::JsonFileFormat"(CFileFormat): + pass + + cdef cppclass CJsonFragmentScanOptions "arrow::dataset::JsonFragmentScanOptions"(CFragmentScanOptions): + CJSONParseOptions parse_options + CJSONReadOptions read_options + + cdef struct CPartitionPathFormat "arrow::dataset::PartitionPathFormat": + c_string directory + c_string filename + + cdef cppclass CPartitioning "arrow::dataset::Partitioning": + c_string type_name() const + CResult[CExpression] Parse(const c_string & path) const + CResult[CPartitionPathFormat] Format(const CExpression & expr) const + const shared_ptr[CSchema] & schema() + c_bool Equals(const CPartitioning& other) const + + cdef cppclass CSegmentEncoding" arrow::dataset::SegmentEncoding": + bint operator==(CSegmentEncoding) + + CSegmentEncoding CSegmentEncoding_None\ + " arrow::dataset::SegmentEncoding::None" + CSegmentEncoding CSegmentEncoding_Uri\ + " arrow::dataset::SegmentEncoding::Uri" + + cdef cppclass CKeyValuePartitioningOptions \ + "arrow::dataset::KeyValuePartitioningOptions": + CSegmentEncoding segment_encoding + + cdef cppclass CHivePartitioningOptions \ + "arrow::dataset::HivePartitioningOptions": + CSegmentEncoding segment_encoding + c_string null_fallback + + cdef cppclass CPartitioningFactoryOptions \ + "arrow::dataset::PartitioningFactoryOptions": + c_bool infer_dictionary + shared_ptr[CSchema] schema + CSegmentEncoding segment_encoding + + cdef cppclass CHivePartitioningFactoryOptions \ + "arrow::dataset::HivePartitioningFactoryOptions": + c_bool infer_dictionary + c_string null_fallback + shared_ptr[CSchema] schema + CSegmentEncoding segment_encoding + + cdef cppclass CPartitioningFactory "arrow::dataset::PartitioningFactory": + c_string type_name() const + + cdef cppclass CKeyValuePartitioning \ + "arrow::dataset::KeyValuePartitioning"(CPartitioning): + CKeyValuePartitioning(shared_ptr[CSchema] schema, + vector[shared_ptr[CArray]] dictionaries, + CKeyValuePartitioningOptions options) + + vector[shared_ptr[CArray]] dictionaries() const + CSegmentEncoding segment_encoding() + + cdef cppclass CDirectoryPartitioning \ + "arrow::dataset::DirectoryPartitioning"(CPartitioning): + CDirectoryPartitioning(shared_ptr[CSchema] schema, + vector[shared_ptr[CArray]] dictionaries) + + @staticmethod + shared_ptr[CPartitioningFactory] MakeFactory( + vector[c_string] field_names, CPartitioningFactoryOptions) + + vector[shared_ptr[CArray]] dictionaries() const + + cdef cppclass CHivePartitioning \ + "arrow::dataset::HivePartitioning"(CPartitioning): + CHivePartitioning(shared_ptr[CSchema] schema, + vector[shared_ptr[CArray]] dictionaries, + CHivePartitioningOptions options) + + @staticmethod + shared_ptr[CPartitioningFactory] MakeFactory( + CHivePartitioningFactoryOptions) + + vector[shared_ptr[CArray]] dictionaries() const + c_string null_fallback() const + + cdef cppclass CFilenamePartitioning \ + "arrow::dataset::FilenamePartitioning"(CPartitioning): + CFilenamePartitioning(shared_ptr[CSchema] schema, + vector[shared_ptr[CArray]] dictionaries) + + @staticmethod + shared_ptr[CPartitioningFactory] MakeFactory( + vector[c_string] field_names, CPartitioningFactoryOptions) + + vector[shared_ptr[CArray]] dictionaries() const + + cdef cppclass CPartitioningOrFactory \ + "arrow::dataset::PartitioningOrFactory": + CPartitioningOrFactory(shared_ptr[CPartitioning]) + CPartitioningOrFactory(shared_ptr[CPartitioningFactory]) + CPartitioningOrFactory & operator = (shared_ptr[CPartitioning]) + CPartitioningOrFactory & operator = ( + shared_ptr[CPartitioningFactory]) + shared_ptr[CPartitioning] partitioning() const + shared_ptr[CPartitioningFactory] factory() const + + cdef cppclass CFileSystemFactoryOptions \ + "arrow::dataset::FileSystemFactoryOptions": + CPartitioningOrFactory partitioning + c_string partition_base_dir + c_bool exclude_invalid_files + vector[c_string] selector_ignore_prefixes + + cdef cppclass CFileSystemDatasetFactory \ + "arrow::dataset::FileSystemDatasetFactory"( + CDatasetFactory): + @staticmethod + CResult[shared_ptr[CDatasetFactory]] MakeFromPaths "Make"( + shared_ptr[CFileSystem] filesystem, + vector[c_string] paths, + shared_ptr[CFileFormat] format, + CFileSystemFactoryOptions options + ) + + @staticmethod + CResult[shared_ptr[CDatasetFactory]] MakeFromSelector "Make"( + shared_ptr[CFileSystem] filesystem, + CFileSelector, + shared_ptr[CFileFormat] format, + CFileSystemFactoryOptions options + ) + + @staticmethod + CResult[shared_ptr[CDatasetFactory]] MakeFromFileInfos "Make"( + shared_ptr[CFileSystem] filesystem, + vector[CFileInfo] files, + shared_ptr[CFileFormat] format, + CFileSystemFactoryOptions options + ) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset_parquet.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset_parquet.pxd new file mode 100644 index 0000000000000000000000000000000000000000..8f4917f6c0c80ad969d79a3ce71e035ddf0b0c57 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_dataset_parquet.pxd @@ -0,0 +1,107 @@ +# 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 pyarrow.includes.libarrow_dataset cimport * +from pyarrow.includes.libparquet_encryption cimport * + +from pyarrow._parquet cimport * + + +cdef extern from "arrow/dataset/parquet_encryption_config.h" namespace "arrow::dataset" nogil: + cdef cppclass CParquetEncryptionConfig "arrow::dataset::ParquetEncryptionConfig": + shared_ptr[CCryptoFactory] crypto_factory + shared_ptr[CKmsConnectionConfig] kms_connection_config + shared_ptr[CEncryptionConfiguration] encryption_config + + cdef cppclass CParquetDecryptionConfig "arrow::dataset::ParquetDecryptionConfig": + shared_ptr[CCryptoFactory] crypto_factory + shared_ptr[CKmsConnectionConfig] kms_connection_config + shared_ptr[CDecryptionConfiguration] decryption_config + + +cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil: + + cdef cppclass CParquetFileWriter \ + "arrow::dataset::ParquetFileWriter"(CFileWriter): + const shared_ptr[FileWriter]& parquet_writer() const + + cdef cppclass CParquetFileWriteOptions \ + "arrow::dataset::ParquetFileWriteOptions"(CFileWriteOptions): + shared_ptr[WriterProperties] writer_properties + shared_ptr[ArrowWriterProperties] arrow_writer_properties + shared_ptr[CParquetEncryptionConfig] parquet_encryption_config + + cdef cppclass CParquetFileFragment "arrow::dataset::ParquetFileFragment"( + CFileFragment): + const vector[int]& row_groups() const + shared_ptr[CFileMetaData] metadata() const + CResult[vector[shared_ptr[CFragment]]] SplitByRowGroup( + CExpression predicate) + CResult[shared_ptr[CFragment]] SubsetWithFilter "Subset"( + CExpression predicate) + CResult[shared_ptr[CFragment]] SubsetWithIds "Subset"( + vector[int] row_group_ids) + CStatus EnsureCompleteMetadata() + + cdef cppclass CParquetFileFormatReaderOptions \ + "arrow::dataset::ParquetFileFormat::ReaderOptions": + unordered_set[c_string] dict_columns + TimeUnit coerce_int96_timestamp_unit + Type binary_type + Type list_type + + cdef cppclass CParquetFileFormat "arrow::dataset::ParquetFileFormat"( + CFileFormat): + CParquetFileFormatReaderOptions reader_options + CResult[shared_ptr[CFileFragment]] MakeFragment( + CFileSource source, + CExpression partition_expression, + shared_ptr[CSchema] physical_schema, + vector[int] row_groups) + + cdef cppclass CParquetFragmentScanOptions \ + "arrow::dataset::ParquetFragmentScanOptions"(CFragmentScanOptions): + shared_ptr[CReaderProperties] reader_properties + shared_ptr[ArrowReaderProperties] arrow_reader_properties + shared_ptr[CParquetDecryptionConfig] parquet_decryption_config + + cdef cppclass CParquetFactoryOptions \ + "arrow::dataset::ParquetFactoryOptions": + CPartitioningOrFactory partitioning + c_string partition_base_dir + c_bool validate_column_chunk_paths + + cdef cppclass CParquetDatasetFactory \ + "arrow::dataset::ParquetDatasetFactory"(CDatasetFactory): + @staticmethod + CResult[shared_ptr[CDatasetFactory]] MakeFromMetaDataPath "Make"( + const c_string& metadata_path, + shared_ptr[CFileSystem] filesystem, + shared_ptr[CParquetFileFormat] format, + CParquetFactoryOptions options + ) + + @staticmethod + CResult[shared_ptr[CDatasetFactory]] MakeFromMetaDataSource "Make"( + const CFileSource& metadata_path, + const c_string& base_path, + shared_ptr[CFileSystem] filesystem, + shared_ptr[CParquetFileFormat] format, + CParquetFactoryOptions options + ) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_feather.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_feather.pxd new file mode 100644 index 0000000000000000000000000000000000000000..722e947bfeca238af0bd6ee002e7f3f9f1063ed6 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_feather.pxd @@ -0,0 +1,50 @@ +# 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 pyarrow.includes.libarrow cimport (CCompressionType, CStatus, CTable, + COutputStream, CResult, shared_ptr, + vector, CRandomAccessFile, CSchema, + c_string, CIpcReadOptions) + + +cdef extern from "arrow/ipc/api.h" namespace "arrow::ipc" nogil: + int kFeatherV1Version" arrow::ipc::feather::kFeatherV1Version" + int kFeatherV2Version" arrow::ipc::feather::kFeatherV2Version" + + cdef cppclass CFeatherProperties" arrow::ipc::feather::WriteProperties": + int version + int chunksize + CCompressionType compression + int compression_level + + CStatus WriteFeather" arrow::ipc::feather::WriteTable" \ + (const CTable& table, COutputStream* out, + CFeatherProperties properties) + + cdef cppclass CFeatherReader" arrow::ipc::feather::Reader": + @staticmethod + CResult[shared_ptr[CFeatherReader]] Open( + const shared_ptr[CRandomAccessFile]& file, + const CIpcReadOptions& options) + int version() + shared_ptr[CSchema] schema() + + CStatus Read(shared_ptr[CTable]* out) + CStatus Read(const vector[int] indices, shared_ptr[CTable]* out) + CStatus Read(const vector[c_string] names, shared_ptr[CTable]* out) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_flight.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_flight.pxd new file mode 100644 index 0000000000000000000000000000000000000000..a89137b8455ee3ae25630ce0f125654610f8c351 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_flight.pxd @@ -0,0 +1,621 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * +from pyarrow.includes.libarrow_python cimport CTimePoint + +from libcpp.map cimport multimap + + +cdef extern from "arrow/flight/api.h" namespace "arrow" nogil: + cdef char* CTracingServerMiddlewareName\ + " arrow::flight::TracingServerMiddleware::kMiddlewareName" + + cdef cppclass CActionType" arrow::flight::ActionType": + c_string type + c_string description + bint operator==(CActionType) + CResult[c_string] SerializeToString() + + @staticmethod + CResult[CActionType] Deserialize(const c_string& serialized) + + cdef cppclass CAction" arrow::flight::Action": + c_string type + shared_ptr[CBuffer] body + bint operator==(CAction) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CAction] Deserialize(const c_string& serialized) + + cdef cppclass CFlightResult" arrow::flight::Result": + CFlightResult() + CFlightResult(CFlightResult) + shared_ptr[CBuffer] body + bint operator==(CFlightResult) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CFlightResult] Deserialize(const c_string& serialized) + + cdef cppclass CBasicAuth" arrow::flight::BasicAuth": + CBasicAuth() + CBasicAuth(CBuffer) + CBasicAuth(CBasicAuth) + c_string username + c_string password + bint operator==(CBasicAuth) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CBasicAuth] Deserialize(const c_string& serialized) + + cdef cppclass CResultStream" arrow::flight::ResultStream": + CResult[unique_ptr[CFlightResult]] Next() + + cdef cppclass CDescriptorType \ + " arrow::flight::FlightDescriptor::DescriptorType": + bint operator==(CDescriptorType) + + CDescriptorType CDescriptorTypeUnknown\ + " arrow::flight::FlightDescriptor::UNKNOWN" + CDescriptorType CDescriptorTypePath\ + " arrow::flight::FlightDescriptor::PATH" + CDescriptorType CDescriptorTypeCmd\ + " arrow::flight::FlightDescriptor::CMD" + + cdef cppclass CFlightDescriptor" arrow::flight::FlightDescriptor": + CDescriptorType type + c_string cmd + vector[c_string] path + bint operator==(CFlightDescriptor) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CFlightDescriptor] Deserialize(const c_string& serialized) + + cdef cppclass CTicket" arrow::flight::Ticket": + CTicket() + c_string ticket + bint operator==(CTicket) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CTicket] Deserialize(const c_string& serialized) + + cdef cppclass CCriteria" arrow::flight::Criteria": + CCriteria() + c_string expression + bint operator==(CCriteria) + CResult[c_string] SerializeToString() + + @staticmethod + CResult[CCriteria] Deserialize(const c_string& serialized) + + cdef cppclass CLocation" arrow::flight::Location": + CLocation() + c_string ToString() + c_bool Equals(const CLocation& other) + + @staticmethod + CResult[CLocation] Parse(const c_string& uri_string) + + @staticmethod + CResult[CLocation] ForGrpcTcp(const c_string& host, int port) + + @staticmethod + CResult[CLocation] ForGrpcTls(const c_string& host, int port) + + @staticmethod + CResult[CLocation] ForGrpcUnix(const c_string& path) + + cdef cppclass CFlightEndpoint" arrow::flight::FlightEndpoint": + CFlightEndpoint() + + CTicket ticket + vector[CLocation] locations + optional[CTimePoint] expiration_time + c_string app_metadata + + bint operator==(CFlightEndpoint) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CFlightEndpoint] Deserialize(const c_string& serialized) + + cdef cppclass CFlightInfo" arrow::flight::FlightInfo": + CFlightInfo(CFlightInfo info) + int64_t total_records() + int64_t total_bytes() + c_bool ordered() + c_string app_metadata() + CResult[shared_ptr[CSchema]] GetSchema(CDictionaryMemo* memo) + CFlightDescriptor& descriptor() + const vector[CFlightEndpoint]& endpoints() + CResult[c_string] SerializeToString() + c_string ToString() + bint operator==(CFlightInfo) + + @staticmethod + CResult[unique_ptr[CFlightInfo]] Deserialize( + const c_string& serialized) + + cdef cppclass CSchemaResult" arrow::flight::SchemaResult": + CSchemaResult() + CSchemaResult(CSchemaResult result) + CResult[shared_ptr[CSchema]] GetSchema(CDictionaryMemo* memo) + bint operator==(CSchemaResult) + CResult[c_string] SerializeToString() + c_string ToString() + + @staticmethod + CResult[CSchemaResult] Deserialize(const c_string& serialized) + + cdef cppclass CFlightListing" arrow::flight::FlightListing": + CResult[unique_ptr[CFlightInfo]] Next() + + cdef cppclass CSimpleFlightListing" arrow::flight::SimpleFlightListing": + # This doesn't work with Cython >= 3 + # CSimpleFlightListing(vector[CFlightInfo]&& info) + CSimpleFlightListing(const vector[CFlightInfo]& info) + + cdef cppclass CFlightPayload" arrow::flight::FlightPayload": + shared_ptr[CBuffer] descriptor + shared_ptr[CBuffer] app_metadata + CIpcPayload ipc_message + + cdef cppclass CFlightDataStream" arrow::flight::FlightDataStream": + shared_ptr[CSchema] schema() + CResult[CFlightPayload] Next() + + cdef cppclass CFlightStreamChunk" arrow::flight::FlightStreamChunk": + CFlightStreamChunk() + shared_ptr[CRecordBatch] data + shared_ptr[CBuffer] app_metadata + + cdef cppclass CMetadataRecordBatchReader \ + " arrow::flight::MetadataRecordBatchReader": + CResult[shared_ptr[CSchema]] GetSchema() + CResult[CFlightStreamChunk] Next() + CResult[shared_ptr[CTable]] ToTable() + CIpcReadStats stats() const + + CResult[shared_ptr[CRecordBatchReader]] MakeRecordBatchReader\ + " arrow::flight::MakeRecordBatchReader"( + shared_ptr[CMetadataRecordBatchReader]) + + cdef cppclass CMetadataRecordBatchWriter \ + " arrow::flight::MetadataRecordBatchWriter"(CRecordBatchWriter): + CStatus Begin(shared_ptr[CSchema] schema, + const CIpcWriteOptions& options) + CStatus WriteMetadata(shared_ptr[CBuffer] app_metadata) + CStatus WriteWithMetadata(const CRecordBatch& batch, + shared_ptr[CBuffer] app_metadata) + + cdef cppclass CFlightStreamReader \ + " arrow::flight::FlightStreamReader"(CMetadataRecordBatchReader): + void Cancel() + CResult[shared_ptr[CTable]] ToTableWithStopToken" ToTable"\ + (const CStopToken& stop_token) + + cdef cppclass CFlightMessageReader \ + " arrow::flight::FlightMessageReader"(CMetadataRecordBatchReader): + CFlightDescriptor& descriptor() + + cdef cppclass CFlightMessageWriter \ + " arrow::flight::FlightMessageWriter"(CMetadataRecordBatchWriter): + pass + + cdef cppclass CFlightStreamWriter \ + " arrow::flight::FlightStreamWriter"(CMetadataRecordBatchWriter): + CStatus DoneWriting() + + cdef cppclass CRecordBatchStream \ + " arrow::flight::RecordBatchStream"(CFlightDataStream): + CRecordBatchStream(shared_ptr[CRecordBatchReader]& reader, + const CIpcWriteOptions& options) + + cdef cppclass CFlightMetadataReader" arrow::flight::FlightMetadataReader": + CStatus ReadMetadata(shared_ptr[CBuffer]* out) + + cdef cppclass CFlightMetadataWriter" arrow::flight::FlightMetadataWriter": + CStatus WriteMetadata(const CBuffer& message) + + cdef cppclass CServerAuthReader" arrow::flight::ServerAuthReader": + CStatus Read(c_string* token) + + cdef cppclass CServerAuthSender" arrow::flight::ServerAuthSender": + CStatus Write(c_string& token) + + cdef cppclass CClientAuthReader" arrow::flight::ClientAuthReader": + CStatus Read(c_string* token) + + cdef cppclass CClientAuthSender" arrow::flight::ClientAuthSender": + CStatus Write(c_string& token) + + cdef cppclass CServerAuthHandler" arrow::flight::ServerAuthHandler": + pass + + cdef cppclass CClientAuthHandler" arrow::flight::ClientAuthHandler": + pass + + cdef cppclass CServerCallContext" arrow::flight::ServerCallContext": + c_string& peer_identity() + c_string& peer() + c_bool is_cancelled() + void AddHeader(const c_string& key, const c_string& value) + void AddTrailer(const c_string& key, const c_string& value) + CServerMiddleware* GetMiddleware(const c_string& key) + + cdef cppclass CTimeoutDuration" arrow::flight::TimeoutDuration": + CTimeoutDuration(double) + double count() + + cdef cppclass CFlightCallOptions" arrow::flight::FlightCallOptions": + CFlightCallOptions() + CTimeoutDuration timeout + CIpcWriteOptions write_options + CIpcReadOptions read_options + vector[pair[c_string, c_string]] headers + CStopToken stop_token + + cdef cppclass CCertKeyPair" arrow::flight::CertKeyPair": + CCertKeyPair() + c_string pem_cert + c_string pem_key + + cdef cppclass CFlightMethod" arrow::flight::FlightMethod": + bint operator==(CFlightMethod) + + CFlightMethod CFlightMethodInvalid\ + " arrow::flight::FlightMethod::Invalid" + CFlightMethod CFlightMethodHandshake\ + " arrow::flight::FlightMethod::Handshake" + CFlightMethod CFlightMethodListFlights\ + " arrow::flight::FlightMethod::ListFlights" + CFlightMethod CFlightMethodGetFlightInfo\ + " arrow::flight::FlightMethod::GetFlightInfo" + CFlightMethod CFlightMethodGetSchema\ + " arrow::flight::FlightMethod::GetSchema" + CFlightMethod CFlightMethodDoGet\ + " arrow::flight::FlightMethod::DoGet" + CFlightMethod CFlightMethodDoPut\ + " arrow::flight::FlightMethod::DoPut" + CFlightMethod CFlightMethodDoAction\ + " arrow::flight::FlightMethod::DoAction" + CFlightMethod CFlightMethodListActions\ + " arrow::flight::FlightMethod::ListActions" + CFlightMethod CFlightMethodDoExchange\ + " arrow::flight::FlightMethod::DoExchange" + + cdef cppclass CCallInfo" arrow::flight::CallInfo": + CFlightMethod method + + ctypedef multimap[cpp_string_view, cpp_string_view] CCallHeaders\ + " arrow::flight::CallHeaders" + + cdef cppclass CAddCallHeaders" arrow::flight::AddCallHeaders": + void AddHeader(const c_string& key, const c_string& value) + + cdef cppclass CServerMiddleware" arrow::flight::ServerMiddleware": + c_string name() + + cdef cppclass CServerMiddlewareFactory\ + " arrow::flight::ServerMiddlewareFactory": + pass + + cdef cppclass CClientMiddleware" arrow::flight::ClientMiddleware": + pass + + cdef cppclass CClientMiddlewareFactory\ + " arrow::flight::ClientMiddlewareFactory": + pass + + cpdef cppclass CTracingServerMiddlewareTraceKey\ + " arrow::flight::TracingServerMiddleware::TraceKey": + CTracingServerMiddlewareTraceKey() + c_string key + c_string value + + cdef cppclass CTracingServerMiddleware\ + " arrow::flight::TracingServerMiddleware"(CServerMiddleware): + vector[CTracingServerMiddlewareTraceKey] GetTraceContext() + + cdef shared_ptr[CServerMiddlewareFactory] \ + MakeTracingServerMiddlewareFactory\ + " arrow::flight::MakeTracingServerMiddlewareFactory"() + + cdef cppclass CFlightServerOptions" arrow::flight::FlightServerOptions": + CFlightServerOptions(const CLocation& location) + CLocation location + unique_ptr[CServerAuthHandler] auth_handler + vector[CCertKeyPair] tls_certificates + c_bool verify_client + c_string root_certificates + vector[pair[c_string, shared_ptr[CServerMiddlewareFactory]]] middleware + + cdef cppclass CFlightClientOptions" arrow::flight::FlightClientOptions": + c_string tls_root_certs + c_string cert_chain + c_string private_key + c_string override_hostname + vector[shared_ptr[CClientMiddlewareFactory]] middleware + int64_t write_size_limit_bytes + vector[pair[c_string, CIntStringVariant]] generic_options + c_bool disable_server_verification + + @staticmethod + CFlightClientOptions Defaults() + + cdef cppclass CDoPutResult" arrow::flight::FlightClient::DoPutResult": + unique_ptr[CFlightStreamWriter] writer + unique_ptr[CFlightMetadataReader] reader + + cdef cppclass CDoExchangeResult" arrow::flight::FlightClient::DoExchangeResult": + unique_ptr[CFlightStreamWriter] writer + unique_ptr[CFlightStreamReader] reader + + cdef cppclass CFlightClient" arrow::flight::FlightClient": + @staticmethod + CResult[unique_ptr[CFlightClient]] Connect(const CLocation& location, + const CFlightClientOptions& options) + + c_bool supports_async() + CStatus CheckAsyncSupport() + + CStatus Authenticate(CFlightCallOptions& options, + unique_ptr[CClientAuthHandler] auth_handler) + + CResult[pair[c_string, c_string]] AuthenticateBasicToken( + CFlightCallOptions& options, + const c_string& username, + const c_string& password) + + CResult[unique_ptr[CResultStream]] DoAction(CFlightCallOptions& options, CAction& action) + CResult[vector[CActionType]] ListActions(CFlightCallOptions& options) + + CResult[unique_ptr[CFlightListing]] ListFlights(CFlightCallOptions& options, CCriteria criteria) + CResult[unique_ptr[CFlightInfo]] GetFlightInfo(CFlightCallOptions& options, + CFlightDescriptor& descriptor) + CFuture[CFlightInfo] GetFlightInfoAsync(CFlightCallOptions& options, + CFlightDescriptor& descriptor) + CResult[unique_ptr[CSchemaResult]] GetSchema(CFlightCallOptions& options, + CFlightDescriptor& descriptor) + CResult[unique_ptr[CFlightStreamReader]] DoGet(CFlightCallOptions& options, CTicket& ticket) + CResult[CDoPutResult] DoPut(CFlightCallOptions& options, + CFlightDescriptor& descriptor, + shared_ptr[CSchema]& schema) + CResult[CDoExchangeResult] DoExchange(CFlightCallOptions& options, + CFlightDescriptor& descriptor) + CStatus Close() + + cdef cppclass CFlightStatusCode" arrow::flight::FlightStatusCode": + bint operator==(CFlightStatusCode) + + CFlightStatusCode CFlightStatusInternal \ + " arrow::flight::FlightStatusCode::Internal" + CFlightStatusCode CFlightStatusTimedOut \ + " arrow::flight::FlightStatusCode::TimedOut" + CFlightStatusCode CFlightStatusCancelled \ + " arrow::flight::FlightStatusCode::Cancelled" + CFlightStatusCode CFlightStatusUnauthenticated \ + " arrow::flight::FlightStatusCode::Unauthenticated" + CFlightStatusCode CFlightStatusUnauthorized \ + " arrow::flight::FlightStatusCode::Unauthorized" + CFlightStatusCode CFlightStatusUnavailable \ + " arrow::flight::FlightStatusCode::Unavailable" + CFlightStatusCode CFlightStatusFailed \ + " arrow::flight::FlightStatusCode::Failed" + + cdef cppclass FlightStatusDetail" arrow::flight::FlightStatusDetail": + CFlightStatusCode code() + c_string extra_info() + + @staticmethod + shared_ptr[FlightStatusDetail] UnwrapStatus(const CStatus& status) + + cdef cppclass FlightWriteSizeStatusDetail\ + " arrow::flight::FlightWriteSizeStatusDetail": + int64_t limit() + int64_t actual() + + @staticmethod + shared_ptr[FlightWriteSizeStatusDetail] UnwrapStatus( + const CStatus& status) + + cdef CStatus MakeFlightError" arrow::flight::MakeFlightError" \ + (CFlightStatusCode code, const c_string& message) + + cdef CStatus MakeFlightError" arrow::flight::MakeFlightError" \ + (CFlightStatusCode code, + const c_string& message, + const c_string& extra_info) + +# Callbacks for implementing Flight servers +# Use typedef to emulate syntax for std::function +ctypedef CStatus cb_list_flights(object, const CServerCallContext&, + const CCriteria*, + unique_ptr[CFlightListing]*) +ctypedef CStatus cb_get_flight_info(object, const CServerCallContext&, + const CFlightDescriptor&, + unique_ptr[CFlightInfo]*) +ctypedef CStatus cb_get_schema(object, const CServerCallContext&, + const CFlightDescriptor&, + unique_ptr[CSchemaResult]*) +ctypedef CStatus cb_do_put(object, const CServerCallContext&, + unique_ptr[CFlightMessageReader], + unique_ptr[CFlightMetadataWriter]) +ctypedef CStatus cb_do_get(object, const CServerCallContext&, + const CTicket&, + unique_ptr[CFlightDataStream]*) +ctypedef CStatus cb_do_exchange(object, const CServerCallContext&, + unique_ptr[CFlightMessageReader], + unique_ptr[CFlightMessageWriter]) +ctypedef CStatus cb_do_action(object, const CServerCallContext&, + const CAction&, + unique_ptr[CResultStream]*) +ctypedef CStatus cb_list_actions(object, const CServerCallContext&, + vector[CActionType]*) +ctypedef CStatus cb_result_next(object, unique_ptr[CFlightResult]*) +ctypedef CStatus cb_data_stream_next(object, CFlightPayload*) +ctypedef CStatus cb_server_authenticate(object, CServerAuthSender*, + CServerAuthReader*) +ctypedef CStatus cb_is_valid(object, const c_string&, c_string*) +ctypedef CStatus cb_client_authenticate(object, CClientAuthSender*, + CClientAuthReader*) +ctypedef CStatus cb_get_token(object, c_string*) + +ctypedef CStatus cb_middleware_sending_headers(object, CAddCallHeaders*) +ctypedef CStatus cb_middleware_call_completed(object, const CStatus&) +ctypedef CStatus cb_client_middleware_received_headers( + object, const CCallHeaders&) +ctypedef CStatus cb_server_middleware_start_call( + object, + const CCallInfo&, + const CCallHeaders&, + shared_ptr[CServerMiddleware]*) +ctypedef CStatus cb_client_middleware_start_call( + object, + const CCallInfo&, + unique_ptr[CClientMiddleware]*) + +cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil: + cdef char* CPyServerMiddlewareName\ + " arrow::py::flight::kPyServerMiddlewareName" + + cdef cppclass PyFlightServerVtable: + PyFlightServerVtable() + function[cb_list_flights] list_flights + function[cb_get_flight_info] get_flight_info + function[cb_get_schema] get_schema + function[cb_do_put] do_put + function[cb_do_get] do_get + function[cb_do_exchange] do_exchange + function[cb_do_action] do_action + function[cb_list_actions] list_actions + + cdef cppclass PyServerAuthHandlerVtable: + PyServerAuthHandlerVtable() + function[cb_server_authenticate] authenticate + function[cb_is_valid] is_valid + + cdef cppclass PyClientAuthHandlerVtable: + PyClientAuthHandlerVtable() + function[cb_client_authenticate] authenticate + function[cb_get_token] get_token + + cdef cppclass PyFlightServer: + PyFlightServer(object server, PyFlightServerVtable vtable) + + CStatus Init(CFlightServerOptions& options) + int port() + CStatus ServeWithSignals() except * + CStatus Shutdown() + CStatus Wait() + + cdef cppclass PyServerAuthHandler\ + " arrow::py::flight::PyServerAuthHandler"(CServerAuthHandler): + PyServerAuthHandler(object handler, PyServerAuthHandlerVtable vtable) + + cdef cppclass PyClientAuthHandler\ + " arrow::py::flight::PyClientAuthHandler"(CClientAuthHandler): + PyClientAuthHandler(object handler, PyClientAuthHandlerVtable vtable) + + cdef cppclass CPyFlightResultStream\ + " arrow::py::flight::PyFlightResultStream"(CResultStream): + CPyFlightResultStream(object generator, + function[cb_result_next] callback) + + cdef cppclass CPyFlightDataStream\ + " arrow::py::flight::PyFlightDataStream"(CFlightDataStream): + CPyFlightDataStream(object data_source, + unique_ptr[CFlightDataStream] stream) + + cdef cppclass CPyGeneratorFlightDataStream\ + " arrow::py::flight::PyGeneratorFlightDataStream"\ + (CFlightDataStream): + CPyGeneratorFlightDataStream(object generator, + shared_ptr[CSchema] schema, + function[cb_data_stream_next] callback, + const CIpcWriteOptions& options) + + cdef cppclass PyServerMiddlewareVtable\ + " arrow::py::flight::PyServerMiddleware::Vtable": + PyServerMiddlewareVtable() + function[cb_middleware_sending_headers] sending_headers + function[cb_middleware_call_completed] call_completed + + cdef cppclass PyClientMiddlewareVtable\ + " arrow::py::flight::PyClientMiddleware::Vtable": + PyClientMiddlewareVtable() + function[cb_middleware_sending_headers] sending_headers + function[cb_client_middleware_received_headers] received_headers + function[cb_middleware_call_completed] call_completed + + cdef cppclass CPyServerMiddleware\ + " arrow::py::flight::PyServerMiddleware"(CServerMiddleware): + CPyServerMiddleware(object middleware, PyServerMiddlewareVtable vtable) + void* py_object() + + cdef cppclass CPyServerMiddlewareFactory\ + " arrow::py::flight::PyServerMiddlewareFactory"\ + (CServerMiddlewareFactory): + CPyServerMiddlewareFactory( + object factory, + function[cb_server_middleware_start_call] start_call) + + cdef cppclass CPyClientMiddleware\ + " arrow::py::flight::PyClientMiddleware"(CClientMiddleware): + CPyClientMiddleware(object middleware, PyClientMiddlewareVtable vtable) + + cdef cppclass CPyClientMiddlewareFactory\ + " arrow::py::flight::PyClientMiddlewareFactory"\ + (CClientMiddlewareFactory): + CPyClientMiddlewareFactory( + object factory, + function[cb_client_middleware_start_call] start_call) + + cdef CStatus CreateFlightInfo" arrow::py::flight::CreateFlightInfo"( + shared_ptr[CSchema] schema, + CFlightDescriptor& descriptor, + vector[CFlightEndpoint] endpoints, + int64_t total_records, + int64_t total_bytes, + c_bool ordered, + const c_string& app_metadata, + unique_ptr[CFlightInfo]* out) + + cdef CStatus CreateSchemaResult" arrow::py::flight::CreateSchemaResult"( + shared_ptr[CSchema] schema, + unique_ptr[CSchemaResult]* out) + + +cdef extern from "" namespace "std" nogil: + cdef cppclass CIntStringVariant" std::variant": + CIntStringVariant() + CIntStringVariant(int) + CIntStringVariant(c_string) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_fs.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_fs.pxd new file mode 100644 index 0000000000000000000000000000000000000000..af01c47c8c7b037b6a9f799f367872391f755343 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_fs.pxd @@ -0,0 +1,364 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * +from pyarrow.includes.libarrow_python cimport CTimePoint + +cdef extern from "arrow/filesystem/api.h" namespace "arrow::fs" nogil: + + ctypedef enum CFileType "arrow::fs::FileType": + CFileType_NotFound "arrow::fs::FileType::NotFound" + CFileType_Unknown "arrow::fs::FileType::Unknown" + CFileType_File "arrow::fs::FileType::File" + CFileType_Directory "arrow::fs::FileType::Directory" + + cdef cppclass CFileInfo "arrow::fs::FileInfo": + CFileInfo() + CFileInfo(CFileInfo) + CFileInfo& operator=(CFileInfo) + CFileInfo(const CFileInfo&) + CFileInfo& operator=(const CFileInfo&) + + CFileType type() + void set_type(CFileType type) + c_string path() + void set_path(const c_string& path) + c_string base_name() + int64_t size() + void set_size(int64_t size) + c_string extension() + CTimePoint mtime() + void set_mtime(CTimePoint mtime) + + cdef cppclass CFileSelector "arrow::fs::FileSelector": + CFileSelector() + c_string base_dir + c_bool allow_not_found + c_bool recursive + + cdef cppclass CFileLocator "arrow::fs::FileLocator": + shared_ptr[CFileSystem] filesystem + c_string path + + cdef cppclass CFileSystem "arrow::fs::FileSystem": + shared_ptr[CFileSystem] shared_from_this() + c_string type_name() const + CResult[c_string] NormalizePath(c_string path) + CResult[c_string] MakeUri(c_string path) + CResult[CFileInfo] GetFileInfo(const c_string& path) + CResult[vector[CFileInfo]] GetFileInfo( + const vector[c_string]& paths) + CResult[vector[CFileInfo]] GetFileInfo(const CFileSelector& select) + CStatus CreateDir(const c_string& path, c_bool recursive) + CStatus DeleteDir(const c_string& path) + CStatus DeleteDirContents(const c_string& path, c_bool missing_dir_ok) + CStatus DeleteRootDirContents() + CStatus DeleteFile(const c_string& path) + CStatus DeleteFiles(const vector[c_string]& paths) + CStatus Move(const c_string& src, const c_string& dest) + CStatus CopyFile(const c_string& src, const c_string& dest) + CResult[shared_ptr[CInputStream]] OpenInputStream( + const c_string& path) + CResult[shared_ptr[CRandomAccessFile]] OpenInputFile( + const c_string& path) + CResult[shared_ptr[COutputStream]] OpenOutputStream( + const c_string& path, const shared_ptr[const CKeyValueMetadata]&) + CResult[shared_ptr[COutputStream]] OpenAppendStream( + const c_string& path, const shared_ptr[const CKeyValueMetadata]&) + c_bool Equals(const CFileSystem& other) + c_bool Equals(shared_ptr[CFileSystem] other) + + CResult[shared_ptr[CFileSystem]] CFileSystemFromUri \ + "arrow::fs::FileSystemFromUri"(const c_string& uri) + CResult[shared_ptr[CFileSystem]] CFileSystemFromUri \ + "arrow::fs::FileSystemFromUri"(const c_string& uri, c_string* out_path) + CResult[shared_ptr[CFileSystem]] CFileSystemFromUriOrPath \ + "arrow::fs::FileSystemFromUriOrPath"(const c_string& uri, + c_string* out_path) + + cdef cppclass CFileSystemGlobalOptions \ + "arrow::fs::FileSystemGlobalOptions": + c_string tls_ca_file_path + c_string tls_ca_dir_path + + CStatus CFileSystemsInitialize "arrow::fs::Initialize" \ + (const CFileSystemGlobalOptions& options) + + cdef cppclass CSubTreeFileSystem \ + "arrow::fs::SubTreeFileSystem"(CFileSystem): + CSubTreeFileSystem(const c_string& base_path, + shared_ptr[CFileSystem] base_fs) + c_string base_path() + shared_ptr[CFileSystem] base_fs() + + ctypedef enum CS3LogLevel "arrow::fs::S3LogLevel": + CS3LogLevel_Off "arrow::fs::S3LogLevel::Off" + CS3LogLevel_Fatal "arrow::fs::S3LogLevel::Fatal" + CS3LogLevel_Error "arrow::fs::S3LogLevel::Error" + CS3LogLevel_Warn "arrow::fs::S3LogLevel::Warn" + CS3LogLevel_Info "arrow::fs::S3LogLevel::Info" + CS3LogLevel_Debug "arrow::fs::S3LogLevel::Debug" + CS3LogLevel_Trace "arrow::fs::S3LogLevel::Trace" + + cdef struct CS3GlobalOptions "arrow::fs::S3GlobalOptions": + CS3LogLevel log_level + int num_event_loop_threads + + cdef cppclass CS3ProxyOptions "arrow::fs::S3ProxyOptions": + c_string scheme + c_string host + int port + c_string username + c_string password + c_bool Equals(const CS3ProxyOptions& other) + + @staticmethod + CResult[CS3ProxyOptions] FromUriString "FromUri"( + const c_string& uri_string) + + ctypedef enum CS3CredentialsKind "arrow::fs::S3CredentialsKind": + CS3CredentialsKind_Anonymous "arrow::fs::S3CredentialsKind::Anonymous" + CS3CredentialsKind_Default "arrow::fs::S3CredentialsKind::Default" + CS3CredentialsKind_Explicit "arrow::fs::S3CredentialsKind::Explicit" + CS3CredentialsKind_Role "arrow::fs::S3CredentialsKind::Role" + CS3CredentialsKind_WebIdentity \ + "arrow::fs::S3CredentialsKind::WebIdentity" + + cdef cppclass CS3RetryStrategy "arrow::fs::S3RetryStrategy": + @staticmethod + shared_ptr[CS3RetryStrategy] GetAwsDefaultRetryStrategy(int64_t max_attempts) + + @staticmethod + shared_ptr[CS3RetryStrategy] GetAwsStandardRetryStrategy(int64_t max_attempts) + + cdef cppclass CS3Options "arrow::fs::S3Options": + c_string region + double connect_timeout + double request_timeout + c_string endpoint_override + c_string scheme + c_bool background_writes + c_bool allow_delayed_open + c_bool allow_bucket_creation + c_bool allow_bucket_deletion + c_bool check_directory_existence_before_creation + c_bool force_virtual_addressing + c_string tls_ca_file_path + shared_ptr[const CKeyValueMetadata] default_metadata + c_string role_arn + c_string session_name + c_string external_id + int load_frequency + CS3ProxyOptions proxy_options + CS3CredentialsKind credentials_kind + shared_ptr[CS3RetryStrategy] retry_strategy + void ConfigureDefaultCredentials() + void ConfigureAccessKey(const c_string& access_key, + const c_string& secret_key, + const c_string& session_token) + c_string GetAccessKey() + c_string GetSecretKey() + c_string GetSessionToken() + c_bool Equals(const CS3Options& other) + + @staticmethod + CS3Options Defaults() + + @staticmethod + CS3Options Anonymous() + + @staticmethod + CS3Options FromAccessKey(const c_string& access_key, + const c_string& secret_key, + const c_string& session_token) + + @staticmethod + CS3Options FromAssumeRole(const c_string& role_arn, + const c_string& session_name, + const c_string& external_id, + const int load_frequency) + + cdef cppclass CS3FileSystem "arrow::fs::S3FileSystem"(CFileSystem): + @staticmethod + CResult[shared_ptr[CS3FileSystem]] Make(const CS3Options& options) + CS3Options options() + c_string region() + + cdef CStatus CInitializeS3 "arrow::fs::InitializeS3"( + const CS3GlobalOptions& options) + cdef CStatus CEnsureS3Initialized "arrow::fs::EnsureS3Initialized"() + cdef CStatus CFinalizeS3 "arrow::fs::FinalizeS3"() + cdef CStatus CEnsureS3Finalized "arrow::fs::EnsureS3Finalized"() + + cdef CResult[c_string] ResolveS3BucketRegion(const c_string& bucket) + + cdef cppclass CGcsCredentials "arrow::fs::GcsCredentials": + c_bool anonymous() + CTimePoint expiration() + c_string access_token() + c_string target_service_account() + + cdef cppclass CGcsOptions "arrow::fs::GcsOptions": + CGcsCredentials credentials + c_string endpoint_override + c_string scheme + c_string default_bucket_location + optional[c_string] project_id + optional[double] retry_limit_seconds + shared_ptr[const CKeyValueMetadata] default_metadata + c_bool Equals(const CS3Options& other) + + @staticmethod + CGcsOptions Defaults() + + @staticmethod + CGcsOptions Anonymous() + + @staticmethod + CGcsOptions FromAccessToken(const c_string& access_token, + CTimePoint expiration) + + @staticmethod + CGcsOptions FromImpersonatedServiceAccount(const CGcsCredentials& base_credentials, + c_string& target_service_account) + + cdef cppclass CGcsFileSystem "arrow::fs::GcsFileSystem": + @staticmethod + CResult[shared_ptr[CGcsFileSystem]] Make(const CGcsOptions& options) + CGcsOptions options() + + cdef cppclass CAzureOptions "arrow::fs::AzureOptions": + c_string account_name + c_string blob_storage_authority + c_string dfs_storage_authority + c_string blob_storage_scheme + c_string dfs_storage_scheme + + c_bool Equals(const CAzureOptions& other) + CStatus ConfigureDefaultCredential() + CStatus ConfigureAccountKeyCredential(c_string account_key) + CStatus ConfigureSASCredential(c_string sas_token) + CStatus ConfigureManagedIdentityCredential(c_string client_id) + CStatus ConfigureClientSecretCredential(c_string tenant_id, + c_string client_id, + c_string client_secret) + + cdef cppclass CAzureFileSystem "arrow::fs::AzureFileSystem": + @staticmethod + CResult[shared_ptr[CAzureFileSystem]] Make(const CAzureOptions& options) + CAzureOptions options() + + cdef cppclass CHdfsOptions "arrow::fs::HdfsOptions": + HdfsConnectionConfig connection_config + int32_t buffer_size + int16_t replication + int64_t default_block_size + + @staticmethod + CResult[CHdfsOptions] FromUriString "FromUri"( + const c_string& uri_string) + void ConfigureEndPoint(c_string host, int port) + void ConfigureDriver(c_bool use_hdfs3) + void ConfigureReplication(int16_t replication) + void ConfigureUser(c_string user_name) + void ConfigureBufferSize(int32_t buffer_size) + void ConfigureBlockSize(int64_t default_block_size) + void ConfigureKerberosTicketCachePath(c_string path) + void ConfigureExtraConf(c_string key, c_string value) + + cdef cppclass CHadoopFileSystem "arrow::fs::HadoopFileSystem"(CFileSystem): + @staticmethod + CResult[shared_ptr[CHadoopFileSystem]] Make( + const CHdfsOptions& options) + CHdfsOptions options() + + cdef cppclass CMockFileSystem "arrow::fs::internal::MockFileSystem"( + CFileSystem): + CMockFileSystem(CTimePoint current_time) + + CStatus CCopyFiles "arrow::fs::CopyFiles"( + const vector[CFileLocator]& sources, + const vector[CFileLocator]& destinations, + const CIOContext& io_context, + int64_t chunk_size, c_bool use_threads) + CStatus CCopyFilesWithSelector "arrow::fs::CopyFiles"( + const shared_ptr[CFileSystem]& source_fs, + const CFileSelector& source_sel, + const shared_ptr[CFileSystem]& destination_fs, + const c_string& destination_base_dir, + const CIOContext& io_context, + int64_t chunk_size, c_bool use_threads) + + +# Callbacks for implementing Python filesystems +# Use typedef to emulate syntax for std::function +ctypedef void CallbackGetTypeName(object, c_string*) +ctypedef c_bool CallbackEquals(object, const CFileSystem&) + +ctypedef void CallbackGetFileInfo(object, const c_string&, CFileInfo*) +ctypedef void CallbackGetFileInfoVector(object, const vector[c_string]&, + vector[CFileInfo]*) +ctypedef void CallbackGetFileInfoSelector(object, const CFileSelector&, + vector[CFileInfo]*) +ctypedef void CallbackCreateDir(object, const c_string&, c_bool) +ctypedef void CallbackDeleteDir(object, const c_string&) +ctypedef void CallbackDeleteDirContents(object, const c_string&, c_bool) +ctypedef void CallbackDeleteRootDirContents(object) +ctypedef void CallbackDeleteFile(object, const c_string&) +ctypedef void CallbackMove(object, const c_string&, const c_string&) +ctypedef void CallbackCopyFile(object, const c_string&, const c_string&) + +ctypedef void CallbackOpenInputStream(object, const c_string&, + shared_ptr[CInputStream]*) +ctypedef void CallbackOpenInputFile(object, const c_string&, + shared_ptr[CRandomAccessFile]*) +ctypedef void CallbackOpenOutputStream( + object, const c_string&, const shared_ptr[const CKeyValueMetadata]&, + shared_ptr[COutputStream]*) +ctypedef void CallbackNormalizePath(object, const c_string&, c_string*) + +cdef extern from "arrow/python/filesystem.h" namespace "arrow::py::fs" nogil: + + cdef cppclass CPyFileSystemVtable "arrow::py::fs::PyFileSystemVtable": + PyFileSystemVtable() + function[CallbackGetTypeName] get_type_name + function[CallbackEquals] equals + function[CallbackGetFileInfo] get_file_info + function[CallbackGetFileInfoVector] get_file_info_vector + function[CallbackGetFileInfoSelector] get_file_info_selector + function[CallbackCreateDir] create_dir + function[CallbackDeleteDir] delete_dir + function[CallbackDeleteDirContents] delete_dir_contents + function[CallbackDeleteRootDirContents] delete_root_dir_contents + function[CallbackDeleteFile] delete_file + function[CallbackMove] move + function[CallbackCopyFile] copy_file + function[CallbackOpenInputStream] open_input_stream + function[CallbackOpenInputFile] open_input_file + function[CallbackOpenOutputStream] open_output_stream + function[CallbackOpenOutputStream] open_append_stream + function[CallbackNormalizePath] normalize_path + + cdef cppclass CPyFileSystem "arrow::py::fs::PyFileSystem": + @staticmethod + shared_ptr[CPyFileSystem] Make(object handler, + CPyFileSystemVtable vtable) + + PyObject* handler() diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_python.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_python.pxd new file mode 100644 index 0000000000000000000000000000000000000000..4724c52ccb53460f738ada3fd5d17d1226ad9373 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_python.pxd @@ -0,0 +1,296 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * + + +ctypedef CInvalidRowResult PyInvalidRowCallback(object, + const CCSVInvalidRow&) + + +cdef extern from "arrow/python/csv.h" namespace "arrow::py::csv": + + function[CInvalidRowHandler] MakeInvalidRowHandler( + function[PyInvalidRowCallback], object handler) + + +cdef extern from "arrow/python/api.h" namespace "arrow::py": + # Requires GIL + CResult[shared_ptr[CDataType]] InferArrowType( + object obj, object mask, c_bool pandas_null_sentinels) + + +cdef extern from "arrow/python/api.h" namespace "arrow::py::internal": + object NewMonthDayNanoTupleType() + CResult[PyObject*] MonthDayNanoIntervalArrayToPyList( + const CMonthDayNanoIntervalArray& array) + CResult[PyObject*] MonthDayNanoIntervalScalarToPyObject( + const CMonthDayNanoIntervalScalar& scalar) + + +cdef extern from "arrow/python/arrow_to_pandas.h" namespace "arrow::py::MapConversionType": + cdef enum MapConversionType "arrow::py::MapConversionType": + DEFAULT, + LOSSY, + STRICT_ + + +cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil: + shared_ptr[CDataType] GetPrimitiveType(Type type) + + object PyFloat_FromHalf(uint16_t value) + + cdef cppclass PyConversionOptions: + PyConversionOptions() + + shared_ptr[CDataType] type + int64_t size + CMemoryPool* pool + c_bool from_pandas + c_bool ignore_timezone + c_bool strict + + # TODO Some functions below are not actually "nogil" + + CResult[shared_ptr[CChunkedArray]] ConvertPySequence( + object obj, object mask, const PyConversionOptions& options, + CMemoryPool* pool) + + CResult[shared_ptr[CArray]] Arange(int64_t start, int64_t stop, + int64_t step, CMemoryPool* pool) + + CResult[shared_ptr[CDataType]] NumPyDtypeToArrow(object dtype) + + CStatus NdarrayToArrow(CMemoryPool* pool, object ao, object mo, + c_bool from_pandas, + const shared_ptr[CDataType]& type, + shared_ptr[CChunkedArray]* out) + + CStatus NdarrayToArrow(CMemoryPool* pool, object ao, object mo, + c_bool from_pandas, + const shared_ptr[CDataType]& type, + const CCastOptions& cast_options, + shared_ptr[CChunkedArray]* out) + + CStatus NdarrayToTensor(CMemoryPool* pool, object ao, + const vector[c_string]& dim_names, + shared_ptr[CTensor]* out) + + CStatus TensorToNdarray(const shared_ptr[CTensor]& tensor, object base, + PyObject** out) + + CStatus SparseCOOTensorToNdarray( + const shared_ptr[CSparseCOOTensor]& sparse_tensor, object base, + PyObject** out_data, PyObject** out_coords) + + CStatus SparseCSRMatrixToNdarray( + const shared_ptr[CSparseCSRMatrix]& sparse_tensor, object base, + PyObject** out_data, PyObject** out_indptr, PyObject** out_indices) + + CStatus SparseCSCMatrixToNdarray( + const shared_ptr[CSparseCSCMatrix]& sparse_tensor, object base, + PyObject** out_data, PyObject** out_indptr, PyObject** out_indices) + + CStatus SparseCSFTensorToNdarray( + const shared_ptr[CSparseCSFTensor]& sparse_tensor, object base, + PyObject** out_data, PyObject** out_indptr, PyObject** out_indices) + + CStatus NdarraysToSparseCOOTensor(CMemoryPool* pool, object data_ao, + object coords_ao, + const vector[int64_t]& shape, + const vector[c_string]& dim_names, + shared_ptr[CSparseCOOTensor]* out) + + CStatus NdarraysToSparseCSRMatrix(CMemoryPool* pool, object data_ao, + object indptr_ao, object indices_ao, + const vector[int64_t]& shape, + const vector[c_string]& dim_names, + shared_ptr[CSparseCSRMatrix]* out) + + CStatus NdarraysToSparseCSCMatrix(CMemoryPool* pool, object data_ao, + object indptr_ao, object indices_ao, + const vector[int64_t]& shape, + const vector[c_string]& dim_names, + shared_ptr[CSparseCSCMatrix]* out) + + CStatus NdarraysToSparseCSFTensor(CMemoryPool* pool, object data_ao, + object indptr_ao, object indices_ao, + const vector[int64_t]& shape, + const vector[int64_t]& axis_order, + const vector[c_string]& dim_names, + shared_ptr[CSparseCSFTensor]* out) + + CStatus TensorToSparseCOOTensor(shared_ptr[CTensor], + shared_ptr[CSparseCOOTensor]* out) + + CStatus TensorToSparseCSRMatrix(shared_ptr[CTensor], + shared_ptr[CSparseCSRMatrix]* out) + + CStatus TensorToSparseCSCMatrix(shared_ptr[CTensor], + shared_ptr[CSparseCSCMatrix]* out) + + CStatus TensorToSparseCSFTensor(shared_ptr[CTensor], + shared_ptr[CSparseCSFTensor]* out) + + CStatus ConvertArrayToPandas(const PandasOptions& options, + shared_ptr[CArray] arr, + object py_ref, PyObject** out) + + CStatus ConvertChunkedArrayToPandas(const PandasOptions& options, + shared_ptr[CChunkedArray] arr, + object py_ref, PyObject** out) + + CStatus ConvertTableToPandas(const PandasOptions& options, + shared_ptr[CTable] table, + PyObject** out) + + void c_set_default_memory_pool \ + " arrow::py::set_default_memory_pool"(CMemoryPool* pool)\ + + CMemoryPool* c_get_memory_pool \ + " arrow::py::get_memory_pool"() + + cdef cppclass PyBuffer(CBuffer): + @staticmethod + CResult[shared_ptr[CBuffer]] FromPyObject(object obj) + + cdef cppclass PyForeignBuffer(CBuffer): + @staticmethod + CStatus Make(const uint8_t* data, int64_t size, object base, + shared_ptr[CBuffer]* out) + + cdef cppclass PyReadableFile(CRandomAccessFile): + PyReadableFile(object fo) + + cdef cppclass PyOutputStream(COutputStream): + PyOutputStream(object fo) + + cdef cppclass PandasOptions: + CMemoryPool* pool + c_bool strings_to_categorical + c_bool zero_copy_only + c_bool integer_object_nulls + c_bool date_as_object + c_bool timestamp_as_object + c_bool use_threads + c_bool coerce_temporal_nanoseconds + c_bool ignore_timezone + c_bool deduplicate_objects + c_bool safe_cast + c_bool split_blocks + c_bool self_destruct + MapConversionType maps_as_pydicts + c_bool decode_dictionaries + shared_ptr[const unordered_set[c_string]] categorical_columns + shared_ptr[const unordered_set[c_string]] extension_columns + c_bool to_numpy + + +cdef extern from "arrow/python/api.h" namespace "arrow::py::internal" nogil: + cdef cppclass CTimePoint "arrow::py::internal::TimePoint": + pass + + CTimePoint PyDateTime_to_TimePoint(PyDateTime_DateTime* pydatetime) + int64_t TimePoint_to_ns(CTimePoint val) + CTimePoint TimePoint_from_s(double val) + CTimePoint TimePoint_from_ns(int64_t val) + + CResult[c_string] TzinfoToString(PyObject* pytzinfo) + CResult[PyObject*] StringToTzinfo(c_string) + + +cdef extern from "arrow/python/numpy_init.h" namespace "arrow::py": + int arrow_init_numpy() except -1 + + +cdef extern from "arrow/python/pyarrow.h" namespace "arrow::py": + int import_pyarrow() except -1 + + +cdef extern from "arrow/python/common.h" namespace "arrow::py": + c_bool IsPyError(const CStatus& status) + void RestorePyError(const CStatus& status) except * + + +cdef extern from "arrow/python/common.h" namespace "arrow::py" nogil: + cdef cppclass SharedPtrNoGIL[T](shared_ptr[T]): + # This looks like the only way to satisfy both Cython 2 and Cython 3 + SharedPtrNoGIL& operator=(...) + cdef cppclass UniquePtrNoGIL[T, DELETER=*](unique_ptr[T, DELETER]): + UniquePtrNoGIL& operator=(...) + + +cdef extern from "arrow/python/inference.h" namespace "arrow::py": + c_bool IsPyBool(object o) + c_bool IsPyInt(object o) + c_bool IsPyFloat(object o) + + +cdef extern from "arrow/python/ipc.h" namespace "arrow::py": + cdef cppclass CPyRecordBatchReader" arrow::py::PyRecordBatchReader" \ + (CRecordBatchReader): + @staticmethod + CResult[shared_ptr[CRecordBatchReader]] Make(shared_ptr[CSchema], + object) + + +cdef extern from "arrow/python/ipc.h" namespace "arrow::py" nogil: + cdef cppclass CCastingRecordBatchReader" arrow::py::CastingRecordBatchReader" \ + (CRecordBatchReader): + @staticmethod + CResult[shared_ptr[CRecordBatchReader]] Make(shared_ptr[CRecordBatchReader], + shared_ptr[CSchema]) + + +cdef extern from "arrow/python/extension_type.h" namespace "arrow::py": + cdef cppclass CPyExtensionType \ + " arrow::py::PyExtensionType"(CExtensionType): + @staticmethod + CStatus FromClass(const shared_ptr[CDataType] storage_type, + const c_string extension_name, object typ, + shared_ptr[CExtensionType]* out) + + @staticmethod + CStatus FromInstance(shared_ptr[CDataType] storage_type, + object inst, shared_ptr[CExtensionType]* out) + + object GetInstance() + CStatus SetInstance(object) + + c_string PyExtensionName() + CStatus RegisterPyExtensionType(shared_ptr[CDataType]) + CStatus UnregisterPyExtensionType(c_string type_name) + + +cdef extern from "arrow/python/benchmark.h" namespace "arrow::py::benchmark": + void Benchmark_PandasObjectIsNull(object lst) except * + + +cdef extern from "arrow/python/gdb.h" namespace "arrow::gdb" nogil: + void GdbTestSession "arrow::gdb::TestSession"() + +cdef extern from "arrow/python/helpers.h" namespace "arrow::py::internal": + c_bool IsThreadingEnabled() + +cdef extern from "arrow/python/config.h" namespace "arrow::py": + cdef cppclass CBuildInfo "arrow::py::BuildInfo": + c_string build_type + + const CBuildInfo& GetBuildInfo "arrow::py::GetBuildInfo"() diff --git a/lib/python3.12/site-packages/pyarrow/includes/libarrow_substrait.pxd b/lib/python3.12/site-packages/pyarrow/includes/libarrow_substrait.pxd new file mode 100644 index 0000000000000000000000000000000000000000..865568e2ba6f141c29271527093e33547bafb263 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libarrow_substrait.pxd @@ -0,0 +1,100 @@ +# 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) diff --git a/lib/python3.12/site-packages/pyarrow/includes/libgandiva.pxd b/lib/python3.12/site-packages/pyarrow/includes/libgandiva.pxd new file mode 100644 index 0000000000000000000000000000000000000000..782d83423e827fd1fbfcf66fa73fe4f3ab3d1a1c --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libgandiva.pxd @@ -0,0 +1,298 @@ +# 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.string cimport string as c_string +from libcpp.unordered_set cimport unordered_set as c_unordered_set +from libc.stdint cimport int64_t, int32_t, uint8_t, uintptr_t + +from pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport * + +cdef extern from "gandiva/node.h" namespace "gandiva" nogil: + + cdef cppclass CNode" gandiva::Node": + c_string ToString() + shared_ptr[CDataType] return_type() + + cdef cppclass CGandivaExpression" gandiva::Expression": + c_string ToString() + shared_ptr[CNode] root() + shared_ptr[CField] result() + + ctypedef vector[shared_ptr[CNode]] CNodeVector" gandiva::NodeVector" + + ctypedef vector[shared_ptr[CGandivaExpression]] \ + CExpressionVector" gandiva::ExpressionVector" + +cdef extern from "gandiva/selection_vector.h" namespace "gandiva" nogil: + + cdef cppclass CSelectionVector" gandiva::SelectionVector": + + shared_ptr[CArray] ToArray() + + enum CSelectionVector_Mode" gandiva::SelectionVector::Mode": + CSelectionVector_Mode_NONE" gandiva::SelectionVector::Mode::MODE_NONE" + CSelectionVector_Mode_UINT16" \ + gandiva::SelectionVector::Mode::MODE_UINT16" + CSelectionVector_Mode_UINT32" \ + gandiva::SelectionVector::Mode::MODE_UINT32" + CSelectionVector_Mode_UINT64" \ + gandiva::SelectionVector::Mode::MODE_UINT64" + + cdef CStatus SelectionVector_MakeInt16\ + "gandiva::SelectionVector::MakeInt16"( + int64_t max_slots, CMemoryPool* pool, + shared_ptr[CSelectionVector]* selection_vector) + + cdef CStatus SelectionVector_MakeInt32\ + "gandiva::SelectionVector::MakeInt32"( + int64_t max_slots, CMemoryPool* pool, + shared_ptr[CSelectionVector]* selection_vector) + + cdef CStatus SelectionVector_MakeInt64\ + "gandiva::SelectionVector::MakeInt64"( + int64_t max_slots, CMemoryPool* pool, + shared_ptr[CSelectionVector]* selection_vector) + +cdef inline CSelectionVector_Mode _ensure_selection_mode(str name) except *: + uppercase = name.upper() + if uppercase == 'NONE': + return CSelectionVector_Mode_NONE + elif uppercase == 'UINT16': + return CSelectionVector_Mode_UINT16 + elif uppercase == 'UINT32': + return CSelectionVector_Mode_UINT32 + elif uppercase == 'UINT64': + return CSelectionVector_Mode_UINT64 + else: + raise ValueError(f'Invalid value for Selection Mode: {name!r}') + +cdef inline str _selection_mode_name(CSelectionVector_Mode ctype): + if ctype == CSelectionVector_Mode_NONE: + return 'NONE' + elif ctype == CSelectionVector_Mode_UINT16: + return 'UINT16' + elif ctype == CSelectionVector_Mode_UINT32: + return 'UINT32' + elif ctype == CSelectionVector_Mode_UINT64: + return 'UINT64' + else: + raise RuntimeError('Unexpected CSelectionVector_Mode value') + +cdef extern from "gandiva/condition.h" namespace "gandiva" nogil: + + cdef cppclass CCondition" gandiva::Condition": + c_string ToString() + shared_ptr[CNode] root() + shared_ptr[CField] result() + +cdef extern from "gandiva/arrow.h" namespace "gandiva" nogil: + + ctypedef vector[shared_ptr[CArray]] CArrayVector" gandiva::ArrayVector" + + +cdef extern from "gandiva/tree_expr_builder.h" namespace "gandiva" nogil: + + cdef shared_ptr[CNode] TreeExprBuilder_MakeBoolLiteral \ + "gandiva::TreeExprBuilder::MakeLiteral"(c_bool value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt8Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(uint8_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt16Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(uint16_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt32Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(uint32_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt64Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(uint64_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInt8Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(int8_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInt16Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(int16_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInt32Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(int32_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInt64Literal \ + "gandiva::TreeExprBuilder::MakeLiteral"(int64_t value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeFloatLiteral \ + "gandiva::TreeExprBuilder::MakeLiteral"(float value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeDoubleLiteral \ + "gandiva::TreeExprBuilder::MakeLiteral"(double value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeStringLiteral \ + "gandiva::TreeExprBuilder::MakeStringLiteral"(const c_string& value) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeBinaryLiteral \ + "gandiva::TreeExprBuilder::MakeBinaryLiteral"(const c_string& value) + + cdef shared_ptr[CGandivaExpression] TreeExprBuilder_MakeExpression\ + "gandiva::TreeExprBuilder::MakeExpression"( + shared_ptr[CNode] root_node, shared_ptr[CField] result_field) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeFunction \ + "gandiva::TreeExprBuilder::MakeFunction"( + const c_string& name, const CNodeVector& children, + shared_ptr[CDataType] return_type) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeField \ + "gandiva::TreeExprBuilder::MakeField"(shared_ptr[CField] field) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeIf \ + "gandiva::TreeExprBuilder::MakeIf"( + shared_ptr[CNode] condition, shared_ptr[CNode] this_node, + shared_ptr[CNode] else_node, shared_ptr[CDataType] return_type) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeAnd \ + "gandiva::TreeExprBuilder::MakeAnd"(const CNodeVector& children) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeOr \ + "gandiva::TreeExprBuilder::MakeOr"(const CNodeVector& children) + + cdef shared_ptr[CCondition] TreeExprBuilder_MakeCondition \ + "gandiva::TreeExprBuilder::MakeCondition"( + shared_ptr[CNode] condition) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionInt32 \ + "gandiva::TreeExprBuilder::MakeInExpressionInt32"( + shared_ptr[CNode] node, const c_unordered_set[int32_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionInt64 \ + "gandiva::TreeExprBuilder::MakeInExpressionInt64"( + shared_ptr[CNode] node, const c_unordered_set[int64_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTime32 \ + "gandiva::TreeExprBuilder::MakeInExpressionTime32"( + shared_ptr[CNode] node, const c_unordered_set[int32_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTime64 \ + "gandiva::TreeExprBuilder::MakeInExpressionTime64"( + shared_ptr[CNode] node, const c_unordered_set[int64_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionDate32 \ + "gandiva::TreeExprBuilder::MakeInExpressionDate32"( + shared_ptr[CNode] node, const c_unordered_set[int32_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionDate64 \ + "gandiva::TreeExprBuilder::MakeInExpressionDate64"( + shared_ptr[CNode] node, const c_unordered_set[int64_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTimeStamp \ + "gandiva::TreeExprBuilder::MakeInExpressionTimeStamp"( + shared_ptr[CNode] node, const c_unordered_set[int64_t]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionString \ + "gandiva::TreeExprBuilder::MakeInExpressionString"( + shared_ptr[CNode] node, const c_unordered_set[c_string]& values) + + cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionBinary \ + "gandiva::TreeExprBuilder::MakeInExpressionBinary"( + shared_ptr[CNode] node, const c_unordered_set[c_string]& values) + +cdef extern from "gandiva/projector.h" namespace "gandiva" nogil: + + cdef cppclass CProjector" gandiva::Projector": + + CStatus Evaluate( + const CRecordBatch& batch, CMemoryPool* pool, + const CArrayVector* output) + + CStatus Evaluate( + const CRecordBatch& batch, + const CSelectionVector* selection, + CMemoryPool* pool, + const CArrayVector* output) + + c_string DumpIR() + + cdef CStatus Projector_Make \ + "gandiva::Projector::Make"( + shared_ptr[CSchema] schema, const CExpressionVector& children, + shared_ptr[CProjector]* projector) + + cdef CStatus Projector_Make \ + "gandiva::Projector::Make"( + shared_ptr[CSchema] schema, const CExpressionVector& children, + CSelectionVector_Mode mode, + shared_ptr[CConfiguration] configuration, + shared_ptr[CProjector]* projector) + +cdef extern from "gandiva/filter.h" namespace "gandiva" nogil: + + cdef cppclass CFilter" gandiva::Filter": + + CStatus Evaluate( + const CRecordBatch& batch, + shared_ptr[CSelectionVector] out_selection) + + c_string DumpIR() + + cdef CStatus Filter_Make \ + "gandiva::Filter::Make"( + shared_ptr[CSchema] schema, shared_ptr[CCondition] condition, + shared_ptr[CConfiguration] configuration, + shared_ptr[CFilter]* filter) + +cdef extern from "gandiva/function_signature.h" namespace "gandiva" nogil: + + cdef cppclass CFunctionSignature" gandiva::FunctionSignature": + + CFunctionSignature(const c_string& base_name, + vector[shared_ptr[CDataType]] param_types, + shared_ptr[CDataType] ret_type) + + shared_ptr[CDataType] ret_type() const + + const c_string& base_name() const + + vector[shared_ptr[CDataType]] param_types() const + + c_string ToString() const + +cdef extern from "gandiva/expression_registry.h" namespace "gandiva" nogil: + + cdef vector[shared_ptr[CFunctionSignature]] \ + GetRegisteredFunctionSignatures() + +cdef extern from "gandiva/configuration.h" namespace "gandiva" nogil: + + cdef cppclass CConfiguration" gandiva::Configuration": + + CConfiguration() + + CConfiguration(bint optimize, bint dump_ir) + + void set_optimize(bint optimize) + + void set_dump_ir(bint dump_ir) + + cdef cppclass CConfigurationBuilder \ + " gandiva::ConfigurationBuilder": + @staticmethod + shared_ptr[CConfiguration] DefaultConfiguration() + + CConfigurationBuilder() + + shared_ptr[CConfiguration] build() diff --git a/lib/python3.12/site-packages/pyarrow/includes/libparquet.pxd b/lib/python3.12/site-packages/pyarrow/includes/libparquet.pxd new file mode 100644 index 0000000000000000000000000000000000000000..81901a00acd55c6c00f64b9e1611ca3fe5a2e9f0 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libparquet.pxd @@ -0,0 +1,634 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport (Type, CChunkedArray, CScalar, CSchema, + CStatus, CTable, CMemoryPool, CBuffer, + CKeyValueMetadata, CRandomAccessFile, + COutputStream, CCacheOptions, + TimeUnit, CRecordBatchReader) + + +cdef extern from "parquet/api/schema.h" namespace "parquet::schema" nogil: + cdef cppclass Node: + pass + + cdef cppclass GroupNode(Node): + pass + + cdef cppclass PrimitiveNode(Node): + pass + + cdef cppclass ColumnPath: + c_string ToDotString() + vector[c_string] ToDotVector() + + +cdef extern from "parquet/api/schema.h" namespace "parquet" nogil: + enum ParquetType" parquet::Type::type": + ParquetType_BOOLEAN" parquet::Type::BOOLEAN" + ParquetType_INT32" parquet::Type::INT32" + ParquetType_INT64" parquet::Type::INT64" + ParquetType_INT96" parquet::Type::INT96" + ParquetType_FLOAT" parquet::Type::FLOAT" + ParquetType_DOUBLE" parquet::Type::DOUBLE" + ParquetType_BYTE_ARRAY" parquet::Type::BYTE_ARRAY" + ParquetType_FIXED_LEN_BYTE_ARRAY" parquet::Type::FIXED_LEN_BYTE_ARRAY" + + enum ParquetLogicalTypeId" parquet::LogicalType::Type::type": + ParquetLogicalType_UNDEFINED" parquet::LogicalType::Type::UNDEFINED" + ParquetLogicalType_STRING" parquet::LogicalType::Type::STRING" + ParquetLogicalType_MAP" parquet::LogicalType::Type::MAP" + ParquetLogicalType_LIST" parquet::LogicalType::Type::LIST" + ParquetLogicalType_ENUM" parquet::LogicalType::Type::ENUM" + ParquetLogicalType_DECIMAL" parquet::LogicalType::Type::DECIMAL" + ParquetLogicalType_DATE" parquet::LogicalType::Type::DATE" + ParquetLogicalType_TIME" parquet::LogicalType::Type::TIME" + ParquetLogicalType_TIMESTAMP" parquet::LogicalType::Type::TIMESTAMP" + ParquetLogicalType_INT" parquet::LogicalType::Type::INT" + ParquetLogicalType_FLOAT16" parquet::LogicalType::Type::FLOAT16" + ParquetLogicalType_JSON" parquet::LogicalType::Type::JSON" + ParquetLogicalType_BSON" parquet::LogicalType::Type::BSON" + ParquetLogicalType_UUID" parquet::LogicalType::Type::UUID" + ParquetLogicalType_GEOMETRY" parquet::LogicalType::Type::GEOMETRY" + ParquetLogicalType_GEOGRAPHY" parquet::LogicalType::Type::GEOGRAPHY" + ParquetLogicalType_NONE" parquet::LogicalType::Type::NONE" + + enum ParquetTimeUnit" parquet::LogicalType::TimeUnit::unit": + ParquetTimeUnit_UNKNOWN" parquet::LogicalType::TimeUnit::UNKNOWN" + ParquetTimeUnit_MILLIS" parquet::LogicalType::TimeUnit::MILLIS" + ParquetTimeUnit_MICROS" parquet::LogicalType::TimeUnit::MICROS" + ParquetTimeUnit_NANOS" parquet::LogicalType::TimeUnit::NANOS" + + enum ParquetEdgeInterpolationAlgorithm" parquet::LogicalType::EdgeInterpolationAlgorithm": + ParquetEdgeInterpolationAlgorithm_UNKNOWN" parquet::LogicalType::EdgeInterpolationAlgorithm::UNKNOWN" + ParquetEdgeInterpolationAlgorithm_SPHERICAL" parquet::LogicalType::EdgeInterpolationAlgorithm::SPHERICAL" + ParquetEdgeInterpolationAlgorithm_VINCENTY" parquet::LogicalType::EdgeInterpolationAlgorithm::VINCENTY" + ParquetEdgeInterpolationAlgorithm_THOMAS" parquet::LogicalType::EdgeInterpolationAlgorithm::THOMAS" + ParquetEdgeInterpolationAlgorithm_ANDOYER" parquet::LogicalType::EdgeInterpolationAlgorithm::ANDOYER" + ParquetEdgeInterpolationAlgorithm_KARNEY" parquet::LogicalType::EdgeInterpolationAlgorithm::KARNEY" + + enum ParquetConvertedType" parquet::ConvertedType::type": + ParquetConvertedType_NONE" parquet::ConvertedType::NONE" + ParquetConvertedType_UTF8" parquet::ConvertedType::UTF8" + ParquetConvertedType_MAP" parquet::ConvertedType::MAP" + ParquetConvertedType_MAP_KEY_VALUE \ + " parquet::ConvertedType::MAP_KEY_VALUE" + ParquetConvertedType_LIST" parquet::ConvertedType::LIST" + ParquetConvertedType_ENUM" parquet::ConvertedType::ENUM" + ParquetConvertedType_DECIMAL" parquet::ConvertedType::DECIMAL" + ParquetConvertedType_DATE" parquet::ConvertedType::DATE" + ParquetConvertedType_TIME_MILLIS" parquet::ConvertedType::TIME_MILLIS" + ParquetConvertedType_TIME_MICROS" parquet::ConvertedType::TIME_MICROS" + ParquetConvertedType_TIMESTAMP_MILLIS \ + " parquet::ConvertedType::TIMESTAMP_MILLIS" + ParquetConvertedType_TIMESTAMP_MICROS \ + " parquet::ConvertedType::TIMESTAMP_MICROS" + ParquetConvertedType_UINT_8" parquet::ConvertedType::UINT_8" + ParquetConvertedType_UINT_16" parquet::ConvertedType::UINT_16" + ParquetConvertedType_UINT_32" parquet::ConvertedType::UINT_32" + ParquetConvertedType_UINT_64" parquet::ConvertedType::UINT_64" + ParquetConvertedType_INT_8" parquet::ConvertedType::INT_8" + ParquetConvertedType_INT_16" parquet::ConvertedType::INT_16" + ParquetConvertedType_INT_32" parquet::ConvertedType::INT_32" + ParquetConvertedType_INT_64" parquet::ConvertedType::INT_64" + ParquetConvertedType_JSON" parquet::ConvertedType::JSON" + ParquetConvertedType_BSON" parquet::ConvertedType::BSON" + ParquetConvertedType_INTERVAL" parquet::ConvertedType::INTERVAL" + + enum ParquetRepetition" parquet::Repetition::type": + ParquetRepetition_REQUIRED" parquet::REPETITION::REQUIRED" + ParquetRepetition_OPTIONAL" parquet::REPETITION::OPTIONAL" + ParquetRepetition_REPEATED" parquet::REPETITION::REPEATED" + + enum ParquetEncoding" parquet::Encoding::type": + ParquetEncoding_PLAIN" parquet::Encoding::PLAIN" + ParquetEncoding_PLAIN_DICTIONARY" parquet::Encoding::PLAIN_DICTIONARY" + ParquetEncoding_RLE" parquet::Encoding::RLE" + ParquetEncoding_BIT_PACKED" parquet::Encoding::BIT_PACKED" + ParquetEncoding_DELTA_BINARY_PACKED \ + " parquet::Encoding::DELTA_BINARY_PACKED" + ParquetEncoding_DELTA_LENGTH_BYTE_ARRAY \ + " parquet::Encoding::DELTA_LENGTH_BYTE_ARRAY" + ParquetEncoding_DELTA_BYTE_ARRAY" parquet::Encoding::DELTA_BYTE_ARRAY" + ParquetEncoding_RLE_DICTIONARY" parquet::Encoding::RLE_DICTIONARY" + ParquetEncoding_BYTE_STREAM_SPLIT \ + " parquet::Encoding::BYTE_STREAM_SPLIT" + + enum ParquetCompression" parquet::Compression::type": + ParquetCompression_UNCOMPRESSED" parquet::Compression::UNCOMPRESSED" + ParquetCompression_SNAPPY" parquet::Compression::SNAPPY" + ParquetCompression_GZIP" parquet::Compression::GZIP" + ParquetCompression_LZO" parquet::Compression::LZO" + ParquetCompression_BROTLI" parquet::Compression::BROTLI" + ParquetCompression_LZ4" parquet::Compression::LZ4" + ParquetCompression_ZSTD" parquet::Compression::ZSTD" + + enum ParquetVersion" parquet::ParquetVersion::type": + ParquetVersion_V1" parquet::ParquetVersion::PARQUET_1_0" + ParquetVersion_V2_4" parquet::ParquetVersion::PARQUET_2_4" + ParquetVersion_V2_6" parquet::ParquetVersion::PARQUET_2_6" + + enum ParquetSortOrder" parquet::SortOrder::type": + ParquetSortOrder_SIGNED" parquet::SortOrder::SIGNED" + ParquetSortOrder_UNSIGNED" parquet::SortOrder::UNSIGNED" + ParquetSortOrder_UNKNOWN" parquet::SortOrder::UNKNOWN" + + cdef cppclass CParquetLogicalType" parquet::LogicalType": + c_string ToString() const + c_string ToJSON() const + ParquetLogicalTypeId type() const + + cdef cppclass CParquetDecimalType \ + " parquet::DecimalLogicalType"(CParquetLogicalType): + int32_t precision() const + int32_t scale() const + + cdef cppclass CParquetIntType \ + " parquet::IntLogicalType"(CParquetLogicalType): + int bit_width() const + c_bool is_signed() const + + cdef cppclass CParquetTimeType \ + " parquet::TimeLogicalType"(CParquetLogicalType): + c_bool is_adjusted_to_utc() const + ParquetTimeUnit time_unit() const + + cdef cppclass CParquetTimestampType \ + " parquet::TimestampLogicalType"(CParquetLogicalType): + c_bool is_adjusted_to_utc() const + ParquetTimeUnit time_unit() const + + cdef cppclass CParquetGeometryType \ + " parquet::GeometryLogicalType"(CParquetLogicalType): + c_string crs() const + + cdef cppclass CParquetGeographyType \ + " parquet::GeographyLogicalType"(CParquetLogicalType): + c_string crs() const + ParquetEdgeInterpolationAlgorithm algorithm() const + + cdef cppclass ColumnDescriptor" parquet::ColumnDescriptor": + c_bool Equals(const ColumnDescriptor& other) + + shared_ptr[ColumnPath] path() + int16_t max_definition_level() + int16_t max_repetition_level() + + ParquetType physical_type() + const shared_ptr[const CParquetLogicalType]& logical_type() + ParquetConvertedType converted_type() + const c_string& name() + int type_length() + int type_precision() + int type_scale() + + cdef cppclass SchemaDescriptor: + const ColumnDescriptor* Column(int i) + shared_ptr[Node] schema() + GroupNode* group() + c_bool Equals(const SchemaDescriptor& other) + c_string ToString() + int num_columns() + + cdef c_string FormatStatValue(ParquetType parquet_type, c_string val) + + enum ParquetCipher" parquet::ParquetCipher::type": + ParquetCipher_AES_GCM_V1" parquet::ParquetCipher::AES_GCM_V1" + ParquetCipher_AES_GCM_CTR_V1" parquet::ParquetCipher::AES_GCM_CTR_V1" + + struct AadMetadata: + c_string aad_prefix + c_string aad_file_unique + c_bool supply_aad_prefix + + struct EncryptionAlgorithm: + ParquetCipher algorithm + AadMetadata aad + + +# Specific array<> types needed for GeoStatistics +cdef extern from "" namespace "std" nogil: + cdef cppclass double_array4 "std::array": + double_array4() except + + double& operator[](size_t) + + cdef cppclass bool_array4 "std::array": + bool_array4() except + + c_bool& operator[](size_t) + + +cdef extern from "parquet/geospatial/statistics.h" namespace "parquet" nogil: + + cdef cppclass CParquetGeoStatistics" parquet::geospatial::GeoStatistics": + c_bool is_valid() const + + double_array4 lower_bound() const + double_array4 upper_bound() const + bool_array4 dimension_valid() const + bool_array4 dimension_empty() const + + optional[vector[int32_t]] geometry_types() const + + c_string ToString() const + + +cdef extern from "parquet/api/reader.h" namespace "parquet" nogil: + cdef cppclass ColumnReader: + pass + + cdef cppclass BoolReader(ColumnReader): + pass + + cdef cppclass Int32Reader(ColumnReader): + pass + + cdef cppclass Int64Reader(ColumnReader): + pass + + cdef cppclass Int96Reader(ColumnReader): + pass + + cdef cppclass FloatReader(ColumnReader): + pass + + cdef cppclass DoubleReader(ColumnReader): + pass + + cdef cppclass ByteArrayReader(ColumnReader): + pass + + cdef cppclass RowGroupReader: + pass + + cdef cppclass CEncodedStatistics" parquet::EncodedStatistics": + const c_string& max() const + const c_string& min() const + int64_t null_count + int64_t distinct_count + bint has_min + bint has_max + bint has_null_count + bint has_distinct_count + + cdef cppclass ParquetByteArray" parquet::ByteArray": + uint32_t len + const uint8_t* ptr + + cdef cppclass ParquetFLBA" parquet::FLBA": + const uint8_t* ptr + + cdef cppclass CStatistics" parquet::Statistics": + int64_t null_count() const + int64_t distinct_count() const + int64_t num_values() const + bint HasMinMax() + bint HasNullCount() + bint HasDistinctCount() + c_bool Equals(const CStatistics&) const + void Reset() + c_string EncodeMin() + c_string EncodeMax() + CEncodedStatistics Encode() + void SetComparator() + ParquetType physical_type() const + const ColumnDescriptor* descr() const + + cdef cppclass CBoolStatistics" parquet::BoolStatistics"(CStatistics): + c_bool min() + c_bool max() + + cdef cppclass CInt32Statistics" parquet::Int32Statistics"(CStatistics): + int32_t min() + int32_t max() + + cdef cppclass CInt64Statistics" parquet::Int64Statistics"(CStatistics): + int64_t min() + int64_t max() + + cdef cppclass CFloatStatistics" parquet::FloatStatistics"(CStatistics): + float min() + float max() + + cdef cppclass CDoubleStatistics" parquet::DoubleStatistics"(CStatistics): + double min() + double max() + + cdef cppclass CByteArrayStatistics \ + " parquet::ByteArrayStatistics"(CStatistics): + ParquetByteArray min() + ParquetByteArray max() + + cdef cppclass CFLBAStatistics" parquet::FLBAStatistics"(CStatistics): + ParquetFLBA min() + ParquetFLBA max() + + cdef cppclass CColumnCryptoMetaData" parquet::ColumnCryptoMetaData": + shared_ptr[ColumnPath] path_in_schema() const + c_bool encrypted_with_footer_key() const + const c_string& key_metadata() const + + cdef cppclass ParquetIndexLocation" parquet::IndexLocation": + int64_t offset + int32_t length + + cdef cppclass CColumnChunkMetaData" parquet::ColumnChunkMetaData": + int64_t file_offset() const + const c_string& file_path() const + + c_bool is_metadata_set() const + ParquetType type() const + int64_t num_values() const + shared_ptr[ColumnPath] path_in_schema() const + bint is_stats_set() const + shared_ptr[CStatistics] statistics() const + c_bool is_geo_stats_set() const + shared_ptr[CParquetGeoStatistics] geo_statistics() const + ParquetCompression compression() const + const vector[ParquetEncoding]& encodings() const + c_bool Equals(const CColumnChunkMetaData&) const + + int64_t has_dictionary_page() const + int64_t dictionary_page_offset() const + int64_t data_page_offset() const + int64_t index_page_offset() const + int64_t total_compressed_size() const + int64_t total_uncompressed_size() const + unique_ptr[CColumnCryptoMetaData] crypto_metadata() const + optional[ParquetIndexLocation] GetColumnIndexLocation() const + optional[ParquetIndexLocation] GetOffsetIndexLocation() const + shared_ptr[const CKeyValueMetadata] key_value_metadata() const + + struct CSortingColumn" parquet::SortingColumn": + int column_idx + c_bool descending + c_bool nulls_first + + cdef cppclass CRowGroupMetaData" parquet::RowGroupMetaData": + c_bool Equals(const CRowGroupMetaData&) const + int num_columns() const + int64_t num_rows() const + int64_t total_byte_size() const + vector[CSortingColumn] sorting_columns() const + unique_ptr[CColumnChunkMetaData] ColumnChunk(int i) const + + cdef cppclass CFileMetaData" parquet::FileMetaData": + c_bool Equals(const CFileMetaData&) const + uint32_t size() + int num_columns() + int64_t num_rows() + int num_row_groups() + ParquetVersion version() + const c_string created_by() + int num_schema_elements() + + void set_file_path(const c_string& path) + void AppendRowGroups(const CFileMetaData& other) except + + + unique_ptr[CRowGroupMetaData] RowGroup(int i) + const SchemaDescriptor* schema() + shared_ptr[const CKeyValueMetadata] key_value_metadata() const + void WriteTo(COutputStream* dst) const + + inline c_bool is_encryption_algorithm_set() const + inline EncryptionAlgorithm encryption_algorithm() const + inline const c_string& footer_signing_key_metadata() const + + cdef shared_ptr[CFileMetaData] CFileMetaData_Make \ + " parquet::FileMetaData::Make"(const void* serialized_metadata, + uint32_t* metadata_len) + + cdef cppclass CReaderProperties" parquet::ReaderProperties": + c_bool is_buffered_stream_enabled() const + void enable_buffered_stream() + void disable_buffered_stream() + + void set_buffer_size(int64_t buf_size) + int64_t buffer_size() const + + void set_thrift_string_size_limit(int32_t size) + int32_t thrift_string_size_limit() const + + void set_thrift_container_size_limit(int32_t size) + int32_t thrift_container_size_limit() const + + void file_decryption_properties(shared_ptr[CFileDecryptionProperties] + decryption) + shared_ptr[CFileDecryptionProperties] file_decryption_properties() \ + const + + c_bool page_checksum_verification() const + void set_page_checksum_verification(c_bool check_crc) + + CReaderProperties default_reader_properties() + + cdef cppclass ArrowReaderProperties: + ArrowReaderProperties() + void set_binary_type(Type binary_type) + Type binary_type() + void set_list_type(Type list_type) + Type list_type() + void set_read_dictionary(int column_index, c_bool read_dict) + c_bool read_dictionary(int column_index) + void set_batch_size(int64_t batch_size) + int64_t batch_size() + void set_pre_buffer(c_bool pre_buffer) + c_bool pre_buffer() const + void set_cache_options(CCacheOptions options) + CCacheOptions cache_options() const + void set_coerce_int96_timestamp_unit(TimeUnit unit) + TimeUnit coerce_int96_timestamp_unit() const + void set_arrow_extensions_enabled(c_bool extensions_enabled) + c_bool get_arrow_extensions_enabled() const + + ArrowReaderProperties default_arrow_reader_properties() + + cdef cppclass ParquetFileReader: + shared_ptr[CFileMetaData] metadata() + + +cdef extern from "parquet/api/writer.h" namespace "parquet" nogil: + cdef cppclass CdcOptions: + int64_t min_chunk_size + int64_t max_chunk_size + int norm_level + + cdef cppclass WriterProperties: + cppclass Builder: + Builder* data_page_version(ParquetDataPageVersion version) + Builder* version(ParquetVersion version) + Builder* compression(ParquetCompression codec) + Builder* compression(const c_string& path, + ParquetCompression codec) + Builder* compression_level(int compression_level) + Builder* compression_level(const c_string& path, + int compression_level) + Builder* encryption( + shared_ptr[CFileEncryptionProperties] + file_encryption_properties) + Builder* disable_dictionary() + Builder* enable_dictionary() + Builder* enable_dictionary(const c_string& path) + Builder* set_sorting_columns(vector[CSortingColumn] sorting_columns) + Builder* disable_statistics() + Builder* enable_statistics() + Builder* enable_statistics(const c_string& path) + Builder* enable_store_decimal_as_integer() + Builder* disable_store_decimal_as_integer() + Builder* data_pagesize(int64_t size) + Builder* max_rows_per_page(int64_t max_rows) + Builder* encoding(ParquetEncoding encoding) + Builder* encoding(const c_string& path, + ParquetEncoding encoding) + Builder* max_row_group_length(int64_t size) + Builder* write_batch_size(int64_t batch_size) + Builder* dictionary_pagesize_limit(int64_t dictionary_pagesize_limit) + Builder* enable_write_page_index() + Builder* disable_write_page_index() + Builder* enable_page_checksum() + Builder* disable_page_checksum() + Builder* enable_content_defined_chunking() + Builder* disable_content_defined_chunking() + Builder* content_defined_chunking_options(CdcOptions options) + shared_ptr[WriterProperties] build() + + cdef cppclass ArrowWriterProperties: + cppclass Builder: + Builder() + Builder* disable_deprecated_int96_timestamps() + Builder* enable_deprecated_int96_timestamps() + Builder* coerce_timestamps(TimeUnit unit) + Builder* allow_truncated_timestamps() + Builder* disallow_truncated_timestamps() + Builder* store_schema() + Builder* enable_compliant_nested_types() + Builder* disable_compliant_nested_types() + Builder* set_engine_version(ArrowWriterEngineVersion version) + Builder* set_time_adjusted_to_utc(c_bool adjusted) + shared_ptr[ArrowWriterProperties] build() + c_bool support_deprecated_int96_timestamps() + +cdef extern from "parquet/arrow/reader.h" namespace "parquet::arrow" nogil: + cdef cppclass FileReader: + FileReader(CMemoryPool* pool, unique_ptr[ParquetFileReader] reader) + + CStatus GetSchema(shared_ptr[CSchema]* out) + + CStatus ReadColumn(int i, shared_ptr[CChunkedArray]* out) + CStatus ReadSchemaField(int i, shared_ptr[CChunkedArray]* out) + + int num_row_groups() + CStatus ReadRowGroup(int i, shared_ptr[CTable]* out) + CStatus ReadRowGroup(int i, const vector[int]& column_indices, + shared_ptr[CTable]* out) + + CStatus ReadRowGroups(const vector[int]& row_groups, + shared_ptr[CTable]* out) + CStatus ReadRowGroups(const vector[int]& row_groups, + const vector[int]& column_indices, + shared_ptr[CTable]* out) + + CResult[unique_ptr[CRecordBatchReader]] GetRecordBatchReader(const vector[int]& row_group_indices, + const vector[int]& column_indices) + CResult[unique_ptr[CRecordBatchReader]] GetRecordBatchReader(const vector[int]& row_group_indices) + + CStatus ReadTable(shared_ptr[CTable]* out) + CStatus ReadTable(const vector[int]& column_indices, + shared_ptr[CTable]* out) + + CStatus ScanContents(vector[int] columns, int32_t column_batch_size, + int64_t* num_rows) + + const ParquetFileReader* parquet_reader() + + void set_use_threads(c_bool use_threads) + + void set_batch_size(int64_t batch_size) + + cdef cppclass FileReaderBuilder: + FileReaderBuilder() + CStatus Open(const shared_ptr[CRandomAccessFile]& file, + const CReaderProperties& properties, + const shared_ptr[CFileMetaData]& metadata) + + ParquetFileReader* raw_reader() + FileReaderBuilder* memory_pool(CMemoryPool*) + FileReaderBuilder* properties(const ArrowReaderProperties&) + CStatus Build(unique_ptr[FileReader]* out) + + CStatus FromParquetSchema( + const SchemaDescriptor* parquet_schema, + const ArrowReaderProperties& properties, + const shared_ptr[const CKeyValueMetadata]& key_value_metadata, + shared_ptr[CSchema]* out) + + CStatus StatisticsAsScalars(const CStatistics& Statistics, + shared_ptr[CScalar]* min, + shared_ptr[CScalar]* max) + +cdef extern from "parquet/arrow/schema.h" namespace "parquet::arrow" nogil: + + CStatus ToParquetSchema( + const CSchema* arrow_schema, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, + shared_ptr[SchemaDescriptor]* out) + + +cdef extern from "parquet/properties.h" namespace "parquet" nogil: + cdef enum ArrowWriterEngineVersion: + V1 "parquet::ArrowWriterProperties::V1", + V2 "parquet::ArrowWriterProperties::V2" + + cdef cppclass ParquetDataPageVersion: + pass + + cdef ParquetDataPageVersion ParquetDataPageVersion_V1 \ + " parquet::ParquetDataPageVersion::V1" + cdef ParquetDataPageVersion ParquetDataPageVersion_V2 \ + " parquet::ParquetDataPageVersion::V2" + +cdef extern from "parquet/arrow/writer.h" namespace "parquet::arrow" nogil: + cdef cppclass FileWriter: + + @staticmethod + CResult[unique_ptr[FileWriter]] Open(const CSchema& schema, CMemoryPool* pool, + const shared_ptr[COutputStream]& sink, + const shared_ptr[WriterProperties]& properties, + const shared_ptr[ArrowWriterProperties]& arrow_properties) + + CStatus WriteTable(const CTable& table, int64_t chunk_size) + CStatus NewRowGroup() + CStatus Close() + CStatus AddKeyValueMetadata(const shared_ptr[const CKeyValueMetadata]& key_value_metadata) + + const shared_ptr[CFileMetaData] metadata() const + + CStatus WriteMetaDataFile( + const CFileMetaData& file_metadata, + const COutputStream* sink) + +cdef extern from "parquet/encryption/encryption.h" namespace "parquet" nogil: + cdef cppclass CFileDecryptionProperties\ + " parquet::FileDecryptionProperties": + pass + + cdef cppclass CFileEncryptionProperties\ + " parquet::FileEncryptionProperties": + pass diff --git a/lib/python3.12/site-packages/pyarrow/includes/libparquet_encryption.pxd b/lib/python3.12/site-packages/pyarrow/includes/libparquet_encryption.pxd new file mode 100644 index 0000000000000000000000000000000000000000..7024f14ac27e8ba1c65ac6a06d3bbb35bd75d886 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/includes/libparquet_encryption.pxd @@ -0,0 +1,187 @@ +# 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 pyarrow.includes.common cimport * +from pyarrow.includes.libarrow cimport CSecureString +from pyarrow.includes.libarrow_fs cimport CFileSystem +from pyarrow._parquet cimport (ParquetCipher, + CFileEncryptionProperties, + CFileDecryptionProperties, + ParquetCipher_AES_GCM_V1, + ParquetCipher_AES_GCM_CTR_V1) + + +cdef extern from "parquet/encryption/kms_client.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CKmsClient" parquet::encryption::KmsClient": + c_string WrapKey(const CSecureString& key, + const c_string& master_key_identifier) except + + CSecureString UnwrapKey(const c_string& wrapped_key, + const c_string& master_key_identifier) except + + + cdef cppclass CKeyAccessToken" parquet::encryption::KeyAccessToken": + CKeyAccessToken(const c_string value) + void Refresh(const c_string& new_value) + const c_string& value() const + + cdef cppclass CKmsConnectionConfig \ + " parquet::encryption::KmsConnectionConfig": + CKmsConnectionConfig() + c_string kms_instance_id + c_string kms_instance_url + shared_ptr[CKeyAccessToken] refreshable_key_access_token + unordered_map[c_string, c_string] custom_kms_conf + +# Callbacks for implementing Python kms clients +# Use typedef to emulate syntax for std::function +ctypedef void CallbackWrapKey( + object, const CSecureString&, const c_string&, c_string*) +ctypedef void CallbackUnwrapKey( + object, const c_string&, const c_string&, CSecureString*) + +cdef extern from "parquet/encryption/kms_client_factory.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CKmsClientFactory" parquet::encryption::KmsClientFactory": + shared_ptr[CKmsClient] CreateKmsClient( + const CKmsConnectionConfig& kms_connection_config) except + + +# Callbacks for implementing Python kms client factories +# Use typedef to emulate syntax for std::function +ctypedef void CallbackCreateKmsClient( + object, + const CKmsConnectionConfig&, shared_ptr[CKmsClient]*) + +cdef extern from "parquet/encryption/crypto_factory.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CEncryptionConfiguration\ + " parquet::encryption::EncryptionConfiguration": + CEncryptionConfiguration(const c_string& footer_key) except + + c_string footer_key + c_string column_keys + c_bool uniform_encryption + ParquetCipher encryption_algorithm + c_bool plaintext_footer + c_bool double_wrapping + double cache_lifetime_seconds + c_bool internal_key_material + int32_t data_key_length_bits + + cdef cppclass CDecryptionConfiguration\ + " parquet::encryption::DecryptionConfiguration": + CDecryptionConfiguration() except + + double cache_lifetime_seconds + + cdef cppclass CCryptoFactory" parquet::encryption::CryptoFactory": + void RegisterKmsClientFactory( + shared_ptr[CKmsClientFactory] kms_client_factory) except + + shared_ptr[CFileEncryptionProperties] GetFileEncryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CEncryptionConfiguration& encryption_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] file_system) except +* + shared_ptr[CFileDecryptionProperties] GetFileDecryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CDecryptionConfiguration& decryption_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] file_system) except +* + void RemoveCacheEntriesForToken(const c_string& access_token) except + + void RemoveCacheEntriesForAllTokens() except + + void RotateMasterKeys(const CKmsConnectionConfig& kms_connection_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] file_system, + c_bool double_wrapping, + double cache_lifetime_seconds) + +cdef extern from "parquet/encryption/file_key_material_store.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CFileKeyMaterialStore\ + "parquet::encryption::FileKeyMaterialStore": + @staticmethod + c_string GetKeyMaterial(c_string key_id_in_file) except + + vector[c_string] GetKeyIDSet() except + + +cdef extern from "parquet/encryption/file_system_key_material_store.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CFileSystemKeyMaterialStore\ + "parquet::encryption::FileSystemKeyMaterialStore": + + @staticmethod + shared_ptr[CFileSystemKeyMaterialStore] Make(c_string parquet_file_path, + shared_ptr[CFileSystem] file_system, + c_bool use_tmp_prefix) except + + + c_string GetKeyMaterial(c_string key_id_in_file) except + + + vector[c_string] GetKeyIDSet() except + + +cdef extern from "parquet/encryption/key_material.h" \ + namespace "parquet::encryption" nogil: + cdef cppclass CKeyMaterial "parquet::encryption::KeyMaterial": + @staticmethod + CKeyMaterial Parse(const c_string& key_material_string) + c_bool is_footer_key() + c_bool is_double_wrapped() + const c_string& master_key_id() + const c_string& wrapped_dek() + const c_string& kek_id() + const c_string& wrapped_kek() + const c_string& kms_instance_id() + const c_string& kms_instance_url() + +cdef extern from "arrow/python/parquet_encryption.h" \ + namespace "arrow::py::parquet::encryption" nogil: + cdef cppclass CPyKmsClientVtable \ + " arrow::py::parquet::encryption::PyKmsClientVtable": + CPyKmsClientVtable() + function[CallbackWrapKey] wrap_key + function[CallbackUnwrapKey] unwrap_key + + cdef cppclass CPyKmsClient\ + " arrow::py::parquet::encryption::PyKmsClient"(CKmsClient): + CPyKmsClient(object handler, CPyKmsClientVtable vtable) + + cdef cppclass CPyKmsClientFactoryVtable\ + " arrow::py::parquet::encryption::PyKmsClientFactoryVtable": + CPyKmsClientFactoryVtable() + function[CallbackCreateKmsClient] create_kms_client + + cdef cppclass CPyKmsClientFactory\ + " arrow::py::parquet::encryption::PyKmsClientFactory"( + CKmsClientFactory): + CPyKmsClientFactory(object handler, CPyKmsClientFactoryVtable vtable) + + cdef cppclass CPyCryptoFactory\ + " arrow::py::parquet::encryption::PyCryptoFactory"(CCryptoFactory): + CResult[shared_ptr[CFileEncryptionProperties]] \ + SafeGetFileEncryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CEncryptionConfiguration& encryption_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] filesystem) + CResult[shared_ptr[CFileDecryptionProperties]] \ + SafeGetFileDecryptionProperties( + const CKmsConnectionConfig& kms_connection_config, + const CDecryptionConfiguration& decryption_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] filesystem) + CStatus SafeRotateMasterKeys(const CKmsConnectionConfig& kms_connection_config, + const c_string parquet_file_path, + const shared_ptr[CFileSystem] filesystem, + c_bool double_wrapping, + double cache_lifetime_seconds) diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__init__.py b/lib/python3.12/site-packages/pyarrow/interchange/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7ebe59b499c214dd82954bff84824cfea574b415 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/interchange/__init__.py @@ -0,0 +1,20 @@ +# 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. + +# flake8: noqa + +from .from_dataframe import from_dataframe diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ccd1ffcf48ad4a5d639b4af4227cc42d3716c84 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/buffer.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36ca9790f9d05441d2152a225a187a3cfdd2ec0c Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/column.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/column.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aa16f4c9a2bd600d0ec85fdb0e6accaef06b4ca9 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/column.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/dataframe.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/dataframe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..328e63209a6888d411a778353b2556cb9bebdb9d Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/dataframe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/from_dataframe.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/from_dataframe.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c03cd196ad7ff38d253a2fdf59a75020ed59ce71 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/interchange/__pycache__/from_dataframe.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/interchange/buffer.py b/lib/python3.12/site-packages/pyarrow/interchange/buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..1f537798130b9a77bc50e1040ea8046557974894 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/interchange/buffer.py @@ -0,0 +1,107 @@ +# 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. + +from __future__ import annotations +import enum + +import pyarrow as pa + + +class DlpackDeviceType(enum.IntEnum): + """Integer enum for device type codes matching DLPack.""" + + CPU = 1 + CUDA = 2 + CPU_PINNED = 3 + OPENCL = 4 + VULKAN = 7 + METAL = 8 + VPI = 9 + ROCM = 10 + + +class _PyArrowBuffer: + """ + Data in the buffer is guaranteed to be contiguous in memory. + + Note that there is no dtype attribute present, a buffer can be thought of + as simply a block of memory. However, if the column that the buffer is + attached to has a dtype that's supported by DLPack and ``__dlpack__`` is + implemented, then that dtype information will be contained in the return + value from ``__dlpack__``. + + This distinction is useful to support both data exchange via DLPack on a + buffer and (b) dtypes like variable-length strings which do not have a + fixed number of bytes per element. + """ + + def __init__(self, x: pa.Buffer, allow_copy: bool = True) -> None: + """ + Handle PyArrow Buffers. + """ + self._x = x + + @property + def bufsize(self) -> int: + """ + Buffer size in bytes. + """ + return self._x.size + + @property + def ptr(self) -> int: + """ + Pointer to start of the buffer as an integer. + """ + return self._x.address + + def __dlpack__(self): + """ + Produce DLPack capsule (see array API standard). + + Raises: + - TypeError : if the buffer contains unsupported dtypes. + - NotImplementedError : if DLPack support is not implemented + + Useful to have to connect to array libraries. Support optional because + it's not completely trivial to implement for a Python-only library. + """ + raise NotImplementedError("__dlpack__") + + def __dlpack_device__(self) -> tuple[DlpackDeviceType, int | None]: + """ + Device type and device ID for where the data in the buffer resides. + Uses device type codes matching DLPack. + Note: must be implemented even if ``__dlpack__`` is not. + """ + if self._x.is_cpu: + return (DlpackDeviceType.CPU, None) + else: + raise NotImplementedError("__dlpack_device__") + + def __repr__(self) -> str: + return ( + "PyArrowBuffer(" + + str( + { + "bufsize": self.bufsize, + "ptr": self.ptr, + "device": self.__dlpack_device__()[0].name, + } + ) + + ")" + ) diff --git a/lib/python3.12/site-packages/pyarrow/interchange/column.py b/lib/python3.12/site-packages/pyarrow/interchange/column.py new file mode 100644 index 0000000000000000000000000000000000000000..ddbceabcb00ade7a3d760f4e1663819571584d8e --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/interchange/column.py @@ -0,0 +1,529 @@ +# 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. + +from __future__ import annotations + +import enum +from typing import ( + Any, + Dict, + Iterable, + Optional, + Tuple, +) + +import sys +if sys.version_info >= (3, 8): + from typing import TypedDict +else: + from typing_extensions import TypedDict + +import pyarrow as pa +import pyarrow.compute as pc +from pyarrow.interchange.buffer import _PyArrowBuffer + + +class DtypeKind(enum.IntEnum): + """ + Integer enum for data types. + + Attributes + ---------- + INT : int + Matches to signed integer data type. + UINT : int + Matches to unsigned integer data type. + FLOAT : int + Matches to floating point data type. + BOOL : int + Matches to boolean data type. + STRING : int + Matches to string data type (UTF-8 encoded). + DATETIME : int + Matches to datetime data type. + CATEGORICAL : int + Matches to categorical data type. + """ + + INT = 0 + UINT = 1 + FLOAT = 2 + BOOL = 20 + STRING = 21 # UTF-8 + DATETIME = 22 + CATEGORICAL = 23 + + +Dtype = Tuple[DtypeKind, int, str, str] # see Column.dtype + + +_PYARROW_KINDS = { + pa.int8(): (DtypeKind.INT, "c"), + pa.int16(): (DtypeKind.INT, "s"), + pa.int32(): (DtypeKind.INT, "i"), + pa.int64(): (DtypeKind.INT, "l"), + pa.uint8(): (DtypeKind.UINT, "C"), + pa.uint16(): (DtypeKind.UINT, "S"), + pa.uint32(): (DtypeKind.UINT, "I"), + pa.uint64(): (DtypeKind.UINT, "L"), + pa.float16(): (DtypeKind.FLOAT, "e"), + pa.float32(): (DtypeKind.FLOAT, "f"), + pa.float64(): (DtypeKind.FLOAT, "g"), + pa.bool_(): (DtypeKind.BOOL, "b"), + pa.string(): (DtypeKind.STRING, "u"), + pa.large_string(): (DtypeKind.STRING, "U"), +} + + +class ColumnNullType(enum.IntEnum): + """ + Integer enum for null type representation. + + Attributes + ---------- + NON_NULLABLE : int + Non-nullable column. + USE_NAN : int + Use explicit float NaN value. + USE_SENTINEL : int + Sentinel value besides NaN. + USE_BITMASK : int + The bit is set/unset representing a null on a certain position. + USE_BYTEMASK : int + The byte is set/unset representing a null on a certain position. + """ + + NON_NULLABLE = 0 + USE_NAN = 1 + USE_SENTINEL = 2 + USE_BITMASK = 3 + USE_BYTEMASK = 4 + + +class ColumnBuffers(TypedDict): + # first element is a buffer containing the column data; + # second element is the data buffer's associated dtype + data: Tuple[_PyArrowBuffer, Dtype] + + # first element is a buffer containing mask values indicating missing data; + # second element is the mask value buffer's associated dtype. + # None if the null representation is not a bit or byte mask + validity: Optional[Tuple[_PyArrowBuffer, Dtype]] + + # first element is a buffer containing the offset values for + # variable-size binary data (e.g., variable-length strings); + # second element is the offsets buffer's associated dtype. + # None if the data buffer does not have an associated offsets buffer + offsets: Optional[Tuple[_PyArrowBuffer, Dtype]] + + +class CategoricalDescription(TypedDict): + # whether the ordering of dictionary indices is semantically meaningful + is_ordered: bool + # whether a dictionary-style mapping of categorical values to other objects + # exists + is_dictionary: bool + # Python-level only (e.g. ``{int: str}``). + # None if not a dictionary-style categorical. + categories: Optional[_PyArrowColumn] + + +class Endianness: + """Enum indicating the byte-order of a data-type.""" + + LITTLE = "<" + BIG = ">" + NATIVE = "=" + NA = "|" + + +class NoBufferPresent(Exception): + """Exception to signal that there is no requested buffer.""" + + +class _PyArrowColumn: + """ + A column object, with only the methods and properties required by the + interchange protocol defined. + + A column can contain one or more chunks. Each chunk can contain up to three + buffers - a data buffer, a mask buffer (depending on null representation), + and an offsets buffer (if variable-size binary; e.g., variable-length + strings). + + TBD: Arrow has a separate "null" dtype, and has no separate mask concept. + Instead, it seems to use "children" for both columns with a bit mask, + and for nested dtypes. Unclear whether this is elegant or confusing. + This design requires checking the null representation explicitly. + + The Arrow design requires checking: + 1. the ARROW_FLAG_NULLABLE (for sentinel values) + 2. if a column has two children, combined with one of those children + having a null dtype. + + Making the mask concept explicit seems useful. One null dtype would + not be enough to cover both bit and byte masks, so that would mean + even more checking if we did it the Arrow way. + + TBD: there's also the "chunk" concept here, which is implicit in Arrow as + multiple buffers per array (= column here). Semantically it may make + sense to have both: chunks were meant for example for lazy evaluation + of data which doesn't fit in memory, while multiple buffers per column + could also come from doing a selection operation on a single + contiguous buffer. + + Given these concepts, one would expect chunks to be all of the same + size (say a 10,000 row dataframe could have 10 chunks of 1,000 rows), + while multiple buffers could have data-dependent lengths. Not an issue + in pandas if one column is backed by a single NumPy array, but in + Arrow it seems possible. + Are multiple chunks *and* multiple buffers per column necessary for + the purposes of this interchange protocol, or must producers either + reuse the chunk concept for this or copy the data? + + Note: this Column object can only be produced by ``__dataframe__``, so + doesn't need its own version or ``__column__`` protocol. + """ + + def __init__( + self, column: pa.Array | pa.ChunkedArray, allow_copy: bool = True + ) -> None: + """ + Handles PyArrow Arrays and ChunkedArrays. + """ + # Store the column as a private attribute + if isinstance(column, pa.ChunkedArray): + if column.num_chunks == 1: + column = column.chunk(0) + else: + if not allow_copy: + raise RuntimeError( + "Chunks will be combined and a copy is required which " + "is forbidden by allow_copy=False" + ) + column = column.combine_chunks() + + self._allow_copy = allow_copy + + if pa.types.is_boolean(column.type): + if not allow_copy: + raise RuntimeError( + "Boolean column will be casted to uint8 and a copy " + "is required which is forbidden by allow_copy=False" + ) + self._dtype = self._dtype_from_arrowdtype(column.type, 8) + self._col = pc.cast(column, pa.uint8()) + else: + self._col = column + dtype = self._col.type + try: + bit_width = dtype.bit_width + except ValueError: + # in case of a variable-length strings, considered as array + # of bytes (8 bits) + bit_width = 8 + self._dtype = self._dtype_from_arrowdtype(dtype, bit_width) + + def size(self) -> int: + """ + Size of the column, in elements. + + Corresponds to DataFrame.num_rows() if column is a single chunk; + equal to size of this current chunk otherwise. + + Is a method rather than a property because it may cause a (potentially + expensive) computation for some dataframe implementations. + """ + return len(self._col) + + @property + def offset(self) -> int: + """ + Offset of first element. + + May be > 0 if using chunks; for example for a column with N chunks of + equal size M (only the last chunk may be shorter), + ``offset = n * M``, ``n = 0 .. N-1``. + """ + return self._col.offset + + @property + def dtype(self) -> Tuple[DtypeKind, int, str, str]: + """ + Dtype description as a tuple ``(kind, bit-width, format string, + endianness)``. + + Bit-width : the number of bits as an integer + Format string : data type description format string in Apache Arrow C + Data Interface format. + Endianness : current only native endianness (``=``) is supported + + Notes: + - Kind specifiers are aligned with DLPack where possible (hence the + jump to 20, leave enough room for future extension) + - Masks must be specified as boolean with either bit width 1 (for + bit masks) or 8 (for byte masks). + - Dtype width in bits was preferred over bytes + - Endianness isn't too useful, but included now in case in the + future we need to support non-native endianness + - Went with Apache Arrow format strings over NumPy format strings + because they're more complete from a dataframe perspective + - Format strings are mostly useful for datetime specification, and + for categoricals. + - For categoricals, the format string describes the type of the + categorical in the data buffer. In case of a separate encoding of + the categorical (e.g. an integer to string mapping), this can + be derived from ``self.describe_categorical``. + - Data types not included: complex, Arrow-style null, binary, + decimal, and nested (list, struct, map, union) dtypes. + """ + return self._dtype + + def _dtype_from_arrowdtype( + self, dtype: pa.DataType, bit_width: int + ) -> Tuple[DtypeKind, int, str, str]: + """ + See `self.dtype` for details. + """ + # Note: 'c' (complex) not handled yet (not in array spec v1). + # 'b', 'B' (bytes), 'S', 'a', (old-style string) 'V' (void) + # not handled datetime and timedelta both map to datetime + # (is timedelta handled?) + + if pa.types.is_timestamp(dtype): + kind = DtypeKind.DATETIME + ts = dtype.unit[0] + tz = dtype.tz if dtype.tz else "" + f_string = f"ts{ts}:{tz}" + return kind, bit_width, f_string, Endianness.NATIVE + elif pa.types.is_dictionary(dtype): + kind = DtypeKind.CATEGORICAL + arr = self._col + indices_dtype = arr.indices.type + _, f_string = _PYARROW_KINDS.get(indices_dtype) + return kind, bit_width, f_string, Endianness.NATIVE + else: + kind, f_string = _PYARROW_KINDS.get(dtype, (None, None)) + if kind is None: + raise ValueError( + f"Data type {dtype} not supported by interchange protocol") + + return kind, bit_width, f_string, Endianness.NATIVE + + @property + def describe_categorical(self) -> CategoricalDescription: + """ + If the dtype is categorical, there are two options: + - There are only values in the data buffer. + - There is a separate non-categorical Column encoding categorical + values. + + Raises TypeError if the dtype is not categorical + + Returns the dictionary with description on how to interpret the + data buffer: + - "is_ordered" : bool, whether the ordering of dictionary indices + is semantically meaningful. + - "is_dictionary" : bool, whether a mapping of + categorical values to other objects exists + - "categories" : Column representing the (implicit) mapping of + indices to category values (e.g. an array of + cat1, cat2, ...). None if not a dictionary-style + categorical. + + TBD: are there any other in-memory representations that are needed? + """ + arr = self._col + if not pa.types.is_dictionary(arr.type): + raise TypeError( + "describe_categorical only works on a column with " + "categorical dtype!" + ) + + return { + "is_ordered": self._col.type.ordered, + "is_dictionary": True, + "categories": _PyArrowColumn(arr.dictionary), + } + + @property + def describe_null(self) -> Tuple[ColumnNullType, Any]: + """ + Return the missing value (or "null") representation the column dtype + uses, as a tuple ``(kind, value)``. + + Value : if kind is "sentinel value", the actual value. If kind is a bit + mask or a byte mask, the value (0 or 1) indicating a missing value. + None otherwise. + """ + # In case of no missing values, we need to set ColumnNullType to + # non nullable as in the current __dataframe__ protocol bit/byte masks + # cannot be None + if self.null_count == 0: + return ColumnNullType.NON_NULLABLE, None + else: + return ColumnNullType.USE_BITMASK, 0 + + @property + def null_count(self) -> int: + """ + Number of null elements, if known. + + Note: Arrow uses -1 to indicate "unknown", but None seems cleaner. + """ + arrow_null_count = self._col.null_count + n = arrow_null_count if arrow_null_count != -1 else None + return n + + @property + def metadata(self) -> Dict[str, Any]: + """ + The metadata for the column. See `DataFrame.metadata` for more details. + """ + pass + + def num_chunks(self) -> int: + """ + Return the number of chunks the column consists of. + """ + return 1 + + def get_chunks( + self, n_chunks: Optional[int] = None + ) -> Iterable[_PyArrowColumn]: + """ + Return an iterator yielding the chunks. + + See `DataFrame.get_chunks` for details on ``n_chunks``. + """ + if n_chunks and n_chunks > 1: + chunk_size = self.size() // n_chunks + if self.size() % n_chunks != 0: + chunk_size += 1 + + array = self._col + i = 0 + for start in range(0, chunk_size * n_chunks, chunk_size): + yield _PyArrowColumn( + array.slice(start, chunk_size), self._allow_copy + ) + i += 1 + else: + yield self + + def get_buffers(self) -> ColumnBuffers: + """ + Return a dictionary containing the underlying buffers. + + The returned dictionary has the following contents: + + - "data": a two-element tuple whose first element is a buffer + containing the data and whose second element is the data + buffer's associated dtype. + - "validity": a two-element tuple whose first element is a buffer + containing mask values indicating missing data and + whose second element is the mask value buffer's + associated dtype. None if the null representation is + not a bit or byte mask. + - "offsets": a two-element tuple whose first element is a buffer + containing the offset values for variable-size binary + data (e.g., variable-length strings) and whose second + element is the offsets buffer's associated dtype. None + if the data buffer does not have an associated offsets + buffer. + """ + buffers: ColumnBuffers = { + "data": self._get_data_buffer(), + "validity": None, + "offsets": None, + } + + try: + buffers["validity"] = self._get_validity_buffer() + except NoBufferPresent: + pass + + try: + buffers["offsets"] = self._get_offsets_buffer() + except NoBufferPresent: + pass + + return buffers + + def _get_data_buffer( + self, + ) -> Tuple[_PyArrowBuffer, Any]: # Any is for self.dtype tuple + """ + Return the buffer containing the data and the buffer's + associated dtype. + """ + array = self._col + dtype = self.dtype + + # In case of dictionary arrays, use indices + # to define a buffer, codes are transferred through + # describe_categorical() + if pa.types.is_dictionary(array.type): + array = array.indices + dtype = _PyArrowColumn(array).dtype + + n = len(array.buffers()) + if n == 2: + return _PyArrowBuffer(array.buffers()[1]), dtype + elif n == 3: + return _PyArrowBuffer(array.buffers()[2]), dtype + + def _get_validity_buffer(self) -> Tuple[_PyArrowBuffer, Any]: + """ + Return the buffer containing the mask values indicating missing data + and the buffer's associated dtype. + Raises NoBufferPresent if null representation is not a bit or byte + mask. + """ + # Define the dtype of the returned buffer + dtype = (DtypeKind.BOOL, 1, "b", Endianness.NATIVE) + array = self._col + buff = array.buffers()[0] + if buff: + return _PyArrowBuffer(buff), dtype + else: + raise NoBufferPresent( + "There are no missing values so " + "does not have a separate mask") + + def _get_offsets_buffer(self) -> Tuple[_PyArrowBuffer, Any]: + """ + Return the buffer containing the offset values for variable-size binary + data (e.g., variable-length strings) and the buffer's associated dtype. + Raises NoBufferPresent if the data buffer does not have an associated + offsets buffer. + """ + array = self._col + n = len(array.buffers()) + if n == 2: + raise NoBufferPresent( + "This column has a fixed-length dtype so " + "it does not have an offsets buffer" + ) + elif n == 3: + # Define the dtype of the returned buffer + dtype = self._col.type + if pa.types.is_large_string(dtype): + dtype = (DtypeKind.INT, 64, "l", Endianness.NATIVE) + else: + dtype = (DtypeKind.INT, 32, "i", Endianness.NATIVE) + return _PyArrowBuffer(array.buffers()[1]), dtype diff --git a/lib/python3.12/site-packages/pyarrow/interchange/dataframe.py b/lib/python3.12/site-packages/pyarrow/interchange/dataframe.py new file mode 100644 index 0000000000000000000000000000000000000000..59ba765c175ad471274a99bf857c8880a072e0b8 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/interchange/dataframe.py @@ -0,0 +1,217 @@ +# 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. + +from __future__ import annotations +from typing import ( + Any, + Iterable, + Optional, + Sequence, +) + +import pyarrow as pa + +from pyarrow.interchange.column import _PyArrowColumn + + +class _PyArrowDataFrame: + """ + A data frame class, with only the methods required by the interchange + protocol defined. + + A "data frame" represents an ordered collection of named columns. + A column's "name" must be a unique string. + Columns may be accessed by name or by position. + + This could be a public data frame class, or an object with the methods and + attributes defined on this DataFrame class could be returned from the + ``__dataframe__`` method of a public data frame class in a library adhering + to the dataframe interchange protocol specification. + """ + + def __init__( + self, df: pa.Table | pa.RecordBatch, + nan_as_null: bool = False, + allow_copy: bool = True + ) -> None: + """ + Constructor - an instance of this (private) class is returned from + `pa.Table.__dataframe__` or `pa.RecordBatch.__dataframe__`. + """ + self._df = df + # ``nan_as_null`` is a keyword intended for the consumer to tell the + # producer to overwrite null values in the data with ``NaN`` (or + # ``NaT``). + if nan_as_null is True: + raise RuntimeError( + "nan_as_null=True currently has no effect, " + "use the default nan_as_null=False" + ) + self._nan_as_null = nan_as_null + self._allow_copy = allow_copy + + def __dataframe__( + self, nan_as_null: bool = False, allow_copy: bool = True + ) -> _PyArrowDataFrame: + """ + Construct a new exchange object, potentially changing the parameters. + ``nan_as_null`` is a keyword intended for the consumer to tell the + producer to overwrite null values in the data with ``NaN``. + It is intended for cases where the consumer does not support the bit + mask or byte mask that is the producer's native representation. + ``allow_copy`` is a keyword that defines whether or not the library is + allowed to make a copy of the data. For example, copying data would be + necessary if a library supports strided buffers, given that this + protocol specifies contiguous buffers. + """ + return _PyArrowDataFrame(self._df, nan_as_null, allow_copy) + + @property + def metadata(self) -> dict[str, Any]: + """ + The metadata for the data frame, as a dictionary with string keys. The + contents of `metadata` may be anything, they are meant for a library + to store information that it needs to, e.g., roundtrip losslessly or + for two implementations to share data that is not (yet) part of the + interchange protocol specification. For avoiding collisions with other + entries, please add name the keys with the name of the library + followed by a period and the desired name, e.g, ``pandas.indexcol``. + """ + # The metadata for the data frame, as a dictionary with string keys. + # Add schema metadata here (pandas metadata or custom metadata) + if self._df.schema.metadata: + schema_metadata = {"pyarrow." + k.decode('utf8'): v.decode('utf8') + for k, v in self._df.schema.metadata.items()} + return schema_metadata + else: + return {} + + def num_columns(self) -> int: + """ + Return the number of columns in the DataFrame. + """ + return self._df.num_columns + + def num_rows(self) -> int: + """ + Return the number of rows in the DataFrame, if available. + """ + return self._df.num_rows + + def num_chunks(self) -> int: + """ + Return the number of chunks the DataFrame consists of. + """ + if isinstance(self._df, pa.RecordBatch): + return 1 + else: + # pyarrow.Table can have columns with different number + # of chunks so we take the number of chunks that + # .to_batches() returns as it takes the min chunk size + # of all the columns (to_batches is a zero copy method) + batches = self._df.to_batches() + return len(batches) + + def column_names(self) -> Iterable[str]: + """ + Return an iterator yielding the column names. + """ + return self._df.schema.names + + def get_column(self, i: int) -> _PyArrowColumn: + """ + Return the column at the indicated position. + """ + return _PyArrowColumn(self._df.column(i), + allow_copy=self._allow_copy) + + def get_column_by_name(self, name: str) -> _PyArrowColumn: + """ + Return the column whose name is the indicated name. + """ + return _PyArrowColumn(self._df.column(name), + allow_copy=self._allow_copy) + + def get_columns(self) -> Iterable[_PyArrowColumn]: + """ + Return an iterator yielding the columns. + """ + return [ + _PyArrowColumn(col, allow_copy=self._allow_copy) + for col in self._df.columns + ] + + def select_columns(self, indices: Sequence[int]) -> _PyArrowDataFrame: + """ + Create a new DataFrame by selecting a subset of columns by index. + """ + return _PyArrowDataFrame( + self._df.select(list(indices)), self._nan_as_null, self._allow_copy + ) + + def select_columns_by_name( + self, names: Sequence[str] + ) -> _PyArrowDataFrame: + """ + Create a new DataFrame by selecting a subset of columns by name. + """ + return _PyArrowDataFrame( + self._df.select(list(names)), self._nan_as_null, self._allow_copy + ) + + def get_chunks( + self, n_chunks: Optional[int] = None + ) -> Iterable[_PyArrowDataFrame]: + """ + Return an iterator yielding the chunks. + + By default (None), yields the chunks that the data is stored as by the + producer. If given, ``n_chunks`` must be a multiple of + ``self.num_chunks()``, meaning the producer must subdivide each chunk + before yielding it. + + Note that the producer must ensure that all columns are chunked the + same way. + """ + # Subdivide chunks + if n_chunks and n_chunks > 1: + chunk_size = self.num_rows() // n_chunks + if self.num_rows() % n_chunks != 0: + chunk_size += 1 + if isinstance(self._df, pa.Table): + batches = self._df.to_batches(max_chunksize=chunk_size) + else: + batches = [] + for start in range(0, chunk_size * n_chunks, chunk_size): + batches.append(self._df.slice(start, chunk_size)) + # In case when the size of the chunk is such that the resulting + # list is one less chunk then n_chunks -> append an empty chunk + if len(batches) == n_chunks - 1: + batches.append(pa.record_batch([[]], schema=self._df.schema)) + # yields the chunks that the data is stored as + else: + if isinstance(self._df, pa.Table): + batches = self._df.to_batches() + else: + batches = [self._df] + + # Create an iterator of RecordBatches + iterator = [_PyArrowDataFrame(batch, + self._nan_as_null, + self._allow_copy) + for batch in batches] + return iterator diff --git a/lib/python3.12/site-packages/pyarrow/interchange/from_dataframe.py b/lib/python3.12/site-packages/pyarrow/interchange/from_dataframe.py new file mode 100644 index 0000000000000000000000000000000000000000..fcaec41e3dcdf982e19bd45ba4a1941fab5ec34e --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/interchange/from_dataframe.py @@ -0,0 +1,614 @@ +# 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. + +from __future__ import annotations + +from typing import ( + Any, + Tuple, +) + +from pyarrow.interchange.column import ( + DtypeKind, + ColumnBuffers, + ColumnNullType, +) + +import pyarrow as pa +import re + +import pyarrow.compute as pc +from pyarrow.interchange.column import Dtype + + +# A typing protocol could be added later to let Mypy validate code using +# `from_dataframe` better. +DataFrameObject = Any +ColumnObject = Any +BufferObject = Any + + +_PYARROW_DTYPES: dict[DtypeKind, dict[int, Any]] = { + DtypeKind.INT: {8: pa.int8(), + 16: pa.int16(), + 32: pa.int32(), + 64: pa.int64()}, + DtypeKind.UINT: {8: pa.uint8(), + 16: pa.uint16(), + 32: pa.uint32(), + 64: pa.uint64()}, + DtypeKind.FLOAT: {16: pa.float16(), + 32: pa.float32(), + 64: pa.float64()}, + DtypeKind.BOOL: {1: pa.bool_(), + 8: pa.uint8()}, + DtypeKind.STRING: {8: pa.string()}, +} + + +def from_dataframe(df: DataFrameObject, allow_copy=True) -> pa.Table: + """ + Build a ``pa.Table`` from any DataFrame supporting the interchange protocol. + + Parameters + ---------- + df : DataFrameObject + Object supporting the interchange protocol, i.e. `__dataframe__` + method. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Table + + Examples + -------- + >>> import pyarrow + >>> from pyarrow.interchange import from_dataframe + + Convert a pandas dataframe to a pyarrow table: + + >>> import pandas as pd + >>> df = pd.DataFrame({ + ... "n_attendees": [100, 10, 1], + ... "country": ["Italy", "Spain", "Slovenia"], + ... }) + >>> df + n_attendees country + 0 100 Italy + 1 10 Spain + 2 1 Slovenia + >>> from_dataframe(df) + pyarrow.Table + n_attendees: int64 + country: large_string + ---- + n_attendees: [[100,10,1]] + country: [["Italy","Spain","Slovenia"]] + """ + if isinstance(df, pa.Table): + return df + elif isinstance(df, pa.RecordBatch): + return pa.Table.from_batches([df]) + + if not hasattr(df, "__dataframe__"): + raise ValueError("`df` does not support __dataframe__") + + return _from_dataframe(df.__dataframe__(allow_copy=allow_copy), + allow_copy=allow_copy) + + +def _from_dataframe(df: DataFrameObject, allow_copy=True): + """ + Build a ``pa.Table`` from the DataFrame interchange object. + + Parameters + ---------- + df : DataFrameObject + Object supporting the interchange protocol, i.e. `__dataframe__` + method. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Table + """ + batches = [] + for chunk in df.get_chunks(): + batch = protocol_df_chunk_to_pyarrow(chunk, allow_copy) + batches.append(batch) + + if not batches: + batch = protocol_df_chunk_to_pyarrow(df) + batches.append(batch) + + return pa.Table.from_batches(batches) + + +def protocol_df_chunk_to_pyarrow( + df: DataFrameObject, + allow_copy: bool = True +) -> pa.RecordBatch: + """ + Convert interchange protocol chunk to ``pa.RecordBatch``. + + Parameters + ---------- + df : DataFrameObject + Object supporting the interchange protocol, i.e. `__dataframe__` + method. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.RecordBatch + """ + # We need a dict of columns here, with each column being a pa.Array + columns: dict[str, pa.Array] = {} + for name in df.column_names(): + if not isinstance(name, str): + raise ValueError(f"Column {name} is not a string") + if name in columns: + raise ValueError(f"Column {name} is not unique") + col = df.get_column_by_name(name) + dtype = col.dtype[0] + if dtype in ( + DtypeKind.INT, + DtypeKind.UINT, + DtypeKind.FLOAT, + DtypeKind.STRING, + DtypeKind.DATETIME, + ): + columns[name] = column_to_array(col, allow_copy) + elif dtype == DtypeKind.BOOL: + columns[name] = bool_column_to_array(col, allow_copy) + elif dtype == DtypeKind.CATEGORICAL: + columns[name] = categorical_column_to_dictionary(col, allow_copy) + else: + raise NotImplementedError(f"Data type {dtype} not handled yet") + + return pa.RecordBatch.from_pydict(columns) + + +def column_to_array( + col: ColumnObject, + allow_copy: bool = True, +) -> pa.Array: + """ + Convert a column holding one of the primitive dtypes to a PyArrow array. + A primitive type is one of: int, uint, float, bool (1 bit). + + Parameters + ---------- + col : ColumnObject + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Array + """ + buffers = col.get_buffers() + data_type = col.dtype + data = buffers_to_array(buffers, data_type, + col.size(), + col.describe_null, + col.offset, + allow_copy) + return data + + +def bool_column_to_array( + col: ColumnObject, + allow_copy: bool = True, +) -> pa.Array: + """ + Convert a column holding boolean dtype to a PyArrow array. + + Parameters + ---------- + col : ColumnObject + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Array + """ + buffers = col.get_buffers() + size = buffers["data"][1][1] + + # If booleans are byte-packed a copy to bit-packed will be made + if size == 8 and not allow_copy: + raise RuntimeError( + "Boolean column will be casted from uint8 and a copy " + "is required which is forbidden by allow_copy=False" + ) + + data_type = col.dtype + data = buffers_to_array(buffers, data_type, + col.size(), + col.describe_null, + col.offset) + if size == 8: + data = pc.cast(data, pa.bool_()) + + return data + + +def categorical_column_to_dictionary( + col: ColumnObject, + allow_copy: bool = True, +) -> pa.DictionaryArray: + """ + Convert a column holding categorical data to a pa.DictionaryArray. + + Parameters + ---------- + col : ColumnObject + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.DictionaryArray + """ + if not allow_copy: + raise RuntimeError( + "Categorical column will be casted from uint8 and a copy " + "is required which is forbidden by allow_copy=False" + ) + + categorical = col.describe_categorical + + if not categorical["is_dictionary"]: + raise NotImplementedError( + "Non-dictionary categoricals not supported yet") + + # We need to first convert the dictionary column + cat_column = categorical["categories"] + dictionary = column_to_array(cat_column) + # Then we need to convert the indices + # Here we need to use the buffer data type! + buffers = col.get_buffers() + _, data_type = buffers["data"] + indices = buffers_to_array(buffers, data_type, + col.size(), + col.describe_null, + col.offset) + + # Constructing a pa.DictionaryArray + dict_array = pa.DictionaryArray.from_arrays(indices, dictionary) + + return dict_array + + +def parse_datetime_format_str(format_str): + """Parse datetime `format_str` to interpret the `data`.""" + + # timestamp 'ts{unit}:tz' + timestamp_meta = re.match(r"ts([smun]):(.*)", format_str) + if timestamp_meta: + unit, tz = timestamp_meta.group(1), timestamp_meta.group(2) + if unit != "s": + # the format string describes only a first letter of the unit, so + # add one extra letter to convert the unit to numpy-style: + # 'm' -> 'ms', 'u' -> 'us', 'n' -> 'ns' + unit += "s" + + return unit, tz + + raise NotImplementedError(f"DateTime kind is not supported: {format_str}") + + +def map_date_type(data_type): + """Map column date type to pyarrow date type. """ + kind, bit_width, f_string, _ = data_type + + if kind == DtypeKind.DATETIME: + unit, tz = parse_datetime_format_str(f_string) + return pa.timestamp(unit, tz=tz) + else: + pa_dtype = _PYARROW_DTYPES.get(kind, {}).get(bit_width, None) + + # Error if dtype is not supported + if pa_dtype: + return pa_dtype + else: + raise NotImplementedError( + f"Conversion for {data_type} is not yet supported.") + + +def buffers_to_array( + buffers: ColumnBuffers, + data_type: Tuple[DtypeKind, int, str, str], + length: int, + describe_null: ColumnNullType, + offset: int = 0, + allow_copy: bool = True, +) -> pa.Array: + """ + Build a PyArrow array from the passed buffer. + + Parameters + ---------- + buffer : ColumnBuffers + Dictionary containing tuples of underlying buffers and + their associated dtype. + data_type : Tuple[DtypeKind, int, str, str], + Dtype description of the column as a tuple ``(kind, bit-width, format string, + endianness)``. + length : int + The number of values in the array. + describe_null: ColumnNullType + Null representation the column dtype uses, + as a tuple ``(kind, value)`` + offset : int, default: 0 + Number of elements to offset from the start of the buffer. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Array + + Notes + ----- + The returned array doesn't own the memory. The caller of this function + is responsible for keeping the memory owner object alive as long as + the returned PyArrow array is being used. + """ + data_buff, _ = buffers["data"] + try: + validity_buff, validity_dtype = buffers["validity"] + except TypeError: + validity_buff = None + try: + offset_buff, offset_dtype = buffers["offsets"] + except TypeError: + offset_buff = None + + # Construct a pyarrow Buffer + data_pa_buffer = pa.foreign_buffer(data_buff.ptr, data_buff.bufsize, + base=data_buff) + + # Construct a validity pyarrow Buffer, if applicable + if validity_buff: + validity_pa_buff = validity_buffer_from_mask(validity_buff, + validity_dtype, + describe_null, + length, + offset, + allow_copy) + else: + validity_pa_buff = validity_buffer_nan_sentinel(data_pa_buffer, + data_type, + describe_null, + length, + offset, + allow_copy) + + # Construct a pyarrow Array from buffers + data_dtype = map_date_type(data_type) + + if offset_buff: + _, offset_bit_width, _, _ = offset_dtype + # If an offset buffer exists, construct an offset pyarrow Buffer + # and add it to the construction of an array + offset_pa_buffer = pa.foreign_buffer(offset_buff.ptr, + offset_buff.bufsize, + base=offset_buff) + + if data_type[2] == 'U': + string_type = pa.large_string() + else: + if offset_bit_width == 64: + string_type = pa.large_string() + else: + string_type = pa.string() + array = pa.Array.from_buffers( + string_type, + length, + [validity_pa_buff, offset_pa_buffer, data_pa_buffer], + offset=offset, + ) + else: + array = pa.Array.from_buffers( + data_dtype, + length, + [validity_pa_buff, data_pa_buffer], + offset=offset, + ) + + return array + + +def validity_buffer_from_mask( + validity_buff: BufferObject, + validity_dtype: Dtype, + describe_null: ColumnNullType, + length: int, + offset: int = 0, + allow_copy: bool = True, +) -> pa.Buffer: + """ + Build a PyArrow buffer from the passed mask buffer. + + Parameters + ---------- + validity_buff : BufferObject + Tuple of underlying validity buffer and associated dtype. + validity_dtype : Dtype + Dtype description as a tuple ``(kind, bit-width, format string, + endianness)``. + describe_null : ColumnNullType + Null representation the column dtype uses, + as a tuple ``(kind, value)`` + length : int + The number of values in the array. + offset : int, default: 0 + Number of elements to offset from the start of the buffer. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Buffer + """ + null_kind, sentinel_val = describe_null + validity_kind, _, _, _ = validity_dtype + assert validity_kind == DtypeKind.BOOL + + if null_kind == ColumnNullType.NON_NULLABLE: + # Sliced array can have a NON_NULLABLE ColumnNullType due + # to no missing values in that slice of an array though the bitmask + # exists and validity_buff must be set to None in this case + return None + + elif null_kind == ColumnNullType.USE_BYTEMASK or ( + null_kind == ColumnNullType.USE_BITMASK and sentinel_val == 1 + ): + buff = pa.foreign_buffer(validity_buff.ptr, + validity_buff.bufsize, + base=validity_buff) + + if null_kind == ColumnNullType.USE_BYTEMASK: + if not allow_copy: + raise RuntimeError( + "To create a bitmask a copy of the data is " + "required which is forbidden by allow_copy=False" + ) + mask = pa.Array.from_buffers(pa.int8(), length, + [None, buff], + offset=offset) + mask_bool = pc.cast(mask, pa.bool_()) + else: + mask_bool = pa.Array.from_buffers(pa.bool_(), length, + [None, buff], + offset=offset) + + if sentinel_val == 1: + mask_bool = pc.invert(mask_bool) + + return mask_bool.buffers()[1] + + elif null_kind == ColumnNullType.USE_BITMASK and sentinel_val == 0: + return pa.foreign_buffer(validity_buff.ptr, + validity_buff.bufsize, + base=validity_buff) + else: + raise NotImplementedError( + f"{describe_null} null representation is not yet supported.") + + +def validity_buffer_nan_sentinel( + data_pa_buffer: BufferObject, + data_type: Dtype, + describe_null: ColumnNullType, + length: int, + offset: int = 0, + allow_copy: bool = True, +) -> pa.Buffer: + """ + Build a PyArrow buffer from NaN or sentinel values. + + Parameters + ---------- + data_pa_buffer : pa.Buffer + PyArrow buffer for the column data. + data_type : Dtype + Dtype description as a tuple ``(kind, bit-width, format string, + endianness)``. + describe_null : ColumnNullType + Null representation the column dtype uses, + as a tuple ``(kind, value)`` + length : int + The number of values in the array. + offset : int, default: 0 + Number of elements to offset from the start of the buffer. + allow_copy : bool, default: True + Whether to allow copying the memory to perform the conversion + (if false then zero-copy approach is requested). + + Returns + ------- + pa.Buffer + """ + kind, bit_width, _, _ = data_type + data_dtype = map_date_type(data_type) + null_kind, sentinel_val = describe_null + + # Check for float NaN values + if null_kind == ColumnNullType.USE_NAN: + if not allow_copy: + raise RuntimeError( + "To create a bitmask a copy of the data is " + "required which is forbidden by allow_copy=False" + ) + + if kind == DtypeKind.FLOAT and bit_width == 16: + # 'pyarrow.compute.is_nan' kernel not yet implemented + # for float16 + raise NotImplementedError( + f"{data_type} with {null_kind} is not yet supported.") + else: + pyarrow_data = pa.Array.from_buffers( + data_dtype, + length, + [None, data_pa_buffer], + offset=offset, + ) + mask = pc.is_nan(pyarrow_data) + mask = pc.invert(mask) + return mask.buffers()[1] + + # Check for sentinel values + elif null_kind == ColumnNullType.USE_SENTINEL: + if not allow_copy: + raise RuntimeError( + "To create a bitmask a copy of the data is " + "required which is forbidden by allow_copy=False" + ) + + if kind == DtypeKind.DATETIME: + sentinel_dtype = pa.int64() + else: + sentinel_dtype = data_dtype + pyarrow_data = pa.Array.from_buffers(sentinel_dtype, + length, + [None, data_pa_buffer], + offset=offset) + sentinel_arr = pc.equal(pyarrow_data, sentinel_val) + mask_bool = pc.invert(sentinel_arr) + return mask_bool.buffers()[1] + + elif null_kind == ColumnNullType.NON_NULLABLE: + pass + else: + raise NotImplementedError( + f"{describe_null} null representation is not yet supported.") diff --git a/lib/python3.12/site-packages/pyarrow/parquet/__init__.py b/lib/python3.12/site-packages/pyarrow/parquet/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..134f3c097ef004f83fdc8e24e5cb45166c17577e --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/parquet/__init__.py @@ -0,0 +1,20 @@ +# 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. + +# flake8: noqa + +from .core import * diff --git a/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..673992324f47f143fe4e34c4681e02e50ca40b3a Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..130ebe4f2c1e0ec08ff06bf6115c005f0cb2508e Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/encryption.cpython-312.pyc b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/encryption.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffaba67f7c8f3e852f90aa797b6dd8c0ad05a481 Binary files /dev/null and b/lib/python3.12/site-packages/pyarrow/parquet/__pycache__/encryption.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pyarrow/parquet/core.py b/lib/python3.12/site-packages/pyarrow/parquet/core.py new file mode 100644 index 0000000000000000000000000000000000000000..676bc445238e1bead69620272f2405673a48e854 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/parquet/core.py @@ -0,0 +1,2454 @@ +# 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. + + +from collections import defaultdict +from contextlib import nullcontext +from functools import reduce + +import inspect +import json +import os +import re +import operator + +import pyarrow as pa + +try: + import pyarrow._parquet as _parquet +except ImportError as exc: + raise ImportError( + "The pyarrow installation is not built with support " + f"for the Parquet file format ({str(exc)})" + ) from None + +from pyarrow._parquet import (ParquetReader, Statistics, # noqa + FileMetaData, RowGroupMetaData, + ColumnChunkMetaData, + ParquetSchema, ColumnSchema, + ParquetLogicalType, + FileEncryptionProperties, + FileDecryptionProperties, + SortingColumn) +from pyarrow.fs import (LocalFileSystem, FileType, _resolve_filesystem_and_path, + _ensure_filesystem) +from pyarrow.util import guid, _is_path_like, _stringify_path, _deprecate_api + + +def _check_contains_null(val): + if isinstance(val, bytes): + for byte in val: + if isinstance(byte, bytes): + compare_to = chr(0) + else: + compare_to = 0 + if byte == compare_to: + return True + elif isinstance(val, str): + return '\x00' in val + return False + + +def _check_filters(filters, check_null_strings=True): + """ + Check if filters are well-formed. + """ + if filters is not None: + if len(filters) == 0 or any(len(f) == 0 for f in filters): + raise ValueError("Malformed filters") + if isinstance(filters[0][0], str): + # We have encountered the situation where we have one nesting level + # too few: + # We have [(,,), ..] instead of [[(,,), ..]] + filters = [filters] + if check_null_strings: + for conjunction in filters: + for col, op, val in conjunction: + if ( + isinstance(val, list) and + all(_check_contains_null(v) for v in val) or + _check_contains_null(val) + ): + raise NotImplementedError( + "Null-terminated binary strings are not supported " + "as filter values." + ) + return filters + + +_DNF_filter_doc = """Predicates are expressed using an ``Expression`` or using + the disjunctive normal form (DNF), like ``[[('x', '=', 0), ...], ...]``. + DNF allows arbitrary boolean logical combinations of single column predicates. + The innermost tuples each describe a single column predicate. The list of inner + predicates is interpreted as a conjunction (AND), forming a more selective and + multiple column predicate. Finally, the most outer list combines these filters + as a disjunction (OR). + + Predicates may also be passed as List[Tuple]. This form is interpreted + as a single conjunction. To express OR in predicates, one must + use the (preferred) List[List[Tuple]] notation. + + Each tuple has format: (``key``, ``op``, ``value``) and compares the + ``key`` with the ``value``. + The supported ``op`` are: ``=`` or ``==``, ``!=``, ``<``, ``>``, ``<=``, + ``>=``, ``in`` and ``not in``. If the ``op`` is ``in`` or ``not in``, the + ``value`` must be a collection such as a ``list``, a ``set`` or a + ``tuple``. + + Examples: + + Using the ``Expression`` API: + + .. code-block:: python + + import pyarrow.compute as pc + pc.field('x') = 0 + pc.field('y').isin(['a', 'b', 'c']) + ~pc.field('y').isin({'a', 'b'}) + + Using the DNF format: + + .. code-block:: python + + ('x', '=', 0) + ('y', 'in', ['a', 'b', 'c']) + ('z', 'not in', {'a','b'}) + + """ + + +def filters_to_expression(filters): + """ + Check if filters are well-formed and convert to an ``Expression``. + + Parameters + ---------- + filters : List[Tuple] or List[List[Tuple]] + + Notes + ----- + See internal ``pyarrow._DNF_filter_doc`` attribute for more details. + + Examples + -------- + + >>> filters_to_expression([('foo', '==', 'bar')]) + + + Returns + ------- + pyarrow.compute.Expression + An Expression representing the filters + """ + import pyarrow.dataset as ds + + if isinstance(filters, ds.Expression): + return filters + + filters = _check_filters(filters, check_null_strings=False) + + def convert_single_predicate(col, op, val): + field = ds.field(col) + + if op == "=" or op == "==": + return field == val + elif op == "!=": + return field != val + elif op == '<': + return field < val + elif op == '>': + return field > val + elif op == '<=': + return field <= val + elif op == '>=': + return field >= val + elif op == 'in': + return field.isin(val) + elif op == 'not in': + return ~field.isin(val) + else: + raise ValueError(f'"{col}" is not a valid operator in predicates.') + + disjunction_members = [] + + for conjunction in filters: + conjunction_members = [ + convert_single_predicate(col, op, val) + for col, op, val in conjunction + ] + + disjunction_members.append(reduce(operator.and_, conjunction_members)) + + return reduce(operator.or_, disjunction_members) + + +_filters_to_expression = _deprecate_api( + "_filters_to_expression", "filters_to_expression", + filters_to_expression, "10.0.0", DeprecationWarning) + + +# ---------------------------------------------------------------------- +# Reading a single Parquet file + + +class ParquetFile: + """ + Reader interface for a single Parquet file. + + Parameters + ---------- + source : str, pathlib.Path, pyarrow.NativeFile, or file-like object + Readable source. For passing bytes or buffer-like file containing a + Parquet file, use pyarrow.BufferReader. + metadata : FileMetaData, default None + Use existing metadata object, rather than reading from file. + common_metadata : FileMetaData, default None + Will be used in reads for pandas schema metadata if not found in the + main file's metadata, no other uses at the moment. + read_dictionary : list + List of column names to read directly as DictionaryArray. + binary_type : pyarrow.DataType, default None + If given, Parquet binary columns will be read as this datatype. + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. + list_type : subclass of pyarrow.DataType, default None + If given, non-MAP repeated columns will be read as an instance of + this datatype (either pyarrow.ListType or pyarrow.LargeListType). + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. + memory_map : bool, default False + If the source is a file path, use a memory map to read file, which can + improve performance in some environments. + buffer_size : int, default 0 + If positive, perform read buffering when deserializing individual + column chunks. Otherwise IO calls are unbuffered. + pre_buffer : bool, default False + Coalesce and issue file reads in parallel to improve performance on + high-latency filesystems (e.g. S3). If True, Arrow will use a + background I/O thread pool. + coerce_int96_timestamp_unit : str, default None + Cast timestamps that are stored in INT96 format to a particular + resolution (e.g. 'ms'). Setting to None is equivalent to 'ns' + and therefore INT96 timestamps will be inferred as timestamps + in nanoseconds. + decryption_properties : FileDecryptionProperties, default None + File decryption properties for Parquet Modular Encryption. + thrift_string_size_limit : int, default None + If not None, override the maximum total string size allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. + thrift_container_size_limit : int, default None + If not None, override the maximum total size of containers allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. + filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. + page_checksum_verification : bool, default False + If True, verify the checksum for each page read from the file. + arrow_extensions_enabled : bool, default True + If True, read Parquet logical types as Arrow extension types where possible, + (e.g., read JSON as the canonical `arrow.json` extension type or UUID as + the canonical `arrow.uuid` extension type). + + Examples + -------- + + Generate an example PyArrow Table and write it to Parquet file: + + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + + Create a ``ParquetFile`` object from the Parquet file: + + >>> parquet_file = pq.ParquetFile('example.parquet') + + Read the data: + + >>> parquet_file.read() + pyarrow.Table + n_legs: int64 + animal: string + ---- + n_legs: [[2,2,4,4,5,100]] + animal: [["Flamingo","Parrot","Dog","Horse","Brittle stars","Centipede"]] + + Create a ParquetFile object with "animal" column as DictionaryArray: + + >>> parquet_file = pq.ParquetFile('example.parquet', + ... read_dictionary=["animal"]) + >>> parquet_file.read() + pyarrow.Table + n_legs: int64 + animal: dictionary + ---- + n_legs: [[2,2,4,4,5,100]] + animal: [ -- dictionary: + ["Flamingo","Parrot",...,"Brittle stars","Centipede"] -- indices: + [0,1,2,3,4,5]] + """ + + def __init__(self, source, *, metadata=None, common_metadata=None, + read_dictionary=None, binary_type=None, list_type=None, + memory_map=False, buffer_size=0, pre_buffer=False, + coerce_int96_timestamp_unit=None, + decryption_properties=None, thrift_string_size_limit=None, + thrift_container_size_limit=None, filesystem=None, + page_checksum_verification=False, arrow_extensions_enabled=True): + + self._close_source = getattr(source, 'closed', True) + + filesystem, source = _resolve_filesystem_and_path( + source, filesystem, memory_map=memory_map) + if filesystem is not None: + source = filesystem.open_input_file(source) + self._close_source = True # We opened it here, ensure we close it. + + self.reader = ParquetReader() + self.reader.open( + source, use_memory_map=memory_map, + buffer_size=buffer_size, pre_buffer=pre_buffer, + read_dictionary=read_dictionary, metadata=metadata, + binary_type=binary_type, list_type=list_type, + coerce_int96_timestamp_unit=coerce_int96_timestamp_unit, + decryption_properties=decryption_properties, + thrift_string_size_limit=thrift_string_size_limit, + thrift_container_size_limit=thrift_container_size_limit, + page_checksum_verification=page_checksum_verification, + arrow_extensions_enabled=arrow_extensions_enabled, + ) + self.common_metadata = common_metadata + self._nested_paths_by_prefix = self._build_nested_paths() + + def __enter__(self): + return self + + def __exit__(self, *args, **kwargs): + self.close() + + def _build_nested_paths(self): + paths = self.reader.column_paths + + result = defaultdict(list) + + for i, path in enumerate(paths): + key = path[0] + rest = path[1:] + while True: + result[key].append(i) + + if not rest: + break + + key = '.'.join((key, rest[0])) + rest = rest[1:] + + return result + + @property + def metadata(self): + """ + Return the Parquet metadata. + """ + return self.reader.metadata + + @property + def schema(self): + """ + Return the Parquet schema, unconverted to Arrow types + """ + return self.metadata.schema + + @property + def schema_arrow(self): + """ + Return the inferred Arrow schema, converted from the whole Parquet + file's schema + + Examples + -------- + Generate an example Parquet file: + + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + Read the Arrow schema: + + >>> parquet_file.schema_arrow + n_legs: int64 + animal: string + """ + return self.reader.schema_arrow + + @property + def num_row_groups(self): + """ + Return the number of row groups of the Parquet file. + + Examples + -------- + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + >>> parquet_file.num_row_groups + 1 + """ + return self.reader.num_row_groups + + def close(self, force: bool = False): + if self._close_source or force: + self.reader.close() + + @property + def closed(self) -> bool: + return self.reader.closed + + def read_row_group(self, i, columns=None, use_threads=True, + use_pandas_metadata=False): + """ + Read a single row group from a Parquet file. + + Parameters + ---------- + i : int + Index of the individual row group that we want to read. + columns : list + If not None, only these columns will be read from the row group. A + column name may be a prefix of a nested field, e.g. 'a' will select + 'a.b', 'a.c', and 'a.d.e'. + use_threads : bool, default True + Perform multi-threaded column reads. + use_pandas_metadata : bool, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded. + + Returns + ------- + pyarrow.table.Table + Content of the row group as a table (of columns) + + Examples + -------- + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + >>> parquet_file.read_row_group(0) + pyarrow.Table + n_legs: int64 + animal: string + ---- + n_legs: [[2,2,4,4,5,100]] + animal: [["Flamingo","Parrot",...,"Brittle stars","Centipede"]] + """ + column_indices = self._get_column_indices( + columns, use_pandas_metadata=use_pandas_metadata) + return self.reader.read_row_group(i, column_indices=column_indices, + use_threads=use_threads) + + def read_row_groups(self, row_groups, columns=None, use_threads=True, + use_pandas_metadata=False): + """ + Read a multiple row groups from a Parquet file. + + Parameters + ---------- + row_groups : list + Only these row groups will be read from the file. + columns : list + If not None, only these columns will be read from the row group. A + column name may be a prefix of a nested field, e.g. 'a' will select + 'a.b', 'a.c', and 'a.d.e'. + use_threads : bool, default True + Perform multi-threaded column reads. + use_pandas_metadata : bool, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded. + + Returns + ------- + pyarrow.table.Table + Content of the row groups as a table (of columns). + + Examples + -------- + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + >>> parquet_file.read_row_groups([0,0]) + pyarrow.Table + n_legs: int64 + animal: string + ---- + n_legs: [[2,2,4,4,5,...,2,4,4,5,100]] + animal: [["Flamingo","Parrot","Dog",...,"Brittle stars","Centipede"]] + """ + column_indices = self._get_column_indices( + columns, use_pandas_metadata=use_pandas_metadata) + return self.reader.read_row_groups(row_groups, + column_indices=column_indices, + use_threads=use_threads) + + def iter_batches(self, batch_size=65536, row_groups=None, columns=None, + use_threads=True, use_pandas_metadata=False): + """ + Read streaming batches from a Parquet file. + + Parameters + ---------- + batch_size : int, default 64K + Maximum number of records to yield per batch. Batches may be + smaller if there aren't enough rows in the file. + row_groups : list + Only these row groups will be read from the file. + columns : list + If not None, only these columns will be read from the file. A + column name may be a prefix of a nested field, e.g. 'a' will select + 'a.b', 'a.c', and 'a.d.e'. + use_threads : boolean, default True + Perform multi-threaded column reads. + use_pandas_metadata : boolean, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded. + + Yields + ------ + pyarrow.RecordBatch + Contents of each batch as a record batch + + Examples + -------- + Generate an example Parquet file: + + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + >>> for i in parquet_file.iter_batches(): + ... print("RecordBatch") + ... print(i.to_pandas()) + ... + RecordBatch + n_legs animal + 0 2 Flamingo + 1 2 Parrot + 2 4 Dog + 3 4 Horse + 4 5 Brittle stars + 5 100 Centipede + """ + if batch_size <= 0: + raise ValueError("batch_size must be greater than zero") + + if row_groups is None: + row_groups = range(0, self.metadata.num_row_groups) + column_indices = self._get_column_indices( + columns, use_pandas_metadata=use_pandas_metadata) + + batches = self.reader.iter_batches(batch_size, + row_groups=row_groups, + column_indices=column_indices, + use_threads=use_threads) + return batches + + def read(self, columns=None, use_threads=True, use_pandas_metadata=False): + """ + Read a Table from Parquet format. + + Parameters + ---------- + columns : list + If not None, only these columns will be read from the file. A + column name may be a prefix of a nested field, e.g. 'a' will select + 'a.b', 'a.c', and 'a.d.e'. + use_threads : bool, default True + Perform multi-threaded column reads. + use_pandas_metadata : bool, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded. + + Returns + ------- + pyarrow.table.Table + Content of the file as a table (of columns). + + Examples + -------- + Generate an example Parquet file: + + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + Read a Table: + + >>> parquet_file.read(columns=["animal"]) + pyarrow.Table + animal: string + ---- + animal: [["Flamingo","Parrot",...,"Brittle stars","Centipede"]] + """ + column_indices = self._get_column_indices( + columns, use_pandas_metadata=use_pandas_metadata) + return self.reader.read_all(column_indices=column_indices, + use_threads=use_threads) + + def scan_contents(self, columns=None, batch_size=65536): + """ + Read contents of file for the given columns and batch size. + + Notes + ----- + This function's primary purpose is benchmarking. + The scan is executed on a single thread. + + Parameters + ---------- + columns : list of integers, default None + Select columns to read, if None scan all columns. + batch_size : int, default 64K + Number of rows to read at a time internally. + + Returns + ------- + num_rows : int + Number of rows in file + + Examples + -------- + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'example.parquet') + >>> parquet_file = pq.ParquetFile('example.parquet') + + >>> parquet_file.scan_contents() + 6 + """ + column_indices = self._get_column_indices(columns) + return self.reader.scan_contents(column_indices, + batch_size=batch_size) + + def _get_column_indices(self, column_names, use_pandas_metadata=False): + if column_names is None: + return None + + indices = [] + + for name in column_names: + if name in self._nested_paths_by_prefix: + indices.extend(self._nested_paths_by_prefix[name]) + + if use_pandas_metadata: + file_keyvalues = self.metadata.metadata + common_keyvalues = (self.common_metadata.metadata + if self.common_metadata is not None + else None) + + if file_keyvalues and b'pandas' in file_keyvalues: + index_columns = _get_pandas_index_columns(file_keyvalues) + elif common_keyvalues and b'pandas' in common_keyvalues: + index_columns = _get_pandas_index_columns(common_keyvalues) + else: + index_columns = [] + + if indices is not None and index_columns: + indices += [self.reader.column_name_idx(descr) + for descr in index_columns + if not isinstance(descr, dict)] + + return indices + + +_SPARK_DISALLOWED_CHARS = re.compile('[ ,;{}()\n\t=]') + + +def _sanitized_spark_field_name(name): + return _SPARK_DISALLOWED_CHARS.sub('_', name) + + +def _sanitize_schema(schema, flavor): + if 'spark' in flavor: + sanitized_fields = [] + + schema_changed = False + + for field in schema: + name = field.name + sanitized_name = _sanitized_spark_field_name(name) + + if sanitized_name != name: + schema_changed = True + sanitized_field = pa.field(sanitized_name, field.type, + field.nullable, field.metadata) + sanitized_fields.append(sanitized_field) + else: + sanitized_fields.append(field) + + new_schema = pa.schema(sanitized_fields, metadata=schema.metadata) + return new_schema, schema_changed + else: + return schema, False + + +def _sanitize_table(table, new_schema, flavor): + # TODO: This will not handle prohibited characters in nested field names + if 'spark' in flavor: + column_data = [table[i] for i in range(table.num_columns)] + return pa.Table.from_arrays(column_data, schema=new_schema) + else: + return table + + +_parquet_writer_arg_docs = """version : {"1.0", "2.4", "2.6"}, default "2.6" + Determine which Parquet logical types are available for use, whether the + reduced set from the Parquet 1.x.x format or the expanded logical types + added in later format versions. + Files written with version='2.4' or '2.6' may not be readable in all + Parquet implementations, so version='1.0' is likely the choice that + maximizes file compatibility. + UINT32 and some logical types are only available with version '2.4'. + Nanosecond timestamps are only available with version '2.6'. + Other features such as compression algorithms or the new serialized + data page format must be enabled separately (see 'compression' and + 'data_page_version'). +use_dictionary : bool or list, default True + Specify if we should use dictionary encoding in general or only for + some columns. + When encoding the column, if the dictionary size is too large, the + column will fallback to ``PLAIN`` encoding. Specially, ``BOOLEAN`` type + doesn't support dictionary encoding. +compression : str or dict, default 'snappy' + Specify the compression codec, either on a general basis or per-column. + Valid values: {'NONE', 'SNAPPY', 'GZIP', 'BROTLI', 'LZ4', 'ZSTD'}. +write_statistics : bool or list, default True + Specify if we should write statistics in general (default is True) or only + for some columns. +use_deprecated_int96_timestamps : bool, default None + Write timestamps to INT96 Parquet format. Defaults to False unless enabled + by flavor argument. This take priority over the coerce_timestamps option. +coerce_timestamps : str, default None + Cast timestamps to a particular resolution. If omitted, defaults are chosen + depending on `version`. For ``version='1.0'`` and ``version='2.4'``, + nanoseconds are cast to microseconds ('us'), while for + ``version='2.6'`` (the default), they are written natively without loss + of resolution. Seconds are always cast to milliseconds ('ms') by default, + as Parquet does not have any temporal type with seconds resolution. + If the casting results in loss of data, it will raise an exception + unless ``allow_truncated_timestamps=True`` is given. + Valid values: {None, 'ms', 'us'} +allow_truncated_timestamps : bool, default False + Allow loss of data when coercing timestamps to a particular + resolution. E.g. if microsecond or nanosecond data is lost when coercing to + 'ms', do not raise an exception. Passing ``allow_truncated_timestamp=True`` + will NOT result in the truncation exception being ignored unless + ``coerce_timestamps`` is not None. +data_page_size : int, default None + Set a target threshold for the approximate encoded size of data + pages within a column chunk (in bytes). If None, use the default data page + size of 1MByte. +max_rows_per_page : int, default None + Maximum number of rows per page within a column chunk. + If None, use the default of 20000. + Smaller values reduce memory usage during reads but increase metadata overhead. +flavor : {'spark'}, default None + Sanitize schema or set other compatibility options to work with + various target systems. +filesystem : FileSystem, default None + If nothing passed, will be inferred from `where` if path-like, else + `where` is already a file-like object so no filesystem is needed. +compression_level : int or dict, default None + Specify the compression level for a codec, either on a general basis or + per-column. If None is passed, arrow selects the compression level for + the compression codec in use. The compression level has a different + meaning for each codec, so you have to read the documentation of the + codec you are using. + An exception is thrown if the compression codec does not allow specifying + a compression level. +use_byte_stream_split : bool or list, default False + Specify if the byte_stream_split encoding should be used in general or + only for some columns. If both dictionary and byte_stream_stream are + enabled, then dictionary is preferred. + The byte_stream_split encoding is valid for integer, floating-point + and fixed-size binary data types (including decimals); it should be + combined with a compression codec so as to achieve size reduction. +column_encoding : string or dict, default None + Specify the encoding scheme on a per column basis. + Can only be used when ``use_dictionary`` is set to False, and + cannot be used in combination with ``use_byte_stream_split``. + Currently supported values: {'PLAIN', 'BYTE_STREAM_SPLIT', + 'DELTA_BINARY_PACKED', 'DELTA_LENGTH_BYTE_ARRAY', 'DELTA_BYTE_ARRAY'}. + Certain encodings are only compatible with certain data types. + Please refer to the encodings section of `Reading and writing Parquet + files `_. +data_page_version : {"1.0", "2.0"}, default "1.0" + The serialized Parquet data page format version to write, defaults to + 1.0. This does not impact the file schema logical types and Arrow to + Parquet type casting behavior; for that use the "version" option. +use_compliant_nested_type : bool, default True + Whether to write compliant Parquet nested type (lists) as defined + `here `_, defaults to ``True``. + For ``use_compliant_nested_type=True``, this will write into a list + with 3-level structure where the middle level, named ``list``, + is a repeated group with a single field named ``element``:: + + group (LIST) { + repeated group list { + element; + } + } + + For ``use_compliant_nested_type=False``, this will also write into a list + with 3-level structure, where the name of the single field of the middle + level ``list`` is taken from the element name for nested columns in Arrow, + which defaults to ``item``:: + + group (LIST) { + repeated group list { + item; + } + } +encryption_properties : FileEncryptionProperties, default None + File encryption properties for Parquet Modular Encryption. + If None, no encryption will be done. + The encryption properties can be created using: + ``CryptoFactory.file_encryption_properties()``. +write_batch_size : int, default None + Number of values to write to a page at a time. If None, use the default of + 1024. ``write_batch_size`` is complementary to ``data_page_size``. If pages + are exceeding the ``data_page_size`` due to large column values, lowering + the batch size can help keep page sizes closer to the intended size. +dictionary_pagesize_limit : int, default None + Specify the dictionary page size limit per row group. If None, use the + default 1MB. +store_schema : bool, default True + By default, the Arrow schema is serialized and stored in the Parquet + file metadata (in the "ARROW:schema" key). When reading the file, + if this key is available, it will be used to more faithfully recreate + the original Arrow data. For example, for tz-aware timestamp columns + it will restore the timezone (Parquet only stores the UTC values without + timezone), or columns with duration type will be restored from the int64 + Parquet column. +write_page_index : bool, default False + Whether to write a page index in general for all columns. + Writing statistics to the page index disables the old method of writing + statistics to each data page header. The page index makes statistics-based + filtering more efficient than the page header, as it gathers all the + statistics for a Parquet file in a single place, avoiding scattered I/O. + Note that the page index is not yet used on the read size by PyArrow. +write_page_checksum : bool, default False + Whether to write page checksums in general for all columns. + Page checksums enable detection of data corruption, which might occur during + transmission or in the storage. +sorting_columns : Sequence of SortingColumn, default None + Specify the sort order of the data being written. The writer does not sort + the data nor does it verify that the data is sorted. The sort order is + written to the row group metadata, which can then be used by readers. +store_decimal_as_integer : bool, default False + Allow decimals with 1 <= precision <= 18 to be stored as integers. + In Parquet, DECIMAL can be stored in any of the following physical types: + - int32: for 1 <= precision <= 9. + - int64: for 10 <= precision <= 18. + - fixed_len_byte_array: precision is limited by the array size. + Length n can store <= floor(log_10(2^(8*n - 1) - 1)) base-10 digits. + - binary: precision is unlimited. The minimum number of bytes to store the + unscaled value is used. + + By default, this is DISABLED and all decimal types annotate fixed_len_byte_array. + When enabled, the writer will use the following physical types to store decimals: + - int32: for 1 <= precision <= 9. + - int64: for 10 <= precision <= 18. + - fixed_len_byte_array: for precision > 18. + + As a consequence, decimal columns stored in integer types are more compact. +use_content_defined_chunking : bool or dict, default False + Optimize parquet files for content addressable storage (CAS) systems by writing + data pages according to content-defined chunk boundaries. This allows for more + efficient deduplication of data across files, hence more efficient network + transfers and storage. The chunking is based on a rolling hash algorithm that + identifies chunk boundaries based on the actual content of the data. + + Note that it is an experimental feature and the API may change in the future. + + If set to ``True``, a default configuration is used with `min_chunk_size=256 KiB` + and `max_chunk_size=1024 KiB`. The chunk size distribution approximates a normal + distribution between `min_chunk_size` and `max_chunk_size` (sizes are accounted + before any Parquet encodings). + + A `dict` can be passed to adjust the chunker parameters with the following keys: + - `min_chunk_size`: minimum chunk size in bytes, default 256 KiB + The rolling hash will not be updated until this size is reached for each chunk. + Note that all data sent through the hash function is counted towards the chunk + size, including definition and repetition levels if present. + - `max_chunk_size`: maximum chunk size in bytes, default is 1024 KiB + The chunker will create a new chunk whenever the chunk size exceeds this value. + Note that the parquet writer has a related `data_pagesize` property that controls + the maximum size of a parquet data page after encoding. While setting + `data_page_size` to a smaller value than `max_chunk_size` doesn't affect the + chunking effectiveness, it results in more small parquet data pages. + - `norm_level`: normalization level to center the chunk size around the average + size more aggressively, default 0 + Increasing the normalization level increases the probability of finding a chunk, + improving the deduplication ratio, but also increasing the number of small chunks + resulting in many small parquet data pages. The default value provides a good + balance between deduplication ratio and fragmentation. Use norm_level=1 or + norm_level=2 to reach a higher deduplication ratio at the expense of + fragmentation. +write_time_adjusted_to_utc : bool, default False + Set the value of isAdjustedTOUTC when writing a TIME column. + If True, this tells the Parquet reader that the TIME columns + are expressed in reference to midnight in the UTC timezone. + If False (the default), the TIME columns are assumed to be expressed + in reference to midnight in an unknown, presumably local, timezone. +""" + +_parquet_writer_example_doc = """\ +Generate an example PyArrow Table and RecordBatch: + +>>> import pyarrow as pa +>>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], +... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]}) +>>> batch = pa.record_batch([[2, 2, 4, 4, 5, 100], +... ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]], +... names=['n_legs', 'animal']) + +create a ParquetWriter object: + +>>> import pyarrow.parquet as pq +>>> writer = pq.ParquetWriter('example.parquet', table.schema) + +and write the Table into the Parquet file: + +>>> writer.write_table(table) +>>> writer.close() + +>>> pq.read_table('example.parquet').to_pandas() + n_legs animal +0 2 Flamingo +1 2 Parrot +2 4 Dog +3 4 Horse +4 5 Brittle stars +5 100 Centipede + +create a ParquetWriter object for the RecordBatch: + +>>> writer2 = pq.ParquetWriter('example2.parquet', batch.schema) + +and write the RecordBatch into the Parquet file: + +>>> writer2.write_batch(batch) +>>> writer2.close() + +>>> pq.read_table('example2.parquet').to_pandas() + n_legs animal +0 2 Flamingo +1 2 Parrot +2 4 Dog +3 4 Horse +4 5 Brittle stars +5 100 Centipede +""" + + +class ParquetWriter: + + __doc__ = f""" +Class for incrementally building a Parquet file for Arrow tables. + +Parameters +---------- +where : path or file-like object +schema : pyarrow.Schema +{_parquet_writer_arg_docs} +writer_engine_version : unused +**options : dict + If options contains a key `metadata_collector` then the + corresponding value is assumed to be a list (or any object with + `.append` method) that will be filled with the file metadata instance + of the written file. + +Examples +-------- +{_parquet_writer_example_doc} +""" + + def __init__(self, where, schema, filesystem=None, + flavor=None, + version='2.6', + use_dictionary=True, + compression='snappy', + write_statistics=True, + use_deprecated_int96_timestamps=None, + compression_level=None, + use_byte_stream_split=False, + column_encoding=None, + writer_engine_version=None, + data_page_version='1.0', + use_compliant_nested_type=True, + encryption_properties=None, + write_batch_size=None, + dictionary_pagesize_limit=None, + store_schema=True, + write_page_index=False, + write_page_checksum=False, + sorting_columns=None, + store_decimal_as_integer=False, + write_time_adjusted_to_utc=False, + max_rows_per_page=None, + **options): + if use_deprecated_int96_timestamps is None: + # Use int96 timestamps for Spark + if flavor is not None and 'spark' in flavor: + use_deprecated_int96_timestamps = True + else: + use_deprecated_int96_timestamps = False + + self.flavor = flavor + if flavor is not None: + schema, self.schema_changed = _sanitize_schema(schema, flavor) + else: + self.schema_changed = False + + self.schema = schema + self.where = where + + # If we open a file using a filesystem, store file handle so we can be + # sure to close it when `self.close` is called. + self.file_handle = None + + filesystem, path = _resolve_filesystem_and_path(where, filesystem) + if filesystem is not None: + # ARROW-10480: do not auto-detect compression. While + # a filename like foo.parquet.gz is nonconforming, it + # shouldn't implicitly apply compression. + sink = self.file_handle = filesystem.open_output_stream( + path, compression=None) + else: + sink = where + self._metadata_collector = options.pop('metadata_collector', None) + engine_version = 'V2' + self.writer = _parquet.ParquetWriter( + sink, schema, + version=version, + compression=compression, + use_dictionary=use_dictionary, + write_statistics=write_statistics, + use_deprecated_int96_timestamps=use_deprecated_int96_timestamps, + compression_level=compression_level, + use_byte_stream_split=use_byte_stream_split, + column_encoding=column_encoding, + writer_engine_version=engine_version, + data_page_version=data_page_version, + use_compliant_nested_type=use_compliant_nested_type, + encryption_properties=encryption_properties, + write_batch_size=write_batch_size, + dictionary_pagesize_limit=dictionary_pagesize_limit, + store_schema=store_schema, + write_page_index=write_page_index, + write_page_checksum=write_page_checksum, + sorting_columns=sorting_columns, + store_decimal_as_integer=store_decimal_as_integer, + write_time_adjusted_to_utc=write_time_adjusted_to_utc, + max_rows_per_page=max_rows_per_page, + **options) + self.is_open = True + + def __del__(self): + if getattr(self, 'is_open', False): + self.close() + + def __enter__(self): + return self + + def __exit__(self, *args, **kwargs): + self.close() + # return false since we want to propagate exceptions + return False + + def write(self, table_or_batch, row_group_size=None): + """ + Write RecordBatch or Table to the Parquet file. + + Parameters + ---------- + table_or_batch : {RecordBatch, Table} + row_group_size : int, default None + Maximum number of rows in each written row group. If None, the row + group size will be the minimum of the number of rows in the + Table/RecordBatch and 1024 * 1024. + """ + if isinstance(table_or_batch, pa.RecordBatch): + self.write_batch(table_or_batch, row_group_size) + elif isinstance(table_or_batch, pa.Table): + self.write_table(table_or_batch, row_group_size) + else: + raise TypeError(type(table_or_batch)) + + def write_batch(self, batch, row_group_size=None): + """ + Write RecordBatch to the Parquet file. + + Parameters + ---------- + batch : RecordBatch + row_group_size : int, default None + Maximum number of rows in written row group. If None, the + row group size will be the minimum of the RecordBatch + size (in rows) and 1024 * 1024. If set larger than 64 * 1024 * 1024 + then 64 * 1024 * 1024 will be used instead. + """ + table = pa.Table.from_batches([batch], batch.schema) + self.write_table(table, row_group_size) + + def write_table(self, table, row_group_size=None): + """ + Write Table to the Parquet file. + + Parameters + ---------- + table : Table + row_group_size : int, default None + Maximum number of rows in each written row group. If None, + the row group size will be the minimum of the Table size (in rows) + and 1024 * 1024. If set larger than 64 * 1024 * 1024 then + 64 * 1024 * 1024 will be used instead. + + """ + if self.schema_changed: + table = _sanitize_table(table, self.schema, self.flavor) + assert self.is_open + + if not table.schema.equals(self.schema, check_metadata=False): + msg = ( + "Table schema does not match schema used to create file: \n" + f"table:\n{table.schema!s} vs. \nfile:\n{self.schema!s}" + ) + raise ValueError(msg) + + self.writer.write_table(table, row_group_size=row_group_size) + + def close(self): + """ + Close the connection to the Parquet file. + """ + if self.is_open: + self.writer.close() + self.is_open = False + if self._metadata_collector is not None: + self._metadata_collector.append(self.writer.metadata) + if self.file_handle is not None: + self.file_handle.close() + + def add_key_value_metadata(self, key_value_metadata): + """ + Add key-value metadata to the file. + This will overwrite any existing metadata with the same key. + + Parameters + ---------- + key_value_metadata : dict + Keys and values must be string-like / coercible to bytes. + """ + assert self.is_open + self.writer.add_key_value_metadata(key_value_metadata) + + +def _get_pandas_index_columns(keyvalues): + return (json.loads(keyvalues[b'pandas'].decode('utf8')) + ['index_columns']) + + +EXCLUDED_PARQUET_PATHS = {'_SUCCESS'} + + +_read_docstring_common = """\ +read_dictionary : list, default None + List of names or column paths (for nested types) to read directly + as DictionaryArray. Only supported for BYTE_ARRAY storage. To read + a flat column as dictionary-encoded pass the column name. For + nested types, you must pass the full column "path", which could be + something like level1.level2.list.item. Refer to the Parquet + file's schema to obtain the paths. +binary_type : pyarrow.DataType, default None + If given, Parquet binary columns will be read as this datatype. + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. +list_type : subclass of pyarrow.DataType, default None + If given, non-MAP repeated columns will be read as an instance of + this datatype (either pyarrow.ListType or pyarrow.LargeListType). + This setting is ignored if a serialized Arrow schema is found in + the Parquet metadata. +memory_map : bool, default False + If the source is a file path, use a memory map to read file, which can + improve performance in some environments. +buffer_size : int, default 0 + If positive, perform read buffering when deserializing individual + column chunks. Otherwise IO calls are unbuffered. +partitioning : pyarrow.dataset.Partitioning or str or list of str, \ +default "hive" + The partitioning scheme for a partitioned dataset. The default of "hive" + assumes directory names with key=value pairs like "/year=2009/month=11". + In addition, a scheme like "/2009/11" is also supported, in which case + you need to specify the field names or a full schema. See the + ``pyarrow.dataset.partitioning()`` function for more details.""" + + +_parquet_dataset_example = """\ +Generate an example PyArrow Table and write it to a partitioned dataset: + +>>> import pyarrow as pa +>>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], +... 'n_legs': [2, 2, 4, 4, 5, 100], +... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]}) +>>> import pyarrow.parquet as pq +>>> pq.write_to_dataset(table, root_path='dataset_v2', +... partition_cols=['year']) + +create a ParquetDataset object from the dataset source: + +>>> dataset = pq.ParquetDataset('dataset_v2/') + +and read the data: + +>>> dataset.read().to_pandas() + n_legs animal year +0 5 Brittle stars 2019 +1 2 Flamingo 2020 +2 4 Dog 2021 +3 100 Centipede 2021 +4 2 Parrot 2022 +5 4 Horse 2022 + +create a ParquetDataset object with filter: + +>>> dataset = pq.ParquetDataset('dataset_v2/', +... filters=[('n_legs','=',4)]) +>>> dataset.read().to_pandas() + n_legs animal year +0 4 Dog 2021 +1 4 Horse 2022 +""" + + +class ParquetDataset: + __doc__ = f""" +Encapsulates details of reading a complete Parquet dataset possibly +consisting of multiple files and partitions in subdirectories. + +Parameters +---------- +path_or_paths : str or List[str] + A directory name, single file name, or list of file names. +filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. +schema : pyarrow.parquet.Schema + Optionally provide the Schema for the Dataset, in which case it will + not be inferred from the source. +filters : pyarrow.compute.Expression or List[Tuple] or List[List[Tuple]], default None + Rows which do not match the filter predicate will be removed from scanned + data. Partition keys embedded in a nested directory structure will be + exploited to avoid loading files at all if they contain no matching rows. + Within-file level filtering and different partitioning schemes are supported. + + {_DNF_filter_doc} +{_read_docstring_common} +ignore_prefixes : list, optional + Files matching any of these prefixes will be ignored by the + discovery process. + This is matched to the basename of a path. + By default this is ['.', '_']. + Note that discovery happens only if a directory is passed as source. +pre_buffer : bool, default True + Coalesce and issue file reads in parallel to improve performance on + high-latency filesystems (e.g. S3, GCS). If True, Arrow will use a + background I/O thread pool. If using a filesystem layer that itself + performs readahead (e.g. fsspec's S3FS), disable readahead for best + results. Set to False if you want to prioritize minimal memory usage + over maximum speed. +coerce_int96_timestamp_unit : str, default None + Cast timestamps that are stored in INT96 format to a particular resolution + (e.g. 'ms'). Setting to None is equivalent to 'ns' and therefore INT96 + timestamps will be inferred as timestamps in nanoseconds. +decryption_properties : FileDecryptionProperties or None + File-level decryption properties. + The decryption properties can be created using + ``CryptoFactory.file_decryption_properties()``. +thrift_string_size_limit : int, default None + If not None, override the maximum total string size allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +thrift_container_size_limit : int, default None + If not None, override the maximum total size of containers allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +page_checksum_verification : bool, default False + If True, verify the page checksum for each page read from the file. +arrow_extensions_enabled : bool, default True + If True, read Parquet logical types as Arrow extension types where possible, + (e.g., read JSON as the canonical `arrow.json` extension type or UUID as + the canonical `arrow.uuid` extension type). + +Examples +-------- +{_parquet_dataset_example} +""" + + def __init__(self, path_or_paths, filesystem=None, schema=None, *, filters=None, + read_dictionary=None, binary_type=None, list_type=None, + memory_map=False, buffer_size=None, partitioning="hive", + ignore_prefixes=None, + pre_buffer=True, coerce_int96_timestamp_unit=None, + decryption_properties=None, thrift_string_size_limit=None, + thrift_container_size_limit=None, + page_checksum_verification=False, + arrow_extensions_enabled=True): + import pyarrow.dataset as ds + + # map format arguments + read_options = { + "pre_buffer": pre_buffer, + "coerce_int96_timestamp_unit": coerce_int96_timestamp_unit, + "thrift_string_size_limit": thrift_string_size_limit, + "thrift_container_size_limit": thrift_container_size_limit, + "page_checksum_verification": page_checksum_verification, + "arrow_extensions_enabled": arrow_extensions_enabled, + "binary_type": binary_type, + "list_type": list_type, + } + if buffer_size: + read_options.update(use_buffered_stream=True, + buffer_size=buffer_size) + if read_dictionary is not None: + read_options.update(dictionary_columns=read_dictionary) + + if decryption_properties is not None: + read_options.update(decryption_properties=decryption_properties) + + self._filter_expression = None + if filters is not None: + self._filter_expression = filters_to_expression(filters) + + # map old filesystems to new one + if filesystem is not None: + filesystem = _ensure_filesystem( + filesystem, use_mmap=memory_map) + elif filesystem is None and memory_map: + # if memory_map is specified, assume local file system (string + # path can in principle be URI for any filesystem) + filesystem = LocalFileSystem(use_mmap=memory_map) + + # This needs to be checked after _ensure_filesystem, because that + # handles the case of an fsspec LocalFileSystem + if ( + hasattr(path_or_paths, "__fspath__") and + filesystem is not None and + not isinstance(filesystem, LocalFileSystem) + ): + raise TypeError( + "Path-like objects with __fspath__ must only be used with " + f"local file systems, not {type(filesystem)}" + ) + + # check for single fragment dataset or dataset directory + single_file = None + self._base_dir = None + if not isinstance(path_or_paths, list): + if _is_path_like(path_or_paths): + filesystem, path_or_paths = _resolve_filesystem_and_path( + path_or_paths, filesystem, memory_map=memory_map + ) + finfo = filesystem.get_file_info(path_or_paths) + if finfo.type == FileType.Directory: + self._base_dir = path_or_paths + else: + single_file = path_or_paths + + parquet_format = ds.ParquetFileFormat(**read_options) + + if single_file is not None: + fragment = parquet_format.make_fragment(single_file, filesystem) + + self._dataset = ds.FileSystemDataset( + [fragment], schema=schema or fragment.physical_schema, + format=parquet_format, + filesystem=fragment.filesystem + ) + return + + # check partitioning to enable dictionary encoding + if partitioning == "hive": + partitioning = ds.HivePartitioning.discover( + infer_dictionary=True) + + self._dataset = ds.dataset(path_or_paths, filesystem=filesystem, + schema=schema, format=parquet_format, + partitioning=partitioning, + ignore_prefixes=ignore_prefixes) + + def equals(self, other): + if not isinstance(other, ParquetDataset): + raise TypeError('`other` must be an instance of ParquetDataset') + + return (self.schema == other.schema and + self._dataset.format == other._dataset.format and + self.filesystem == other.filesystem and + # self.fragments == other.fragments and + self.files == other.files) + + def __eq__(self, other): + try: + return self.equals(other) + except TypeError: + return NotImplemented + + @property + def schema(self): + """ + Schema of the Dataset. + + Examples + -------- + Generate an example dataset: + + >>> import pyarrow as pa + >>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset(table, root_path='dataset_v2_schema', + ... partition_cols=['year']) + >>> dataset = pq.ParquetDataset('dataset_v2_schema/') + + Read the schema: + + >>> dataset.schema + n_legs: int64 + animal: string + year: dictionary + """ + return self._dataset.schema + + def read(self, columns=None, use_threads=True, use_pandas_metadata=False): + """ + Read (multiple) Parquet files as a single pyarrow.Table. + + Parameters + ---------- + columns : List[str] + Names of columns to read from the dataset. The partition fields + are not automatically included. + use_threads : bool, default True + Perform multi-threaded column reads. + use_pandas_metadata : bool, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded. + + Returns + ------- + pyarrow.Table + Content of the file as a table (of columns). + + Examples + -------- + Generate an example dataset: + + >>> import pyarrow as pa + >>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset(table, root_path='dataset_v2_read', + ... partition_cols=['year']) + >>> dataset = pq.ParquetDataset('dataset_v2_read/') + + Read the dataset: + + >>> dataset.read(columns=["n_legs"]) + pyarrow.Table + n_legs: int64 + ---- + n_legs: [[5],[2],[4,100],[2,4]] + """ + # if use_pandas_metadata, we need to include index columns in the + # column selection, to be able to restore those in the pandas DataFrame + metadata = self.schema.metadata or {} + + if use_pandas_metadata: + # if the dataset schema metadata itself doesn't have pandas + # then try to get this from common file (for backwards compat) + if b"pandas" not in metadata: + common_metadata = self._get_common_pandas_metadata() + if common_metadata: + metadata = common_metadata + + if columns is not None and use_pandas_metadata: + if metadata and b'pandas' in metadata: + # RangeIndex can be represented as dict instead of column name + index_columns = [ + col for col in _get_pandas_index_columns(metadata) + if not isinstance(col, dict) + ] + columns = ( + list(columns) + list(set(index_columns) - set(columns)) + ) + + table = self._dataset.to_table( + columns=columns, filter=self._filter_expression, + use_threads=use_threads + ) + + # if use_pandas_metadata, restore the pandas metadata (which gets + # lost if doing a specific `columns` selection in to_table) + if use_pandas_metadata: + if metadata and b"pandas" in metadata: + new_metadata = table.schema.metadata or {} + new_metadata.update({b"pandas": metadata[b"pandas"]}) + table = table.replace_schema_metadata(new_metadata) + + return table + + def _get_common_pandas_metadata(self): + + if not self._base_dir: + return None + + metadata = None + for name in ["_common_metadata", "_metadata"]: + metadata_path = os.path.join(str(self._base_dir), name) + finfo = self.filesystem.get_file_info(metadata_path) + if finfo.is_file: + pq_meta = read_metadata( + metadata_path, filesystem=self.filesystem) + metadata = pq_meta.metadata + if metadata and b'pandas' in metadata: + break + + return metadata + + def read_pandas(self, **kwargs): + """ + Read dataset including pandas metadata, if any. Other arguments passed + through to :func:`read`, see docstring for further details. + + Parameters + ---------- + **kwargs : optional + Additional options for :func:`read` + + Examples + -------- + Generate an example parquet file: + + >>> import pyarrow as pa + >>> import pandas as pd + >>> df = pd.DataFrame({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> table = pa.Table.from_pandas(df) + >>> import pyarrow.parquet as pq + >>> pq.write_table(table, 'table_V2.parquet') + >>> dataset = pq.ParquetDataset('table_V2.parquet') + + Read the dataset with pandas metadata: + + >>> dataset.read_pandas(columns=["n_legs"]) + pyarrow.Table + n_legs: int64 + ---- + n_legs: [[2,2,4,4,5,100]] + + >>> dataset.read_pandas(columns=["n_legs"]).schema.pandas_metadata + {'index_columns': [{'kind': 'range', 'name': None, 'start': 0, ...} + """ + return self.read(use_pandas_metadata=True, **kwargs) + + @property + def fragments(self): + """ + A list of the Dataset source fragments or pieces with absolute + file paths. + + Examples + -------- + Generate an example dataset: + + >>> import pyarrow as pa + >>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset(table, root_path='dataset_v2_fragments', + ... partition_cols=['year']) + >>> dataset = pq.ParquetDataset('dataset_v2_fragments/') + + List the fragments: + + >>> dataset.fragments + [>> import pyarrow as pa + >>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset(table, root_path='dataset_v2_files', + ... partition_cols=['year']) + >>> dataset = pq.ParquetDataset('dataset_v2_files/') + + List the files: + + >>> dataset.files + ['dataset_v2_files/year=2019/...-0.parquet', ... + """ + return self._dataset.files + + @property + def filesystem(self): + """ + The filesystem type of the Dataset source. + """ + return self._dataset.filesystem + + @property + def partitioning(self): + """ + The partitioning of the Dataset source, if discovered. + """ + return self._dataset.partitioning + + +_read_table_docstring = """ +{0} + +Parameters +---------- +source : str, list of str, pyarrow.NativeFile, or file-like object + If a string is passed, can be a single file name or directory name. If a + list of strings is passed, should be file names. For file-like objects, + only read a single file. Use pyarrow.BufferReader to read a file contained + in a bytes or buffer-like object. +columns : list + If not None, only these columns will be read from the file. A column + name may be a prefix of a nested field, e.g. 'a' will select 'a.b', + 'a.c', and 'a.d.e'. If empty, no columns will be read. Note + that the table will still have the correct num_rows set despite having + no columns. +use_threads : bool, default True + Perform multi-threaded column reads. +schema : Schema, optional + Optionally provide the Schema for the parquet dataset, in which case it + will not be inferred from the source. +{1} +filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. +filters : pyarrow.compute.Expression or List[Tuple] or List[List[Tuple]], default None + Rows which do not match the filter predicate will be removed from scanned + data. Partition keys embedded in a nested directory structure will be + exploited to avoid loading files at all if they contain no matching rows. + Within-file level filtering and different partitioning schemes are supported. + + {3} +ignore_prefixes : list, optional + Files matching any of these prefixes will be ignored by the + discovery process. + This is matched to the basename of a path. + By default this is ['.', '_']. + Note that discovery happens only if a directory is passed as source. +pre_buffer : bool, default True + Coalesce and issue file reads in parallel to improve performance on + high-latency filesystems (e.g. S3). If True, Arrow will use a + background I/O thread pool. If using a filesystem layer that itself + performs readahead (e.g. fsspec's S3FS), disable readahead for best + results. +coerce_int96_timestamp_unit : str, default None + Cast timestamps that are stored in INT96 format to a particular + resolution (e.g. 'ms'). Setting to None is equivalent to 'ns' + and therefore INT96 timestamps will be inferred as timestamps + in nanoseconds. +decryption_properties : FileDecryptionProperties or None + File-level decryption properties. + The decryption properties can be created using + ``CryptoFactory.file_decryption_properties()``. +thrift_string_size_limit : int, default None + If not None, override the maximum total string size allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +thrift_container_size_limit : int, default None + If not None, override the maximum total size of containers allocated + when decoding Thrift structures. The default limit should be + sufficient for most Parquet files. +page_checksum_verification : bool, default False + If True, verify the checksum for each page read from the file. +arrow_extensions_enabled : bool, default True + If True, read Parquet logical types as Arrow extension types where possible, + (e.g., read JSON as the canonical `arrow.json` extension type or UUID as + the canonical `arrow.uuid` extension type). + +Returns +------- +{2} + +{4} +""" + +_read_table_example = """\ + +Examples +-------- + +Generate an example PyArrow Table and write it to a partitioned dataset: + +>>> import pyarrow as pa +>>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], +... 'n_legs': [2, 2, 4, 4, 5, 100], +... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]}) +>>> import pyarrow.parquet as pq +>>> pq.write_to_dataset(table, root_path='dataset_name_2', +... partition_cols=['year']) + +Read the data: + +>>> pq.read_table('dataset_name_2').to_pandas() + n_legs animal year +0 5 Brittle stars 2019 +1 2 Flamingo 2020 +2 4 Dog 2021 +3 100 Centipede 2021 +4 2 Parrot 2022 +5 4 Horse 2022 + + +Read only a subset of columns: + +>>> pq.read_table('dataset_name_2', columns=["n_legs", "animal"]) +pyarrow.Table +n_legs: int64 +animal: string +---- +n_legs: [[5],[2],[4,100],[2,4]] +animal: [["Brittle stars"],["Flamingo"],["Dog","Centipede"],["Parrot","Horse"]] + +Read a subset of columns and read one column as DictionaryArray: + +>>> pq.read_table('dataset_name_2', columns=["n_legs", "animal"], +... read_dictionary=["animal"]) +pyarrow.Table +n_legs: int64 +animal: dictionary +---- +n_legs: [[5],[2],[4,100],[2,4]] +animal: [ -- dictionary: +["Brittle stars"] -- indices: +[0], -- dictionary: +["Flamingo"] -- indices: +[0], -- dictionary: +["Dog","Centipede"] -- indices: +[0,1], -- dictionary: +["Parrot","Horse"] -- indices: +[0,1]] + +Read the table with filter: + +>>> pq.read_table('dataset_name_2', columns=["n_legs", "animal"], +... filters=[('n_legs','<',4)]).to_pandas() + n_legs animal +0 2 Flamingo +1 2 Parrot + +Read data from a single Parquet file: + +>>> pq.write_table(table, 'example.parquet') +>>> pq.read_table('dataset_name_2').to_pandas() + n_legs animal year +0 5 Brittle stars 2019 +1 2 Flamingo 2020 +2 4 Dog 2021 +3 100 Centipede 2021 +4 2 Parrot 2022 +5 4 Horse 2022 +""" + + +def read_table(source, *, columns=None, use_threads=True, + schema=None, use_pandas_metadata=False, read_dictionary=None, + binary_type=None, list_type=None, memory_map=False, buffer_size=0, + partitioning="hive", filesystem=None, filters=None, + ignore_prefixes=None, pre_buffer=True, + coerce_int96_timestamp_unit=None, + decryption_properties=None, thrift_string_size_limit=None, + thrift_container_size_limit=None, + page_checksum_verification=False, + arrow_extensions_enabled=True): + + try: + dataset = ParquetDataset( + source, + schema=schema, + filesystem=filesystem, + partitioning=partitioning, + memory_map=memory_map, + read_dictionary=read_dictionary, + binary_type=binary_type, + list_type=list_type, + buffer_size=buffer_size, + filters=filters, + ignore_prefixes=ignore_prefixes, + pre_buffer=pre_buffer, + coerce_int96_timestamp_unit=coerce_int96_timestamp_unit, + decryption_properties=decryption_properties, + thrift_string_size_limit=thrift_string_size_limit, + thrift_container_size_limit=thrift_container_size_limit, + page_checksum_verification=page_checksum_verification, + arrow_extensions_enabled=arrow_extensions_enabled, + ) + except ImportError: + # fall back on ParquetFile for simple cases when pyarrow.dataset + # module is not available + if filters is not None: + raise ValueError( + "the 'filters' keyword is not supported when the " + "pyarrow.dataset module is not available" + ) + if partitioning != "hive": + raise ValueError( + "the 'partitioning' keyword is not supported when the " + "pyarrow.dataset module is not available" + ) + if schema is not None: + raise ValueError( + "the 'schema' argument is not supported when the " + "pyarrow.dataset module is not available" + ) + if isinstance(source, list): + raise ValueError( + "the 'source' argument cannot be a list of files " + "when the pyarrow.dataset module is not available" + ) + + filesystem, path = _resolve_filesystem_and_path(source, filesystem) + if filesystem is not None: + if not filesystem.get_file_info(path).is_file: + raise ValueError( + "the 'source' argument should be " + "an existing parquet file and not a directory " + "when the pyarrow.dataset module is not available" + ) + + source = filesystem.open_input_file(path) + + dataset = ParquetFile( + source, read_dictionary=read_dictionary, + binary_type=binary_type, + list_type=list_type, + memory_map=memory_map, buffer_size=buffer_size, + pre_buffer=pre_buffer, + coerce_int96_timestamp_unit=coerce_int96_timestamp_unit, + decryption_properties=decryption_properties, + thrift_string_size_limit=thrift_string_size_limit, + thrift_container_size_limit=thrift_container_size_limit, + page_checksum_verification=page_checksum_verification, + ) + + return dataset.read(columns=columns, use_threads=use_threads, + use_pandas_metadata=use_pandas_metadata) + + +read_table.__doc__ = _read_table_docstring.format( + """Read a Table from Parquet format""", + "\n".join(("""use_pandas_metadata : bool, default False + If True and file has custom pandas schema metadata, ensure that + index columns are also loaded.""", _read_docstring_common)), + """pyarrow.Table + Content of the file as a table (of columns)""", + _DNF_filter_doc, _read_table_example) + + +def read_pandas(source, columns=None, **kwargs): + return read_table( + source, columns=columns, use_pandas_metadata=True, **kwargs + ) + + +read_pandas.__doc__ = _read_table_docstring.format( + 'Read a Table from Parquet format, also reading DataFrame\n' + 'index values if known in the file metadata', + "\n".join((_read_docstring_common, + """**kwargs + additional options for :func:`read_table`""")), + """pyarrow.Table + Content of the file as a Table of Columns, including DataFrame + indexes as columns""", + _DNF_filter_doc, "") + + +def write_table(table, where, row_group_size=None, version='2.6', + use_dictionary=True, compression='snappy', + write_statistics=True, + use_deprecated_int96_timestamps=None, + coerce_timestamps=None, + allow_truncated_timestamps=False, + data_page_size=None, flavor=None, + filesystem=None, + compression_level=None, + use_byte_stream_split=False, + column_encoding=None, + data_page_version='1.0', + use_compliant_nested_type=True, + encryption_properties=None, + write_batch_size=None, + dictionary_pagesize_limit=None, + store_schema=True, + write_page_index=False, + write_page_checksum=False, + sorting_columns=None, + store_decimal_as_integer=False, + write_time_adjusted_to_utc=False, + max_rows_per_page=None, + **kwargs): + # Implementor's note: when adding keywords here / updating defaults, also + # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions + row_group_size = kwargs.pop('chunk_size', row_group_size) + use_int96 = use_deprecated_int96_timestamps + try: + with ParquetWriter( + where, table.schema, + filesystem=filesystem, + version=version, + flavor=flavor, + use_dictionary=use_dictionary, + write_statistics=write_statistics, + coerce_timestamps=coerce_timestamps, + data_page_size=data_page_size, + allow_truncated_timestamps=allow_truncated_timestamps, + compression=compression, + use_deprecated_int96_timestamps=use_int96, + compression_level=compression_level, + use_byte_stream_split=use_byte_stream_split, + column_encoding=column_encoding, + data_page_version=data_page_version, + use_compliant_nested_type=use_compliant_nested_type, + encryption_properties=encryption_properties, + write_batch_size=write_batch_size, + dictionary_pagesize_limit=dictionary_pagesize_limit, + store_schema=store_schema, + write_page_index=write_page_index, + write_page_checksum=write_page_checksum, + sorting_columns=sorting_columns, + store_decimal_as_integer=store_decimal_as_integer, + write_time_adjusted_to_utc=write_time_adjusted_to_utc, + max_rows_per_page=max_rows_per_page, + **kwargs) as writer: + writer.write_table(table, row_group_size=row_group_size) + except Exception: + if _is_path_like(where): + try: + os.remove(_stringify_path(where)) + except os.error: + pass + raise + + +_write_table_example = """\ +Generate an example PyArrow Table: + +>>> import pyarrow as pa +>>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], +... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", +... "Brittle stars", "Centipede"]}) + +and write the Table into Parquet file: + +>>> import pyarrow.parquet as pq +>>> pq.write_table(table, 'example.parquet') + +Defining row group size for the Parquet file: + +>>> pq.write_table(table, 'example.parquet', row_group_size=3) + +Defining row group compression (default is Snappy): + +>>> pq.write_table(table, 'example.parquet', compression='none') + +Defining row group compression and encoding per-column: + +>>> pq.write_table(table, 'example.parquet', +... compression={'n_legs': 'snappy', 'animal': 'gzip'}, +... use_dictionary=['n_legs', 'animal']) + +Defining column encoding per-column: + +>>> pq.write_table(table, 'example.parquet', +... column_encoding={'animal':'PLAIN'}, +... use_dictionary=False) +""" + +write_table.__doc__ = f""" +Write a Table to Parquet format. + +Parameters +---------- +table : pyarrow.Table +where : string or pyarrow.NativeFile +row_group_size : int, default None + Maximum number of rows in each written row group. If None, the + row group size will be the minimum of the Table size (in rows) + and 1024 * 1024. If set larger than 64 * 1024 * 1024 then + 64 * 1024 * 1024 will be used instead. +{_parquet_writer_arg_docs} +**kwargs : optional + Additional options for ParquetWriter + +Examples +-------- +{_write_table_example} +""" + + +def write_to_dataset(table, root_path, partition_cols=None, + filesystem=None, schema=None, partitioning=None, + basename_template=None, use_threads=None, + file_visitor=None, existing_data_behavior=None, + **kwargs): + """Wrapper around dataset.write_dataset for writing a Table to + Parquet format by partitions. + For each combination of partition columns and values, + a subdirectories are created in the following + manner: + + root_dir/ + group1=value1 + group2=value1 + .parquet + group2=value2 + .parquet + group1=valueN + group2=value1 + .parquet + group2=valueN + .parquet + + Parameters + ---------- + table : pyarrow.Table + root_path : str, pathlib.Path + The root directory of the dataset. + partition_cols : list, + Column names by which to partition the dataset. + Columns are partitioned in the order they are given. + filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. + schema : Schema, optional + This Schema of the dataset. + partitioning : Partitioning or list[str], optional + The partitioning scheme specified with the + ``pyarrow.dataset.partitioning()`` function or a list of field names. + When providing a list of field names, you can use + ``partitioning_flavor`` to drive which partitioning type should be + used. + basename_template : str, optional + A template string used to generate basenames of written data files. + The token '{i}' will be replaced with an automatically incremented + integer. If not specified, it defaults to "guid-{i}.parquet". + use_threads : bool, default True + Write files in parallel. If enabled, then maximum parallelism will be + used determined by the number of available CPU cores. + file_visitor : function + If set, this function will be called with a WrittenFile instance + for each file created during the call. This object will have both + a path attribute and a metadata attribute. + + The path attribute will be a string containing the path to + the created file. + + The metadata attribute will be the parquet metadata of the file. + This metadata will have the file path attribute set and can be used + to build a _metadata file. The metadata attribute will be None if + the format is not parquet. + + Example visitor which simple collects the filenames created:: + + visited_paths = [] + + def file_visitor(written_file): + visited_paths.append(written_file.path) + + existing_data_behavior : 'overwrite_or_ignore' | 'error' | \ +'delete_matching' + Controls how the dataset will handle data that already exists in + the destination. The default behaviour is 'overwrite_or_ignore'. + + 'overwrite_or_ignore' will ignore any existing data and will + overwrite files with the same name as an output file. Other + existing files will be ignored. This behavior, in combination + with a unique basename_template for each write, will allow for + an append workflow. + + 'error' will raise an error if any data exists in the destination. + + 'delete_matching' is useful when you are writing a partitioned + dataset. The first time each partition directory is encountered + the entire directory will be deleted. This allows you to overwrite + old partitions completely. + **kwargs : dict, + Used as additional kwargs for :func:`pyarrow.dataset.write_dataset` + function for matching kwargs, and remainder to + :func:`pyarrow.dataset.ParquetFileFormat.make_write_options`. + See the docstring of :func:`write_table` and + :func:`pyarrow.dataset.write_dataset` for the available options. + Using `metadata_collector` in kwargs allows one to collect the + file metadata instances of dataset pieces. The file paths in the + ColumnChunkMetaData will be set relative to `root_path`. + + Examples + -------- + Generate an example PyArrow Table: + + >>> import pyarrow as pa + >>> table = pa.table({'year': [2020, 2022, 2021, 2022, 2019, 2021], + ... 'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + + and write it to a partitioned dataset: + + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset(table, root_path='dataset_name_3', + ... partition_cols=['year']) + >>> pq.ParquetDataset('dataset_name_3').files + ['dataset_name_3/year=2019/...-0.parquet', ... + + Write a single Parquet file into the root folder: + + >>> pq.write_to_dataset(table, root_path='dataset_name_4') + >>> pq.ParquetDataset('dataset_name_4/').files + ['dataset_name_4/...-0.parquet'] + """ + metadata_collector = kwargs.pop('metadata_collector', None) + + # Check for conflicting keywords + msg_confl = ( + "The '{1}' argument is not supported. " + "Use only '{0}' instead." + ) + if partition_cols is not None and partitioning is not None: + raise ValueError(msg_confl.format("partitioning", + "partition_cols")) + + if metadata_collector is not None and file_visitor is not None: + raise ValueError(msg_confl.format("file_visitor", + "metadata_collector")) + + import pyarrow.dataset as ds + + # extract write_dataset specific options + # reset assumed to go to make_write_options + write_dataset_kwargs = dict() + for key in inspect.signature(ds.write_dataset).parameters: + if key in kwargs: + write_dataset_kwargs[key] = kwargs.pop(key) + write_dataset_kwargs['max_rows_per_group'] = kwargs.pop( + 'row_group_size', kwargs.pop("chunk_size", None) + ) + + if metadata_collector is not None: + def file_visitor(written_file): + metadata_collector.append(written_file.metadata) + + # map format arguments + parquet_format = ds.ParquetFileFormat() + write_options = parquet_format.make_write_options(**kwargs) + + # map old filesystems to new one + if filesystem is not None: + filesystem = _ensure_filesystem(filesystem) + + if partition_cols: + part_schema = table.select(partition_cols).schema + partitioning = ds.partitioning(part_schema, flavor="hive") + + if basename_template is None: + basename_template = guid() + '-{i}.parquet' + + if existing_data_behavior is None: + existing_data_behavior = 'overwrite_or_ignore' + + ds.write_dataset( + table, root_path, filesystem=filesystem, + format=parquet_format, file_options=write_options, schema=schema, + partitioning=partitioning, use_threads=use_threads, + file_visitor=file_visitor, + basename_template=basename_template, + existing_data_behavior=existing_data_behavior, + **write_dataset_kwargs) + return + + +def write_metadata(schema, where, metadata_collector=None, filesystem=None, + **kwargs): + """ + Write metadata-only Parquet file from schema. This can be used with + `write_to_dataset` to generate `_common_metadata` and `_metadata` sidecar + files. + + Parameters + ---------- + schema : pyarrow.Schema + where : string or pyarrow.NativeFile + metadata_collector : list + where to collect metadata information. + filesystem : FileSystem, default None + If nothing passed, will be inferred from `where` if path-like, else + `where` is already a file-like object so no filesystem is needed. + **kwargs : dict, + Additional kwargs for ParquetWriter class. See docstring for + `ParquetWriter` for more information. + + Examples + -------- + Generate example data: + + >>> import pyarrow as pa + >>> table = pa.table({'n_legs': [2, 2, 4, 4, 5, 100], + ... 'animal': ["Flamingo", "Parrot", "Dog", "Horse", + ... "Brittle stars", "Centipede"]}) + + Write a dataset and collect metadata information. + + >>> metadata_collector = [] + >>> import pyarrow.parquet as pq + >>> pq.write_to_dataset( + ... table, 'dataset_metadata', + ... metadata_collector=metadata_collector) + + Write the `_common_metadata` parquet file without row groups statistics. + + >>> pq.write_metadata( + ... table.schema, 'dataset_metadata/_common_metadata') + + Write the `_metadata` parquet file with row groups statistics. + + >>> pq.write_metadata( + ... table.schema, 'dataset_metadata/_metadata', + ... metadata_collector=metadata_collector) + """ + filesystem, where = _resolve_filesystem_and_path(where, filesystem) + + if hasattr(where, "seek"): # file-like + cursor_position = where.tell() + + writer = ParquetWriter(where, schema, filesystem, **kwargs) + writer.close() + + if metadata_collector is not None: + # ParquetWriter doesn't expose the metadata until it's written. Write + # it and read it again. + metadata = read_metadata(where, filesystem=filesystem) + if hasattr(where, "seek"): + where.seek(cursor_position) # file-like, set cursor back. + + for m in metadata_collector: + metadata.append_row_groups(m) + if filesystem is not None: + with filesystem.open_output_stream(where) as f: + metadata.write_metadata_file(f) + else: + metadata.write_metadata_file(where) + + +def read_metadata(where, memory_map=False, decryption_properties=None, + filesystem=None): + """ + Read FileMetaData from footer of a single Parquet file. + + Parameters + ---------- + where : str (file path) or file-like object + memory_map : bool, default False + Create memory map when the source is a file path. + decryption_properties : FileDecryptionProperties, default None + Decryption properties for reading encrypted Parquet files. + filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. + + Returns + ------- + metadata : FileMetaData + The metadata of the Parquet file + + Examples + -------- + >>> import pyarrow as pa + >>> import pyarrow.parquet as pq + >>> table = pa.table({'n_legs': [4, 5, 100], + ... 'animal': ["Dog", "Brittle stars", "Centipede"]}) + >>> pq.write_table(table, 'example.parquet') + + >>> pq.read_metadata('example.parquet') + + created_by: parquet-cpp-arrow version ... + num_columns: 2 + num_rows: 3 + num_row_groups: 1 + format_version: 2.6 + serialized_size: ... + """ + filesystem, where = _resolve_filesystem_and_path(where, filesystem) + file_ctx = nullcontext() + if filesystem is not None: + file_ctx = where = filesystem.open_input_file(where) + + with file_ctx: + file = ParquetFile(where, memory_map=memory_map, + decryption_properties=decryption_properties) + return file.metadata + + +def read_schema(where, memory_map=False, decryption_properties=None, + filesystem=None): + """ + Read effective Arrow schema from Parquet file metadata. + + Parameters + ---------- + where : str (file path) or file-like object + memory_map : bool, default False + Create memory map when the source is a file path. + decryption_properties : FileDecryptionProperties, default None + Decryption properties for reading encrypted Parquet files. + filesystem : FileSystem, default None + If nothing passed, will be inferred based on path. + Path will try to be found in the local on-disk filesystem otherwise + it will be parsed as an URI to determine the filesystem. + + Returns + ------- + schema : pyarrow.Schema + The schema of the Parquet file + + Examples + -------- + >>> import pyarrow as pa + >>> import pyarrow.parquet as pq + >>> table = pa.table({'n_legs': [4, 5, 100], + ... 'animal': ["Dog", "Brittle stars", "Centipede"]}) + >>> pq.write_table(table, 'example.parquet') + + >>> pq.read_schema('example.parquet') + n_legs: int64 + animal: string + """ + filesystem, where = _resolve_filesystem_and_path(where, filesystem) + file_ctx = nullcontext() + if filesystem is not None: + file_ctx = where = filesystem.open_input_file(where) + + with file_ctx: + file = ParquetFile( + where, memory_map=memory_map, + decryption_properties=decryption_properties) + return file.schema.to_arrow_schema() + + +__all__ = ( + "ColumnChunkMetaData", + "ColumnSchema", + "FileDecryptionProperties", + "FileEncryptionProperties", + "FileMetaData", + "ParquetDataset", + "ParquetFile", + "ParquetLogicalType", + "ParquetReader", + "ParquetSchema", + "ParquetWriter", + "RowGroupMetaData", + "SortingColumn", + "Statistics", + "read_metadata", + "read_pandas", + "read_schema", + "read_table", + "write_metadata", + "write_table", + "write_to_dataset", + "_filters_to_expression", + "filters_to_expression", +) diff --git a/lib/python3.12/site-packages/pyarrow/parquet/encryption.py b/lib/python3.12/site-packages/pyarrow/parquet/encryption.py new file mode 100644 index 0000000000000000000000000000000000000000..df6eed913fa52da8c8c9f497f12953b6178439b6 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/parquet/encryption.py @@ -0,0 +1,23 @@ +# pylint: disable=unused-wildcard-import, unused-import + +# 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. +from pyarrow._parquet_encryption import (CryptoFactory, # noqa + EncryptionConfiguration, + DecryptionConfiguration, + KmsConnectionConfig, + KmsClient) diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/CMakeLists.txt b/lib/python3.12/site-packages/pyarrow/src/arrow/python/CMakeLists.txt new file mode 100644 index 0000000000000000000000000000000000000000..67508982eab82eacb6d00bb28986b79f6fed5078 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/CMakeLists.txt @@ -0,0 +1,19 @@ +# 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. + +arrow_install_all_headers("arrow/python") +add_subdirectory(vendored) diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/api.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/api.h new file mode 100644 index 0000000000000000000000000000000000000000..2af0963a9c0444bb858f10323f914e21747cebaf --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/api.h @@ -0,0 +1,29 @@ +// 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. + +#pragma once + +#include "arrow/python/arrow_to_pandas.h" +#include "arrow/python/common.h" +#include "arrow/python/datetime.h" +#include "arrow/python/helpers.h" +#include "arrow/python/inference.h" +#include "arrow/python/io.h" +#include "arrow/python/numpy_convert.h" +#include "arrow/python/numpy_to_arrow.h" +#include "arrow/python/python_to_arrow.h" +#include "arrow/python/util.h" diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.cc b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.cc new file mode 100644 index 0000000000000000000000000000000000000000..f163266f3b876796f8447d0eb356cc28652c9bac --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.cc @@ -0,0 +1,2659 @@ +// 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. + +// Functions for pandas conversion via NumPy + +#include "arrow/python/arrow_to_pandas.h" +#include "arrow/python/numpy_interop.h" // IWYU pragma: expand + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/array.h" +#include "arrow/buffer.h" +#include "arrow/datum.h" +#include "arrow/status.h" +#include "arrow/table.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/hashing.h" +#include "arrow/util/int_util.h" +#include "arrow/util/logging.h" +#include "arrow/util/macros.h" +#include "arrow/util/parallel.h" +#include "arrow/visit_type_inline.h" + +#include "arrow/compute/api.h" + +#include "arrow/python/arrow_to_python_internal.h" +#include "arrow/python/common.h" +#include "arrow/python/datetime.h" +#include "arrow/python/decimal.h" +#include "arrow/python/helpers.h" +#include "arrow/python/numpy_convert.h" +#include "arrow/python/numpy_internal.h" +#include "arrow/python/pyarrow.h" +#include "arrow/python/python_to_arrow.h" +#include "arrow/python/type_traits.h" + +namespace arrow { + +class MemoryPool; + +using internal::checked_cast; +using internal::CheckIndexBounds; +using internal::OptionalParallelFor; + +namespace py { +namespace { + +// Fix options for conversion of an inner (child) array. +PandasOptions MakeInnerOptions(PandasOptions options) { + // Make sure conversion of inner dictionary arrays always returns an array, + // not a dict {'indices': array, 'dictionary': array, 'ordered': bool} + options.decode_dictionaries = true; + options.categorical_columns.reset(); + options.strings_to_categorical = false; + + // In ARROW-7723, we found as a result of ARROW-3789 that second + // through microsecond resolution tz-aware timestamps were being promoted to + // use the DATETIME_NANO_TZ conversion path, yielding a datetime64[ns] NumPy + // array in this function. PyArray_GETITEM returns datetime.datetime for + // units second through microsecond but PyLong for nanosecond (because + // datetime.datetime does not support nanoseconds). + // We force the object conversion to preserve the value of the timezone. + // Nanoseconds are returned as integers. + options.coerce_temporal_nanoseconds = false; + + return options; +} + +// ---------------------------------------------------------------------- +// PyCapsule code for setting ndarray base to reference C++ object + +struct ArrayCapsule { + std::shared_ptr array; +}; + +struct BufferCapsule { + std::shared_ptr buffer; +}; + +void ArrayCapsule_Destructor(PyObject* capsule) { + delete reinterpret_cast(PyCapsule_GetPointer(capsule, "arrow::Array")); +} + +void BufferCapsule_Destructor(PyObject* capsule) { + delete reinterpret_cast(PyCapsule_GetPointer(capsule, "arrow::Buffer")); +} + +// ---------------------------------------------------------------------- +// pandas 0.x DataFrame conversion internals + +using internal::arrow_traits; +using internal::npy_traits; + +template +struct WrapBytes {}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyUnicode_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyUnicode_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyUnicode_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyBytes_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyBytes_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyBytes_FromStringAndSize(data, length); + } +}; + +template <> +struct WrapBytes { + static inline PyObject* Wrap(const char* data, int64_t length) { + return PyBytes_FromStringAndSize(data, length); + } +}; + +static inline bool ListTypeSupported(const DataType& type) { + switch (type.id()) { + case Type::BOOL: + case Type::UINT8: + case Type::INT8: + case Type::UINT16: + case Type::INT16: + case Type::UINT32: + case Type::INT32: + case Type::INT64: + case Type::UINT64: + case Type::HALF_FLOAT: + case Type::FLOAT: + case Type::DOUBLE: + case Type::DECIMAL128: + case Type::DECIMAL256: + case Type::BINARY: + case Type::LARGE_BINARY: + case Type::STRING: + case Type::LARGE_STRING: + case Type::DATE32: + case Type::DATE64: + case Type::STRUCT: + case Type::MAP: + case Type::TIME32: + case Type::TIME64: + case Type::TIMESTAMP: + case Type::DURATION: + case Type::DICTIONARY: + case Type::INTERVAL_MONTH_DAY_NANO: + case Type::NA: // empty list + // The above types are all supported. + return true; + case Type::FIXED_SIZE_LIST: + case Type::LIST: + case Type::LARGE_LIST: + case Type::LIST_VIEW: + case Type::LARGE_LIST_VIEW: { + const auto& list_type = checked_cast(type); + return ListTypeSupported(*list_type.value_type()); + } + case Type::EXTENSION: { + const auto& ext = checked_cast(*type.GetSharedPtr()); + return ListTypeSupported(*(ext.storage_type())); + } + default: + break; + } + return false; +} + +Status CapsulizeArray(const std::shared_ptr& arr, PyObject** out) { + auto capsule = new ArrayCapsule{{arr}}; + *out = PyCapsule_New(reinterpret_cast(capsule), "arrow::Array", + &ArrayCapsule_Destructor); + if (*out == nullptr) { + delete capsule; + RETURN_IF_PYERROR(); + } + return Status::OK(); +} + +Status CapsulizeBuffer(const std::shared_ptr& buffer, PyObject** out) { + auto capsule = new BufferCapsule{{buffer}}; + *out = PyCapsule_New(reinterpret_cast(capsule), "arrow::Buffer", + &BufferCapsule_Destructor); + if (*out == nullptr) { + delete capsule; + RETURN_IF_PYERROR(); + } + return Status::OK(); +} + +Status SetNdarrayBase(PyArrayObject* arr, PyObject* base) { + if (PyArray_SetBaseObject(arr, base) == -1) { + // Error occurred, trust that SetBaseObject sets the error state + Py_XDECREF(base); + RETURN_IF_PYERROR(); + } + return Status::OK(); +} + +Status SetBufferBase(PyArrayObject* arr, const std::shared_ptr& buffer) { + PyObject* base; + RETURN_NOT_OK(CapsulizeBuffer(buffer, &base)); + return SetNdarrayBase(arr, base); +} + +inline void set_numpy_metadata(int type, const DataType* datatype, PyArray_Descr* out) { + auto metadata = + reinterpret_cast(PyDataType_C_METADATA(out)); + if (type == NPY_DATETIME) { + if (datatype->id() == Type::TIMESTAMP) { + const auto& timestamp_type = checked_cast(*datatype); + metadata->meta.base = internal::NumPyFrequency(timestamp_type.unit()); + } else { + ARROW_DCHECK(false) + << "NPY_DATETIME views only supported for Arrow TIMESTAMP types"; + } + } else if (type == NPY_TIMEDELTA) { + ARROW_DCHECK_EQ(datatype->id(), Type::DURATION); + const auto& duration_type = checked_cast(*datatype); + metadata->meta.base = internal::NumPyFrequency(duration_type.unit()); + } +} + +Status PyArray_NewFromPool(int nd, npy_intp* dims, PyArray_Descr* descr, MemoryPool* pool, + PyObject** out) { + // ARROW-6570: Allocate memory from MemoryPool for a couple reasons + // + // * Track allocations + // * Get better performance through custom allocators + int64_t total_size = PyDataType_ELSIZE(descr); + for (int i = 0; i < nd; ++i) { + total_size *= dims[i]; + } + + ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBuffer(total_size, pool)); + *out = PyArray_NewFromDescr(&PyArray_Type, descr, nd, dims, + /*strides=*/nullptr, + /*data=*/buffer->mutable_data(), + /*flags=*/NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEABLE, + /*obj=*/nullptr); + if (*out == nullptr) { + RETURN_IF_PYERROR(); + // Trust that error set if NULL returned + } + return SetBufferBase(reinterpret_cast(*out), std::move(buffer)); +} + +template +inline const T* GetPrimitiveValues(const Array& arr) { + if (arr.length() == 0) { + return nullptr; + } + const int elsize = arr.type()->byte_width(); + const auto& prim_arr = checked_cast(arr); + return reinterpret_cast(prim_arr.values()->data() + arr.offset() * elsize); +} + +Status MakeNumPyView(std::shared_ptr arr, PyObject* py_ref, int npy_type, int ndim, + npy_intp* dims, PyObject** out) { + PyAcquireGIL lock; + + PyArray_Descr* descr = internal::GetSafeNumPyDtype(npy_type); + set_numpy_metadata(npy_type, arr->type().get(), descr); + PyObject* result = PyArray_NewFromDescr( + &PyArray_Type, descr, ndim, dims, /*strides=*/nullptr, + const_cast(GetPrimitiveValues(*arr)), /*flags=*/0, nullptr); + PyArrayObject* np_arr = reinterpret_cast(result); + if (np_arr == nullptr) { + // Error occurred, trust that error set + return Status::OK(); + } + + PyObject* base; + if (py_ref == nullptr) { + // Capsule will be owned by the ndarray, no incref necessary. See + // ARROW-1973 + RETURN_NOT_OK(CapsulizeArray(arr, &base)); + } else { + Py_INCREF(py_ref); + base = py_ref; + } + RETURN_NOT_OK(SetNdarrayBase(np_arr, base)); + + // Do not allow Arrow data to be mutated + PyArray_CLEARFLAGS(np_arr, NPY_ARRAY_WRITEABLE); + *out = result; + return Status::OK(); +} + +class PandasWriter { + public: + enum type { + OBJECT, + UINT8, + INT8, + UINT16, + INT16, + UINT32, + INT32, + UINT64, + INT64, + HALF_FLOAT, + FLOAT, + DOUBLE, + BOOL, + DATETIME_DAY, + DATETIME_SECOND, + DATETIME_MILLI, + DATETIME_MICRO, + DATETIME_NANO, + DATETIME_SECOND_TZ, + DATETIME_MILLI_TZ, + DATETIME_MICRO_TZ, + DATETIME_NANO_TZ, + TIMEDELTA_SECOND, + TIMEDELTA_MILLI, + TIMEDELTA_MICRO, + TIMEDELTA_NANO, + CATEGORICAL, + EXTENSION + }; + + PandasWriter(const PandasOptions& options, int64_t num_rows, int num_columns) + : options_(options), num_rows_(num_rows), num_columns_(num_columns) { + PyAcquireGIL lock; + internal::InitPandasStaticData(); + } + virtual ~PandasWriter() {} + + void SetBlockData(PyObject* arr) { + block_arr_.reset(arr); + block_data_ = + reinterpret_cast(PyArray_DATA(reinterpret_cast(arr))); + } + + /// \brief Either copy or wrap single array to create pandas-compatible array + /// for Series or DataFrame. num_columns_ can only be 1. Will try to zero + /// copy if possible (or error if not possible and zero_copy_only=True) + virtual Status TransferSingle(std::shared_ptr data, PyObject* py_ref) = 0; + + /// \brief Copy ChunkedArray into a multi-column block + virtual Status CopyInto(std::shared_ptr data, int64_t rel_placement) = 0; + + Status EnsurePlacementAllocated() { + std::lock_guard guard(allocation_lock_); + if (placement_data_ != nullptr) { + return Status::OK(); + } + PyAcquireGIL lock; + npy_intp placement_dims[1] = {num_columns_}; + PyObject* placement_arr = PyArray_SimpleNew(1, placement_dims, NPY_INT64); + RETURN_IF_PYERROR(); + placement_arr_.reset(placement_arr); + placement_data_ = reinterpret_cast( + PyArray_DATA(reinterpret_cast(placement_arr))); + return Status::OK(); + } + + Status EnsureAllocated() { + std::lock_guard guard(allocation_lock_); + if (block_data_ != nullptr) { + return Status::OK(); + } + RETURN_NOT_OK(Allocate()); + return Status::OK(); + } + + virtual bool CanZeroCopy(const ChunkedArray& data) const { return false; } + + virtual Status Write(std::shared_ptr data, int64_t abs_placement, + int64_t rel_placement) { + RETURN_NOT_OK(EnsurePlacementAllocated()); + if (num_columns_ == 1 && options_.allow_zero_copy_blocks) { + RETURN_NOT_OK(TransferSingle(data, /*py_ref=*/nullptr)); + } else { + RETURN_NOT_OK( + CheckNoZeroCopy("Cannot do zero copy conversion into " + "multi-column DataFrame block")); + RETURN_NOT_OK(EnsureAllocated()); + RETURN_NOT_OK(CopyInto(data, rel_placement)); + } + placement_data_[rel_placement] = abs_placement; + return Status::OK(); + } + + virtual Status GetDataFrameResult(PyObject** out) { + PyObject* result = PyDict_New(); + RETURN_IF_PYERROR(); + + PyObject* block; + RETURN_NOT_OK(GetResultBlock(&block)); + + PyDict_SetItemString(result, "block", block); + PyDict_SetItemString(result, "placement", placement_arr_.obj()); + + RETURN_NOT_OK(AddResultMetadata(result)); + *out = result; + return Status::OK(); + } + + // Caller steals the reference to this object + virtual Status GetSeriesResult(PyObject** out) { + RETURN_NOT_OK(MakeBlock1D()); + // Caller owns the object now + *out = block_arr_.detach(); + return Status::OK(); + } + + protected: + virtual Status AddResultMetadata(PyObject* result) { return Status::OK(); } + + Status MakeBlock1D() { + // For Series or for certain DataFrame block types, we need to shape to a + // 1D array when there is only one column + PyAcquireGIL lock; + + ARROW_DCHECK_EQ(1, num_columns_); + + npy_intp new_dims[1] = {static_cast(num_rows_)}; + PyArray_Dims dims; + dims.ptr = new_dims; + dims.len = 1; + + PyObject* reshaped = PyArray_Newshape( + reinterpret_cast(block_arr_.obj()), &dims, NPY_ANYORDER); + RETURN_IF_PYERROR(); + + // ARROW-8801: Here a PyArrayObject is created that is not being managed by + // any OwnedRef object. This object is then put in the resulting object + // with PyDict_SetItemString, which increments the reference count, so a + // memory leak ensues. There are several ways to fix the memory leak but a + // simple one is to put the reshaped 1D block array in this OwnedRefNoGIL + // so it will be correctly decref'd when this class is destructed. + block_arr_.reset(reshaped); + return Status::OK(); + } + + virtual Status GetResultBlock(PyObject** out) { + *out = block_arr_.obj(); + return Status::OK(); + } + + Status CheckNoZeroCopy(const std::string& message) { + if (options_.zero_copy_only) { + return Status::Invalid(message); + } + return Status::OK(); + } + + Status CheckNotZeroCopyOnly(const ChunkedArray& data) { + if (options_.zero_copy_only) { + return Status::Invalid("Needed to copy ", data.num_chunks(), " chunks with ", + data.null_count(), " nulls, but zero_copy_only was True"); + } + return Status::OK(); + } + + virtual Status Allocate() { + return Status::NotImplemented("Override Allocate in subclasses"); + } + + Status AllocateNDArray(int npy_type, int ndim = 2) { + PyAcquireGIL lock; + + PyObject* block_arr = nullptr; + npy_intp block_dims[2] = {0, 0}; + + if (ndim == 2) { + block_dims[0] = num_columns_; + block_dims[1] = num_rows_; + } else { + block_dims[0] = num_rows_; + } + PyArray_Descr* descr = internal::GetSafeNumPyDtype(npy_type); + if (PyDataType_REFCHK(descr)) { + // ARROW-6876: if the array has refcounted items, let Numpy + // own the array memory so as to decref elements on array destruction + block_arr = PyArray_SimpleNewFromDescr(ndim, block_dims, descr); + RETURN_IF_PYERROR(); + } else { + RETURN_NOT_OK( + PyArray_NewFromPool(ndim, block_dims, descr, options_.pool, &block_arr)); + } + + SetBlockData(block_arr); + return Status::OK(); + } + + void SetDatetimeUnit(NPY_DATETIMEUNIT unit) { + PyAcquireGIL lock; + auto date_dtype = + reinterpret_cast(PyDataType_C_METADATA( + PyArray_DESCR(reinterpret_cast(block_arr_.obj())))); + date_dtype->meta.base = unit; + } + + PandasOptions options_; + + std::mutex allocation_lock_; + + int64_t num_rows_; + int num_columns_; + + OwnedRefNoGIL block_arr_; + uint8_t* block_data_ = nullptr; + + // ndarray + OwnedRefNoGIL placement_arr_; + int64_t* placement_data_ = nullptr; + + private: + ARROW_DISALLOW_COPY_AND_ASSIGN(PandasWriter); +}; + +template +inline void ConvertIntegerWithNulls(const PandasOptions& options, + const ChunkedArray& data, OutType* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const InType* in_values = GetPrimitiveValues(arr); + // Upcast to double, set NaN as appropriate + + for (int i = 0; i < arr.length(); ++i) { + *out_values++ = + arr.IsNull(i) ? static_cast(NAN) : static_cast(in_values[i]); + } + } +} + +template +inline void ConvertIntegerNoNullsSameType(const PandasOptions& options, + const ChunkedArray& data, T* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + if (arr.length() > 0) { + const T* in_values = GetPrimitiveValues(arr); + memcpy(out_values, in_values, sizeof(T) * arr.length()); + out_values += arr.length(); + } + } +} + +template +inline void ConvertIntegerNoNullsCast(const PandasOptions& options, + const ChunkedArray& data, OutType* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const InType* in_values = GetPrimitiveValues(arr); + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values = in_values[i]; + } + } +} + +template +struct MemoizationTraits { + using Scalar = typename T::c_type; +}; + +template +struct MemoizationTraits> { + // For binary, we memoize string_view as a scalar value to avoid having to + // unnecessarily copy the memory into the memo table data structure + using Scalar = std::string_view; +}; + +// Generic Array -> PyObject** converter that handles object deduplication, if +// requested +template +inline Status ConvertAsPyObjects(const PandasOptions& options, const ChunkedArray& data, + WrapFunction&& wrap_func, PyObject** out_values) { + using ArrayType = typename TypeTraits::ArrayType; + using Scalar = typename MemoizationTraits::Scalar; + + auto convert_chunks = [&](auto&& wrap_func) -> Status { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = arrow::internal::checked_cast(*data.chunk(c)); + RETURN_NOT_OK(internal::WriteArrayObjects(arr, wrap_func, out_values)); + out_values += arr.length(); + } + return Status::OK(); + }; + + if (options.deduplicate_objects) { + // GH-40316: only allocate a memo table if deduplication is enabled. + ::arrow::internal::ScalarMemoTable memo_table(options.pool); + std::vector unique_values; + int32_t memo_size = 0; + + auto WrapMemoized = [&](const Scalar& value, PyObject** out_values) { + int32_t memo_index; + RETURN_NOT_OK(memo_table.GetOrInsert(value, &memo_index)); + if (memo_index == memo_size) { + // New entry + RETURN_NOT_OK(wrap_func(value, out_values)); + unique_values.push_back(*out_values); + ++memo_size; + } else { + // Duplicate entry + Py_INCREF(unique_values[memo_index]); + *out_values = unique_values[memo_index]; + } + return Status::OK(); + }; + return convert_chunks(std::move(WrapMemoized)); + } else { + return convert_chunks(std::forward(wrap_func)); + } +} + +Status ConvertStruct(PandasOptions options, const ChunkedArray& data, + PyObject** out_values) { + if (data.num_chunks() == 0) { + return Status::OK(); + } + // ChunkedArray has at least one chunk + auto arr = checked_cast(data.chunk(0).get()); + // Use it to cache the struct type and number of fields for all chunks + int32_t num_fields = arr->num_fields(); + auto array_type = arr->type(); + std::vector fields_data(num_fields * data.num_chunks()); + OwnedRef dict_item; + + // See notes in MakeInnerOptions. + options = MakeInnerOptions(std::move(options)); + // Don't blindly convert because timestamps in lists are handled differently. + options.timestamp_as_object = true; + + for (int c = 0; c < data.num_chunks(); c++) { + auto fields_data_offset = c * num_fields; + auto arr = checked_cast(data.chunk(c).get()); + // Convert the struct arrays first + for (int32_t i = 0; i < num_fields; i++) { + auto field = arr->field(static_cast(i)); + // In case the field is an extension array, use .storage() to convert to Pandas + if (field->type()->id() == Type::EXTENSION) { + const ExtensionArray& arr_ext = checked_cast(*field); + field = arr_ext.storage(); + } + RETURN_NOT_OK(ConvertArrayToPandas(options, field, nullptr, + fields_data[i + fields_data_offset].ref())); + ARROW_DCHECK(PyArray_Check(fields_data[i + fields_data_offset].obj())); + } + + // Construct a dictionary for each row + const bool has_nulls = data.null_count() > 0; + for (int64_t i = 0; i < arr->length(); ++i) { + if (has_nulls && arr->IsNull(i)) { + Py_INCREF(Py_None); + *out_values = Py_None; + } else { + // Build the new dict object for the row + dict_item.reset(PyDict_New()); + RETURN_IF_PYERROR(); + for (int32_t field_idx = 0; field_idx < num_fields; ++field_idx) { + OwnedRef field_value; + auto name = array_type->field(static_cast(field_idx))->name(); + if (!arr->field(static_cast(field_idx))->IsNull(i)) { + // Value exists in child array, obtain it + auto array = reinterpret_cast( + fields_data[field_idx + fields_data_offset].obj()); + auto ptr = reinterpret_cast(PyArray_GETPTR1(array, i)); + field_value.reset(PyArray_GETITEM(array, ptr)); + RETURN_IF_PYERROR(); + } else { + // Translate the Null to a None + Py_INCREF(Py_None); + field_value.reset(Py_None); + } + // PyDict_SetItemString increments reference count + auto setitem_result = + PyDict_SetItemString(dict_item.obj(), name.c_str(), field_value.obj()); + RETURN_IF_PYERROR(); + ARROW_DCHECK_EQ(setitem_result, 0); + } + *out_values = dict_item.obj(); + // Grant ownership to the resulting array + Py_INCREF(*out_values); + } + ++out_values; + } + } + return Status::OK(); +} + +Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr& dense_type, + ArrayVector* arrays) { + compute::ExecContext ctx(pool); + compute::CastOptions options; + for (size_t i = 0; i < arrays->size(); ++i) { + ARROW_ASSIGN_OR_RAISE((*arrays)[i], + compute::Cast(*(*arrays)[i], dense_type, options, &ctx)); + } + return Status::OK(); +} + +Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr& dense_type, + std::shared_ptr* array) { + auto chunks = (*array)->chunks(); + RETURN_NOT_OK(DecodeDictionaries(pool, dense_type, &chunks)); + *array = std::make_shared(std::move(chunks), dense_type); + return Status::OK(); +} + +template +enable_if_list_like ConvertListsLike(PandasOptions options, + const ChunkedArray& data, + PyObject** out_values) { + using ListArrayT = typename TypeTraits::ArrayType; + // Get column of underlying value arrays + ArrayVector value_arrays; + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + // values() does not account for offsets, so we need to slice into it. + // We can't use Flatten(), because it removes the values behind a null list + // value, and that makes the offsets into original list values and our + // flattened_values array different. + std::shared_ptr flattened_values = arr.values()->Slice( + arr.value_offset(0), arr.value_offset(arr.length()) - arr.value_offset(0)); + if (arr.value_type()->id() == Type::EXTENSION) { + const auto& arr_ext = checked_cast(*flattened_values); + value_arrays.emplace_back(arr_ext.storage()); + } else { + value_arrays.emplace_back(flattened_values); + } + } + + using ListArrayType = typename ListArrayT::TypeClass; + const auto& list_type = checked_cast(*data.type()); + auto value_type = list_type.value_type(); + if (value_type->id() == Type::EXTENSION) { + value_type = checked_cast(*value_type).storage_type(); + } + + auto flat_column = std::make_shared(value_arrays, value_type); + + options = MakeInnerOptions(std::move(options)); + + OwnedRefNoGIL owned_numpy_array; + RETURN_NOT_OK(ConvertChunkedArrayToPandas(options, flat_column, nullptr, + owned_numpy_array.ref())); + PyObject* numpy_array = owned_numpy_array.obj(); + ARROW_DCHECK(PyArray_Check(numpy_array)); + + int64_t chunk_offset = 0; + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + const bool has_nulls = data.null_count() > 0; + for (int64_t i = 0; i < arr.length(); ++i) { + if (has_nulls && arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values = Py_None; + } else { + // Need to subtract value_offset(0) since the original chunk might be a slice + // into another array. + OwnedRef start(PyLong_FromLongLong(arr.value_offset(i) + chunk_offset - + arr.value_offset(0))); + OwnedRef end(PyLong_FromLongLong(arr.value_offset(i + 1) + chunk_offset - + arr.value_offset(0))); + OwnedRef slice(PySlice_New(start.obj(), end.obj(), nullptr)); + + if (ARROW_PREDICT_FALSE(slice.obj() == nullptr)) { + // Fall out of loop, will return from RETURN_IF_PYERROR + break; + } + *out_values = PyObject_GetItem(numpy_array, slice.obj()); + + if (*out_values == nullptr) { + // Fall out of loop, will return from RETURN_IF_PYERROR + break; + } + } + ++out_values; + } + RETURN_IF_PYERROR(); + + chunk_offset += arr.value_offset(arr.length()) - arr.value_offset(0); + } + + return Status::OK(); +} + +// TODO GH-40579: optimize ListView conversion to avoid unnecessary copies +template +enable_if_list_view ConvertListsLike(PandasOptions options, + const ChunkedArray& data, + PyObject** out_values) { + using ListViewArrayType = typename TypeTraits::ArrayType; + using NonViewType = + std::conditional_t; + using NonViewClass = typename TypeTraits::ArrayType; + ArrayVector list_arrays; + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + ARROW_ASSIGN_OR_RAISE(auto non_view_array, + NonViewClass::FromListView(arr, options.pool)); + list_arrays.emplace_back(non_view_array); + } + auto chunked_array = std::make_shared(list_arrays); + return ConvertListsLike(options, *chunked_array, out_values); +} + +template +Status ConvertMapHelper(F1 resetRow, F2 addPairToRow, F3 stealRow, + const ChunkedArray& data, PyArrayObject* py_keys, + PyArrayObject* py_items, + // needed for null checks in items + const std::vector> item_arrays, + PyObject** out_values) { + OwnedRef key_value; + OwnedRef item_value; + + int64_t chunk_offset = 0; + for (int c = 0; c < data.num_chunks(); ++c) { + const auto& arr = checked_cast(*data.chunk(c)); + const bool has_nulls = data.null_count() > 0; + + // Make a list of key/item pairs for each row in array + for (int64_t i = 0; i < arr.length(); ++i) { + if (has_nulls && arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values = Py_None; + } else { + int64_t entry_offset = arr.value_offset(i); + int64_t num_pairs = arr.value_offset(i + 1) - entry_offset; + + // Build the new list object for the row of Python pairs + RETURN_NOT_OK(resetRow(num_pairs)); + + // Add each key/item pair in the row + for (int64_t j = 0; j < num_pairs; ++j) { + // Get key value, key is non-nullable for a valid row + auto ptr_key = reinterpret_cast( + PyArray_GETPTR1(py_keys, chunk_offset + entry_offset + j)); + key_value.reset(PyArray_GETITEM(py_keys, ptr_key)); + RETURN_IF_PYERROR(); + + if (item_arrays[c]->IsNull(entry_offset + j)) { + // Translate the Null to a None + Py_INCREF(Py_None); + item_value.reset(Py_None); + } else { + // Get valid value from item array + auto ptr_item = reinterpret_cast( + PyArray_GETPTR1(py_items, chunk_offset + entry_offset + j)); + item_value.reset(PyArray_GETITEM(py_items, ptr_item)); + RETURN_IF_PYERROR(); + } + + // Add the key/item pair to the row + RETURN_NOT_OK(addPairToRow(j, key_value, item_value)); + } + + // Pass ownership to the resulting array + *out_values = stealRow(); + } + ++out_values; + } + RETURN_IF_PYERROR(); + + chunk_offset += arr.values()->length(); + } + + return Status::OK(); +} + +// A more helpful error message around TypeErrors that may stem from unhashable keys +Status CheckMapAsPydictsTypeError() { + if (ARROW_PREDICT_TRUE(!PyErr_Occurred())) { + return Status::OK(); + } + if (PyErr_ExceptionMatches(PyExc_TypeError)) { + // Modify the error string directly, so it is re-raised + // with our additional info. + // + // There are not many interesting things happening when this + // is hit. This is intended to only be called directly after + // PyDict_SetItem, where a finite set of errors could occur. + PyObject *type, *value, *traceback; + PyErr_Fetch(&type, &value, &traceback); + std::string message; + RETURN_NOT_OK(internal::PyObject_StdStringStr(value, &message)); + message += + ". If keys are not hashable, then you must use the option " + "[maps_as_pydicts=None (default)]"; + + // resets the error + PyErr_SetString(PyExc_TypeError, message.c_str()); + } + return ConvertPyError(); +} + +Status CheckForDuplicateKeys(bool error_on_duplicate_keys, Py_ssize_t total_dict_len, + Py_ssize_t total_raw_len) { + if (total_dict_len < total_raw_len) { + const char* message = + "[maps_as_pydicts] " + "After conversion of Arrow maps to pydicts, " + "detected data loss due to duplicate keys. " + "Original input length is [%lld], total converted pydict length is [%lld]."; + std::array buf; + std::snprintf(buf.data(), buf.size(), message, total_raw_len, total_dict_len); + + if (error_on_duplicate_keys) { + return Status::UnknownError(buf.data()); + } else { + ARROW_LOG(WARNING) << buf.data(); + } + } + return Status::OK(); +} + +Status ConvertMap(PandasOptions options, const ChunkedArray& data, + PyObject** out_values) { + // Get columns of underlying key/item arrays + std::vector> key_arrays; + std::vector> item_arrays; + for (int c = 0; c < data.num_chunks(); ++c) { + const auto& map_arr = checked_cast(*data.chunk(c)); + key_arrays.emplace_back(map_arr.keys()); + item_arrays.emplace_back(map_arr.items()); + } + + const auto& map_type = checked_cast(*data.type()); + auto key_type = map_type.key_type(); + auto item_type = map_type.item_type(); + + // ARROW-6899: Convert dictionary-encoded children to dense instead of + // failing below. A more efficient conversion than this could be done later + if (key_type->id() == Type::DICTIONARY) { + auto dense_type = checked_cast(*key_type).value_type(); + RETURN_NOT_OK(DecodeDictionaries(options.pool, dense_type, &key_arrays)); + key_type = dense_type; + } + if (item_type->id() == Type::DICTIONARY) { + auto dense_type = checked_cast(*item_type).value_type(); + RETURN_NOT_OK(DecodeDictionaries(options.pool, dense_type, &item_arrays)); + item_type = dense_type; + } + + // See notes in MakeInnerOptions. + options = MakeInnerOptions(std::move(options)); + // Don't blindly convert because timestamps in lists are handled differently. + options.timestamp_as_object = true; + + auto flat_keys = std::make_shared(key_arrays, key_type); + auto flat_items = std::make_shared(item_arrays, item_type); + OwnedRefNoGIL owned_numpy_keys; + RETURN_NOT_OK( + ConvertChunkedArrayToPandas(options, flat_keys, nullptr, owned_numpy_keys.ref())); + OwnedRefNoGIL owned_numpy_items; + RETURN_NOT_OK( + ConvertChunkedArrayToPandas(options, flat_items, nullptr, owned_numpy_items.ref())); + PyArrayObject* py_keys = reinterpret_cast(owned_numpy_keys.obj()); + PyArrayObject* py_items = reinterpret_cast(owned_numpy_items.obj()); + + if (options.maps_as_pydicts == MapConversionType::DEFAULT) { + // The default behavior to express an Arrow MAP as a list of [(key, value), ...] pairs + OwnedRef list_item; + return ConvertMapHelper( + [&list_item](int64_t num_pairs) { + list_item.reset(PyList_New(num_pairs)); + return CheckPyError(); + }, + [&list_item](int64_t idx, OwnedRef& key_value, OwnedRef& item_value) { + PyList_SET_ITEM(list_item.obj(), idx, + PyTuple_Pack(2, key_value.obj(), item_value.obj())); + return CheckPyError(); + }, + [&list_item] { return list_item.detach(); }, data, py_keys, py_items, item_arrays, + out_values); + } else { + // Use a native pydict + OwnedRef dict_item; + Py_ssize_t total_dict_len{0}; + Py_ssize_t total_raw_len{0}; + + bool error_on_duplicate_keys; + if (options.maps_as_pydicts == MapConversionType::LOSSY) { + error_on_duplicate_keys = false; + } else if (options.maps_as_pydicts == MapConversionType::STRICT_) { + error_on_duplicate_keys = true; + } else { + auto val = std::underlying_type_t(options.maps_as_pydicts); + return Status::UnknownError("Received unknown option for maps_as_pydicts: " + + std::to_string(val)); + } + + auto status = ConvertMapHelper( + [&dict_item, &total_raw_len](int64_t num_pairs) { + total_raw_len += num_pairs; + dict_item.reset(PyDict_New()); + return CheckPyError(); + }, + [&dict_item]([[maybe_unused]] int64_t idx, OwnedRef& key_value, + OwnedRef& item_value) { + auto setitem_result = + PyDict_SetItem(dict_item.obj(), key_value.obj(), item_value.obj()); + ARROW_RETURN_NOT_OK(CheckMapAsPydictsTypeError()); + // returns -1 if there are internal errors around hashing/resizing + return setitem_result == 0 ? Status::OK() + : Status::UnknownError( + "[maps_as_pydicts] " + "Unexpected failure inserting Arrow (key, " + "value) pair into Python dict"); + }, + [&dict_item, &total_dict_len] { + total_dict_len += PyDict_Size(dict_item.obj()); + return dict_item.detach(); + }, + data, py_keys, py_items, item_arrays, out_values); + + ARROW_RETURN_NOT_OK(status); + // If there were no errors generating the pydicts, + // then check if we detected any data loss from duplicate keys. + return CheckForDuplicateKeys(error_on_duplicate_keys, total_dict_len, total_raw_len); + } +} + +template +inline void ConvertNumericNullable(const ChunkedArray& data, InType na_value, + OutType* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const InType* in_values = GetPrimitiveValues(arr); + + if (arr.null_count() > 0) { + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values++ = arr.IsNull(i) ? na_value : in_values[i]; + } + } else { + memcpy(out_values, in_values, sizeof(InType) * arr.length()); + out_values += arr.length(); + } + } +} + +template +inline void ConvertNumericNullableCast(const ChunkedArray& data, InType na_value, + OutType* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const InType* in_values = GetPrimitiveValues(arr); + + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values++ = arr.IsNull(i) ? static_cast(na_value) + : static_cast(in_values[i]); + } + } +} + +template +class TypedPandasWriter : public PandasWriter { + public: + using T = typename npy_traits::value_type; + + using PandasWriter::PandasWriter; + + Status TransferSingle(std::shared_ptr data, PyObject* py_ref) override { + if (CanZeroCopy(*data)) { + PyObject* wrapped; + npy_intp dims[2] = {static_cast(num_columns_), + static_cast(num_rows_)}; + RETURN_NOT_OK( + MakeNumPyView(data->chunk(0), py_ref, NPY_TYPE, /*ndim=*/2, dims, &wrapped)); + SetBlockData(wrapped); + return Status::OK(); + } else { + RETURN_NOT_OK(CheckNotZeroCopyOnly(*data)); + RETURN_NOT_OK(EnsureAllocated()); + return CopyInto(data, /*rel_placement=*/0); + } + } + + Status CheckTypeExact(const DataType& type, Type::type expected) { + if (type.id() != expected) { + return Status::NotImplemented("Cannot write Arrow data of type ", type.ToString(), + " to pandas block with NumPy type ", + GetNumPyTypeName(NPY_TYPE)); + } + return Status::OK(); + } + + T* GetBlockColumnStart(int64_t rel_placement) { + return reinterpret_cast(block_data_) + rel_placement * num_rows_; + } + + protected: + Status Allocate() override { return AllocateNDArray(NPY_TYPE); } +}; + +struct ObjectWriterVisitor { + const PandasOptions& options; + const ChunkedArray& data; + PyObject** out_values; + + Status Visit(const NullType& type) { + for (int c = 0; c < data.num_chunks(); c++) { + std::shared_ptr arr = data.chunk(c); + + for (int64_t i = 0; i < arr->length(); ++i) { + // All values are null + Py_INCREF(Py_None); + *out_values = Py_None; + ++out_values; + } + } + return Status::OK(); + } + + Status Visit(const BooleanType& type) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + + for (int64_t i = 0; i < arr.length(); ++i) { + if (arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values++ = Py_None; + } else if (arr.Value(i)) { + // True + Py_INCREF(Py_True); + *out_values++ = Py_True; + } else { + // False + Py_INCREF(Py_False); + *out_values++ = Py_False; + } + } + } + return Status::OK(); + } + + template + enable_if_integer Visit(const Type& type) { + using T = typename Type::c_type; + auto WrapValue = [](T value, PyObject** out) { + *out = std::is_signed::value ? PyLong_FromLongLong(value) + : PyLong_FromUnsignedLongLong(value); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + return ConvertAsPyObjects(options, data, WrapValue, out_values); + } + + template + enable_if_t::value || is_binary_view_like_type::value || + is_fixed_size_binary_type::value, + Status> + Visit(const Type& type) { + auto WrapValue = [](const std::string_view& view, PyObject** out) { + *out = WrapBytes::Wrap(view.data(), view.length()); + if (*out == nullptr) { + PyErr_Clear(); + return Status::UnknownError("Wrapping ", view, " failed"); + } + return Status::OK(); + }; + return ConvertAsPyObjects(options, data, WrapValue, out_values); + } + + template + enable_if_date Visit(const Type& type) { + auto WrapValue = [](typename Type::c_type value, PyObject** out) { + RETURN_NOT_OK(internal::PyDate_from_int(value, Type::UNIT, out)); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + return ConvertAsPyObjects(options, data, WrapValue, out_values); + } + + template + enable_if_time Visit(const Type& type) { + const TimeUnit::type unit = type.unit(); + auto WrapValue = [unit](typename Type::c_type value, PyObject** out) { + RETURN_NOT_OK(internal::PyTime_from_int(value, unit, out)); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + return ConvertAsPyObjects(options, data, WrapValue, out_values); + } + + template + enable_if_timestamp Visit(const Type& type) { + const TimeUnit::type unit = type.unit(); + OwnedRef tzinfo; + + auto ConvertTimezoneNaive = [&](typename Type::c_type value, PyObject** out) { + RETURN_NOT_OK(internal::PyDateTime_from_int(value, unit, out)); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + auto ConvertTimezoneAware = [&](typename Type::c_type value, PyObject** out) { + PyObject* naive_datetime; + RETURN_NOT_OK(ConvertTimezoneNaive(value, &naive_datetime)); + + // convert the timezone naive datetime object to timezone aware + // two step conversion of the datetime mimics Python's code: + // dt.replace(tzinfo=datetime.timezone.utc).astimezone(tzinfo) + // first step: replacing timezone with timezone.utc (replace method) + OwnedRef args(PyTuple_New(0)); + OwnedRef keywords(PyDict_New()); + PyDict_SetItemString(keywords.obj(), "tzinfo", PyDateTime_TimeZone_UTC); + OwnedRef naive_datetime_replace(PyObject_GetAttrString(naive_datetime, "replace")); + OwnedRef datetime_utc( + PyObject_Call(naive_datetime_replace.obj(), args.obj(), keywords.obj())); + // second step: adjust the datetime to tzinfo timezone (astimezone method) + *out = PyObject_CallMethod(datetime_utc.obj(), "astimezone", "O", tzinfo.obj()); + + // the timezone naive object is no longer required + Py_DECREF(naive_datetime); + RETURN_IF_PYERROR(); + + return Status::OK(); + }; + + if (!type.timezone().empty() && !options.ignore_timezone) { + // convert timezone aware + PyObject* tzobj; + ARROW_ASSIGN_OR_RAISE(tzobj, internal::StringToTzinfo(type.timezone())); + tzinfo.reset(tzobj); + RETURN_IF_PYERROR(); + RETURN_NOT_OK( + ConvertAsPyObjects(options, data, ConvertTimezoneAware, out_values)); + } else { + // convert timezone naive + RETURN_NOT_OK( + ConvertAsPyObjects(options, data, ConvertTimezoneNaive, out_values)); + } + + return Status::OK(); + } + + template + enable_if_t::value, Status> Visit( + const Type& type) { + OwnedRef args(PyTuple_New(0)); + OwnedRef kwargs(PyDict_New()); + RETURN_IF_PYERROR(); + auto to_date_offset = [&](const MonthDayNanoIntervalType::MonthDayNanos& interval, + PyObject** out) { + ARROW_DCHECK(internal::BorrowPandasDataOffsetType() != nullptr); + // DateOffset objects do not add nanoseconds component to pd.Timestamp. + // as of Pandas 1.3.3 + // (https://github.com/pandas-dev/pandas/issues/43892). + // So convert microseconds and remainder to preserve data + // but give users more expected results. + int64_t microseconds = interval.nanoseconds / 1000; + int64_t nanoseconds; + if (interval.nanoseconds >= 0) { + nanoseconds = interval.nanoseconds % 1000; + } else { + nanoseconds = -((-interval.nanoseconds) % 1000); + } + + PyDict_SetItemString(kwargs.obj(), "months", PyLong_FromLong(interval.months)); + PyDict_SetItemString(kwargs.obj(), "days", PyLong_FromLong(interval.days)); + PyDict_SetItemString(kwargs.obj(), "microseconds", + PyLong_FromLongLong(microseconds)); + PyDict_SetItemString(kwargs.obj(), "nanoseconds", PyLong_FromLongLong(nanoseconds)); + *out = + PyObject_Call(internal::BorrowPandasDataOffsetType(), args.obj(), kwargs.obj()); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + return ConvertAsPyObjects(options, data, to_date_offset, + out_values); + } + + template + Status VisitDecimal(const DecimalT& type) { + OwnedRef decimal; + OwnedRef Decimal; + RETURN_NOT_OK(internal::ImportModule("decimal", &decimal)); + RETURN_NOT_OK(internal::ImportFromModule(decimal.obj(), "Decimal", &Decimal)); + PyObject* decimal_constructor = Decimal.obj(); + + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + + for (int64_t i = 0; i < arr.length(); ++i) { + if (arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values++ = Py_None; + } else { + *out_values++ = + internal::DecimalFromString(decimal_constructor, arr.FormatValue(i)); + RETURN_IF_PYERROR(); + } + } + } + + return Status::OK(); + } + + Status Visit(const Decimal32Type& type) { + return VisitDecimal(type); + } + + Status Visit(const Decimal64Type& type) { + return VisitDecimal(type); + } + + Status Visit(const Decimal128Type& type) { + return VisitDecimal(type); + } + + Status Visit(const Decimal256Type& type) { + return VisitDecimal(type); + } + + template + enable_if_t::value || is_list_view_type::value, Status> Visit( + const T& type) { + if (!ListTypeSupported(*type.value_type())) { + return Status::NotImplemented( + "Not implemented type for conversion from List to Pandas: ", + type.value_type()->ToString()); + } + return ConvertListsLike(options, data, out_values); + } + + Status Visit(const MapType& type) { return ConvertMap(options, data, out_values); } + + Status Visit(const StructType& type) { + return ConvertStruct(options, data, out_values); + } + + template + enable_if_t::value || + std::is_same::value || + std::is_same::value || + std::is_same::value || + std::is_same::value || + (std::is_base_of::value && + !std::is_same::value) || + std::is_base_of::value, + Status> + Visit(const Type& type) { + return Status::NotImplemented("No implemented conversion to object dtype: ", + type.ToString()); + } +}; + +class ObjectWriter : public TypedPandasWriter { + public: + using TypedPandasWriter::TypedPandasWriter; + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + PyAcquireGIL lock; + ObjectWriterVisitor visitor{this->options_, *data, + this->GetBlockColumnStart(rel_placement)}; + return VisitTypeInline(*data->type(), &visitor); + } +}; + +static inline bool IsNonNullContiguous(const ChunkedArray& data) { + return data.num_chunks() == 1 && data.null_count() == 0; +} + +template +class IntWriter : public TypedPandasWriter { + public: + using ArrowType = typename npy_traits::TypeClass; + using TypedPandasWriter::TypedPandasWriter; + + bool CanZeroCopy(const ChunkedArray& data) const override { + return IsNonNullContiguous(data); + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + RETURN_NOT_OK(this->CheckTypeExact(*data->type(), ArrowType::type_id)); + ConvertIntegerNoNullsSameType( + this->options_, *data, this->GetBlockColumnStart(rel_placement)); + return Status::OK(); + } +}; + +template +class FloatWriter : public TypedPandasWriter { + public: + using ArrowType = typename npy_traits::TypeClass; + using TypedPandasWriter::TypedPandasWriter; + using T = typename ArrowType::c_type; + + bool CanZeroCopy(const ChunkedArray& data) const override { + return IsNonNullContiguous(data) && data.type()->id() == ArrowType::type_id; + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + Type::type in_type = data->type()->id(); + auto out_values = this->GetBlockColumnStart(rel_placement); + +#define INTEGER_CASE(IN_TYPE) \ + ConvertIntegerWithNulls(this->options_, *data, out_values); \ + break; + + switch (in_type) { + case Type::UINT8: + INTEGER_CASE(uint8_t); + case Type::INT8: + INTEGER_CASE(int8_t); + case Type::UINT16: + INTEGER_CASE(uint16_t); + case Type::INT16: + INTEGER_CASE(int16_t); + case Type::UINT32: + INTEGER_CASE(uint32_t); + case Type::INT32: + INTEGER_CASE(int32_t); + case Type::UINT64: + INTEGER_CASE(uint64_t); + case Type::INT64: + INTEGER_CASE(int64_t); + case Type::HALF_FLOAT: + ConvertNumericNullableCast(*data, npy_traits::na_sentinel, out_values); + case Type::FLOAT: + ConvertNumericNullableCast(*data, npy_traits::na_sentinel, out_values); + break; + case Type::DOUBLE: + ConvertNumericNullableCast(*data, npy_traits::na_sentinel, out_values); + break; + default: + return Status::NotImplemented("Cannot write Arrow data of type ", + data->type()->ToString(), + " to a Pandas floating point block"); + } + +#undef INTEGER_CASE + + return Status::OK(); + } +}; + +using UInt8Writer = IntWriter; +using Int8Writer = IntWriter; +using UInt16Writer = IntWriter; +using Int16Writer = IntWriter; +using UInt32Writer = IntWriter; +using Int32Writer = IntWriter; +using UInt64Writer = IntWriter; +using Int64Writer = IntWriter; +using Float16Writer = FloatWriter; +using Float32Writer = FloatWriter; +using Float64Writer = FloatWriter; + +class BoolWriter : public TypedPandasWriter { + public: + using TypedPandasWriter::TypedPandasWriter; + + Status TransferSingle(std::shared_ptr data, PyObject* py_ref) override { + RETURN_NOT_OK( + CheckNoZeroCopy("Zero copy conversions not possible with " + "boolean types")); + RETURN_NOT_OK(EnsureAllocated()); + return CopyInto(data, /*rel_placement=*/0); + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + RETURN_NOT_OK(this->CheckTypeExact(*data->type(), Type::BOOL)); + auto out_values = this->GetBlockColumnStart(rel_placement); + for (int c = 0; c < data->num_chunks(); c++) { + const auto& arr = checked_cast(*data->chunk(c)); + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values++ = static_cast(arr.Value(i)); + } + } + return Status::OK(); + } +}; + +// ---------------------------------------------------------------------- +// Date / timestamp types + +template +inline void ConvertDatetime(const ChunkedArray& data, int64_t* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const T* in_values = GetPrimitiveValues(arr); + + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values++ = arr.IsNull(i) ? kPandasTimestampNull + : (static_cast(in_values[i]) * SHIFT); + } + } +} + +template +void ConvertDatesShift(const ChunkedArray& data, int64_t* out_values) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = *data.chunk(c); + const T* in_values = GetPrimitiveValues(arr); + for (int64_t i = 0; i < arr.length(); ++i) { + *out_values++ = arr.IsNull(i) ? kPandasTimestampNull + : static_cast(in_values[i]) / SHIFT; + } + } +} + +class DatetimeDayWriter : public TypedPandasWriter { + public: + using TypedPandasWriter::TypedPandasWriter; + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + int64_t* out_values = this->GetBlockColumnStart(rel_placement); + const auto& type = checked_cast(*data->type()); + switch (type.unit()) { + case DateUnit::DAY: + ConvertDatesShift(*data, out_values); + break; + case DateUnit::MILLI: + ConvertDatesShift(*data, out_values); + break; + } + return Status::OK(); + } + + protected: + Status Allocate() override { + RETURN_NOT_OK(this->AllocateNDArray(NPY_DATETIME)); + SetDatetimeUnit(NPY_FR_D); + return Status::OK(); + } +}; + +template +class DatetimeWriter : public TypedPandasWriter { + public: + using TypedPandasWriter::TypedPandasWriter; + + bool CanZeroCopy(const ChunkedArray& data) const override { + if (data.type()->id() == Type::TIMESTAMP) { + const auto& type = checked_cast(*data.type()); + return IsNonNullContiguous(data) && type.unit() == UNIT; + } else { + return false; + } + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + const auto& ts_type = checked_cast(*data->type()); + ARROW_DCHECK_EQ(UNIT, ts_type.unit()) << "Should only call instances of this writer " + << "with arrays of the correct unit"; + ConvertNumericNullable(*data, kPandasTimestampNull, + this->GetBlockColumnStart(rel_placement)); + return Status::OK(); + } + + protected: + Status Allocate() override { + RETURN_NOT_OK(this->AllocateNDArray(NPY_DATETIME)); + SetDatetimeUnit(internal::NumPyFrequency(UNIT)); + return Status::OK(); + } +}; + +using DatetimeSecondWriter = DatetimeWriter; + +class DatetimeMilliWriter : public DatetimeWriter { + public: + using DatetimeWriter::DatetimeWriter; + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + Type::type type = data->type()->id(); + int64_t* out_values = this->GetBlockColumnStart(rel_placement); + if (type == Type::DATE32) { + // Convert from days since epoch to datetime64[ms] + ConvertDatetime(*data, out_values); + } else if (type == Type::DATE64) { + ConvertNumericNullable(*data, kPandasTimestampNull, out_values); + } else { + const auto& ts_type = checked_cast(*data->type()); + ARROW_DCHECK_EQ(TimeUnit::MILLI, ts_type.unit()) + << "Should only call instances of this writer " + << "with arrays of the correct unit"; + ConvertNumericNullable(*data, kPandasTimestampNull, out_values); + } + return Status::OK(); + } +}; + +using DatetimeMicroWriter = DatetimeWriter; + +class DatetimeNanoWriter : public DatetimeWriter { + public: + using DatetimeWriter::DatetimeWriter; + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + Type::type type = data->type()->id(); + int64_t* out_values = this->GetBlockColumnStart(rel_placement); + compute::ExecContext ctx(options_.pool); + compute::CastOptions options; + if (options_.safe_cast) { + options = compute::CastOptions::Safe(); + } else { + options = compute::CastOptions::Unsafe(); + } + Datum out; + auto target_type = timestamp(TimeUnit::NANO); + + if (type == Type::DATE32) { + // Convert from days since epoch to datetime64[ns] + ConvertDatetime(*data, out_values); + } else if (type == Type::DATE64) { + // Date64Type is millisecond timestamp stored as int64_t + // TODO(wesm): Do we want to make sure to zero out the milliseconds? + ConvertDatetime(*data, out_values); + } else if (type == Type::TIMESTAMP) { + const auto& ts_type = checked_cast(*data->type()); + + if (ts_type.unit() == TimeUnit::NANO) { + ConvertNumericNullable(*data, kPandasTimestampNull, out_values); + } else if (ts_type.unit() == TimeUnit::MICRO || ts_type.unit() == TimeUnit::MILLI || + ts_type.unit() == TimeUnit::SECOND) { + ARROW_ASSIGN_OR_RAISE(out, compute::Cast(data, target_type, options, &ctx)); + ConvertNumericNullable(*out.chunked_array(), kPandasTimestampNull, + out_values); + } else { + return Status::NotImplemented("Unsupported time unit"); + } + } else { + return Status::NotImplemented("Cannot write Arrow data of type ", + data->type()->ToString(), + " to a Pandas datetime block."); + } + return Status::OK(); + } +}; + +template +class DatetimeTZWriter : public BASE { + public: + DatetimeTZWriter(const PandasOptions& options, const std::string& timezone, + int64_t num_rows) + : BASE(options, num_rows, 1), timezone_(timezone) {} + + protected: + Status GetResultBlock(PyObject** out) override { + RETURN_NOT_OK(this->MakeBlock1D()); + *out = this->block_arr_.obj(); + return Status::OK(); + } + + Status AddResultMetadata(PyObject* result) override { + PyObject* py_tz = PyUnicode_FromStringAndSize( + timezone_.c_str(), static_cast(timezone_.size())); + RETURN_IF_PYERROR(); + PyDict_SetItemString(result, "timezone", py_tz); + Py_DECREF(py_tz); + return Status::OK(); + } + + private: + std::string timezone_; +}; + +using DatetimeSecondTZWriter = DatetimeTZWriter; +using DatetimeMilliTZWriter = DatetimeTZWriter; +using DatetimeMicroTZWriter = DatetimeTZWriter; +using DatetimeNanoTZWriter = DatetimeTZWriter; + +template +class TimedeltaWriter : public TypedPandasWriter { + public: + using TypedPandasWriter::TypedPandasWriter; + + Status AllocateTimedelta(int ndim) { + RETURN_NOT_OK(this->AllocateNDArray(NPY_TIMEDELTA, ndim)); + SetDatetimeUnit(internal::NumPyFrequency(UNIT)); + return Status::OK(); + } + + bool CanZeroCopy(const ChunkedArray& data) const override { + const auto& type = checked_cast(*data.type()); + return IsNonNullContiguous(data) && type.unit() == UNIT; + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + const auto& type = checked_cast(*data->type()); + ARROW_DCHECK_EQ(UNIT, type.unit()) << "Should only call instances of this writer " + << "with arrays of the correct unit"; + ConvertNumericNullable(*data, kPandasTimestampNull, + this->GetBlockColumnStart(rel_placement)); + return Status::OK(); + } + + protected: + Status Allocate() override { return AllocateTimedelta(2); } +}; + +using TimedeltaSecondWriter = TimedeltaWriter; +using TimedeltaMilliWriter = TimedeltaWriter; +using TimedeltaMicroWriter = TimedeltaWriter; + +class TimedeltaNanoWriter : public TimedeltaWriter { + public: + using TimedeltaWriter::TimedeltaWriter; + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + Type::type type = data->type()->id(); + int64_t* out_values = this->GetBlockColumnStart(rel_placement); + if (type == Type::DURATION) { + const auto& ts_type = checked_cast(*data->type()); + if (ts_type.unit() == TimeUnit::NANO) { + ConvertNumericNullable(*data, kPandasTimestampNull, out_values); + } else if (ts_type.unit() == TimeUnit::MICRO) { + ConvertDatetime(*data, out_values); + } else if (ts_type.unit() == TimeUnit::MILLI) { + ConvertDatetime(*data, out_values); + } else if (ts_type.unit() == TimeUnit::SECOND) { + ConvertDatetime(*data, out_values); + } else { + return Status::NotImplemented("Unsupported time unit"); + } + } else { + return Status::NotImplemented("Cannot write Arrow data of type ", + data->type()->ToString(), + " to a Pandas timedelta block."); + } + return Status::OK(); + } +}; + +Status MakeZeroLengthArray(const std::shared_ptr& type, + std::shared_ptr* out) { + std::unique_ptr builder; + RETURN_NOT_OK(MakeBuilder(default_memory_pool(), type, &builder)); + RETURN_NOT_OK(builder->Resize(0)); + return builder->Finish(out); +} + +bool NeedDictionaryUnification(const ChunkedArray& data) { + if (data.num_chunks() < 2) { + return false; + } + const auto& arr_first = checked_cast(*data.chunk(0)); + for (int c = 1; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + if (!(arr_first.dictionary()->Equals(arr.dictionary()))) { + return true; + } + } + return false; +} + +template +class CategoricalWriter + : public TypedPandasWriter::npy_type> { + public: + using TRAITS = arrow_traits; + using ArrayType = typename TypeTraits::ArrayType; + using T = typename TRAITS::T; + + explicit CategoricalWriter(const PandasOptions& options, int64_t num_rows) + : TypedPandasWriter(options, num_rows, 1), + ordered_(false), + needs_copy_(false) {} + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + return Status::NotImplemented("categorical type"); + } + + Status TransferSingle(std::shared_ptr data, PyObject* py_ref) override { + const auto& dict_type = checked_cast(*data->type()); + std::shared_ptr dict; + if (data->num_chunks() == 0) { + // no dictionary values => create empty array + RETURN_NOT_OK(this->AllocateNDArray(TRAITS::npy_type, 1)); + RETURN_NOT_OK(MakeZeroLengthArray(dict_type.value_type(), &dict)); + } else { + ARROW_DCHECK_EQ(IndexType::type_id, dict_type.index_type()->id()); + RETURN_NOT_OK(WriteIndices(*data, &dict)); + } + + PyObject* pydict; + RETURN_NOT_OK(ConvertArrayToPandas(this->options_, dict, nullptr, &pydict)); + dictionary_.reset(pydict); + ordered_ = dict_type.ordered(); + return Status::OK(); + } + + Status Write(std::shared_ptr data, int64_t abs_placement, + int64_t rel_placement) override { + RETURN_NOT_OK(this->EnsurePlacementAllocated()); + RETURN_NOT_OK(TransferSingle(data, /*py_ref=*/nullptr)); + this->placement_data_[rel_placement] = abs_placement; + return Status::OK(); + } + + Status GetSeriesResult(PyObject** out) override { + PyAcquireGIL lock; + + PyObject* result = PyDict_New(); + RETURN_IF_PYERROR(); + + // Expected single array dictionary layout + PyDict_SetItemString(result, "indices", this->block_arr_.obj()); + RETURN_IF_PYERROR(); + RETURN_NOT_OK(AddResultMetadata(result)); + + *out = result; + return Status::OK(); + } + + protected: + Status AddResultMetadata(PyObject* result) override { + PyDict_SetItemString(result, "dictionary", dictionary_.obj()); + PyObject* py_ordered = ordered_ ? Py_True : Py_False; + Py_INCREF(py_ordered); + PyDict_SetItemString(result, "ordered", py_ordered); + return Status::OK(); + } + + Status WriteIndicesUniform(const ChunkedArray& data) { + // For unsigned types, upcast to signed since pandas uses -1 for nulls + // uint8 to int16, uint16 to int32, uint32 to int64, signed types unchanged + using OutputType = std::conditional_t< + std::is_same::value, int16_t, + std::conditional_t< + std::is_same::value, int32_t, + std::conditional_t::value, int64_t, T>>>; + const int npy_output_type = std::is_same::value ? NPY_INT16 + : std::is_same::value ? NPY_INT32 + : std::is_same::value + ? NPY_INT64 + : TRAITS::npy_type; + + RETURN_NOT_OK(this->AllocateNDArray(npy_output_type, 1)); + auto out_values = reinterpret_cast(this->block_data_); + + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + const auto& indices = checked_cast(*arr.indices()); + auto values = reinterpret_cast(indices.raw_values()); + + RETURN_NOT_OK(CheckIndexBounds(*indices.data(), arr.dictionary()->length())); + // Null is -1 in CategoricalBlock + for (int i = 0; i < arr.length(); ++i) { + if (indices.IsValid(i)) { + *out_values++ = static_cast(values[i]); + } else { + *out_values++ = -1; + } + } + } + return Status::OK(); + } + + Status WriteIndicesVarying(const ChunkedArray& data, std::shared_ptr* out_dict) { + // Yield int32 indices to allow for dictionary outgrowing the current index + // type + RETURN_NOT_OK(this->AllocateNDArray(NPY_INT32, 1)); + auto out_values = reinterpret_cast(this->block_data_); + + const auto& dict_type = checked_cast(*data.type()); + + ARROW_ASSIGN_OR_RAISE(auto unifier, DictionaryUnifier::Make(dict_type.value_type(), + this->options_.pool)); + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + const auto& indices = checked_cast(*arr.indices()); + auto values = reinterpret_cast(indices.raw_values()); + + std::shared_ptr transpose_buffer; + RETURN_NOT_OK(unifier->Unify(*arr.dictionary(), &transpose_buffer)); + + auto transpose = reinterpret_cast(transpose_buffer->data()); + int64_t dict_length = arr.dictionary()->length(); + + RETURN_NOT_OK(CheckIndexBounds(*indices.data(), dict_length)); + + // Null is -1 in CategoricalBlock + for (int i = 0; i < arr.length(); ++i) { + if (indices.IsValid(i)) { + *out_values++ = transpose[values[i]]; + } else { + *out_values++ = -1; + } + } + } + + std::shared_ptr unused_type; + return unifier->GetResult(&unused_type, out_dict); + } + + Status WriteIndices(const ChunkedArray& data, std::shared_ptr* out_dict) { + ARROW_DCHECK_GT(data.num_chunks(), 0); + + // Sniff the first chunk + const auto& arr_first = checked_cast(*data.chunk(0)); + const auto indices_first = std::static_pointer_cast(arr_first.indices()); + + // For unsigned types, we need to convert to signed for pandas compatibility + // even when there are no nulls, so we skip the fast path + const bool is_unsigned = std::is_unsigned::value; + + if (data.num_chunks() == 1 && indices_first->null_count() == 0 && !is_unsigned) { + RETURN_NOT_OK( + CheckIndexBounds(*indices_first->data(), arr_first.dictionary()->length())); + + PyObject* wrapped; + npy_intp dims[1] = {static_cast(this->num_rows_)}; + RETURN_NOT_OK(MakeNumPyView(indices_first, /*py_ref=*/nullptr, TRAITS::npy_type, + /*ndim=*/1, dims, &wrapped)); + this->SetBlockData(wrapped); + *out_dict = arr_first.dictionary(); + } else { + RETURN_NOT_OK(this->CheckNotZeroCopyOnly(data)); + if (NeedDictionaryUnification(data)) { + RETURN_NOT_OK(WriteIndicesVarying(data, out_dict)); + } else { + RETURN_NOT_OK(WriteIndicesUniform(data)); + *out_dict = arr_first.dictionary(); + } + } + return Status::OK(); + } + + OwnedRefNoGIL dictionary_; + bool ordered_; + bool needs_copy_; +}; + +class ExtensionWriter : public PandasWriter { + public: + using PandasWriter::PandasWriter; + + Status Allocate() override { + // no-op + return Status::OK(); + } + + Status TransferSingle(std::shared_ptr data, PyObject* py_ref) override { + PyAcquireGIL lock; + PyObject* py_array; + py_array = wrap_chunked_array(data); + py_array_.reset(py_array); + + return Status::OK(); + } + + Status CopyInto(std::shared_ptr data, int64_t rel_placement) override { + return TransferSingle(data, nullptr); + } + + Status GetDataFrameResult(PyObject** out) override { + PyAcquireGIL lock; + PyObject* result = PyDict_New(); + RETURN_IF_PYERROR(); + + PyDict_SetItemString(result, "py_array", py_array_.obj()); + PyDict_SetItemString(result, "placement", placement_arr_.obj()); + *out = result; + return Status::OK(); + } + + Status GetSeriesResult(PyObject** out) override { + *out = py_array_.detach(); + return Status::OK(); + } + + protected: + OwnedRefNoGIL py_array_; +}; + +Status MakeWriter(const PandasOptions& options, PandasWriter::type writer_type, + const DataType& type, int64_t num_rows, int num_columns, + std::shared_ptr* writer) { +#define BLOCK_CASE(NAME, TYPE) \ + case PandasWriter::NAME: \ + *writer = std::make_shared(options, num_rows, num_columns); \ + break; + +#define CATEGORICAL_CASE(TYPE) \ + case TYPE::type_id: \ + *writer = std::make_shared>(options, num_rows); \ + break; + +#define TZ_CASE(NAME, TYPE) \ + case PandasWriter::NAME: { \ + const auto& ts_type = checked_cast(type); \ + *writer = std::make_shared(options, ts_type.timezone(), num_rows); \ + } break; + + switch (writer_type) { + case PandasWriter::CATEGORICAL: { + const auto& index_type = *checked_cast(type).index_type(); + switch (index_type.id()) { + CATEGORICAL_CASE(Int8Type); + CATEGORICAL_CASE(Int16Type); + CATEGORICAL_CASE(Int32Type); + CATEGORICAL_CASE(Int64Type); + CATEGORICAL_CASE(UInt8Type); + CATEGORICAL_CASE(UInt16Type); + CATEGORICAL_CASE(UInt32Type); + case Type::UINT64: + return Status::TypeError( + "Converting UInt64 dictionary indices to pandas is not supported."); + default: + // Unreachable + ARROW_DCHECK(false); + break; + } + } break; + case PandasWriter::EXTENSION: + *writer = std::make_shared(options, num_rows, num_columns); + break; + BLOCK_CASE(OBJECT, ObjectWriter); + BLOCK_CASE(UINT8, UInt8Writer); + BLOCK_CASE(INT8, Int8Writer); + BLOCK_CASE(UINT16, UInt16Writer); + BLOCK_CASE(INT16, Int16Writer); + BLOCK_CASE(UINT32, UInt32Writer); + BLOCK_CASE(INT32, Int32Writer); + BLOCK_CASE(UINT64, UInt64Writer); + BLOCK_CASE(INT64, Int64Writer); + BLOCK_CASE(HALF_FLOAT, Float16Writer); + BLOCK_CASE(FLOAT, Float32Writer); + BLOCK_CASE(DOUBLE, Float64Writer); + BLOCK_CASE(BOOL, BoolWriter); + BLOCK_CASE(DATETIME_DAY, DatetimeDayWriter); + BLOCK_CASE(DATETIME_SECOND, DatetimeSecondWriter); + BLOCK_CASE(DATETIME_MILLI, DatetimeMilliWriter); + BLOCK_CASE(DATETIME_MICRO, DatetimeMicroWriter); + BLOCK_CASE(DATETIME_NANO, DatetimeNanoWriter); + BLOCK_CASE(TIMEDELTA_SECOND, TimedeltaSecondWriter); + BLOCK_CASE(TIMEDELTA_MILLI, TimedeltaMilliWriter); + BLOCK_CASE(TIMEDELTA_MICRO, TimedeltaMicroWriter); + BLOCK_CASE(TIMEDELTA_NANO, TimedeltaNanoWriter); + TZ_CASE(DATETIME_SECOND_TZ, DatetimeSecondTZWriter); + TZ_CASE(DATETIME_MILLI_TZ, DatetimeMilliTZWriter); + TZ_CASE(DATETIME_MICRO_TZ, DatetimeMicroTZWriter); + TZ_CASE(DATETIME_NANO_TZ, DatetimeNanoTZWriter); + default: + return Status::NotImplemented("Unsupported block type"); + } + +#undef BLOCK_CASE +#undef CATEGORICAL_CASE + + return Status::OK(); +} + +static Status GetPandasWriterType(const ChunkedArray& data, const PandasOptions& options, + PandasWriter::type* output_type) { +#define INTEGER_CASE(NAME) \ + *output_type = \ + data.null_count() > 0 \ + ? options.integer_object_nulls ? PandasWriter::OBJECT : PandasWriter::DOUBLE \ + : PandasWriter::NAME; \ + break; + + switch (data.type()->id()) { + case Type::BOOL: + *output_type = data.null_count() > 0 ? PandasWriter::OBJECT : PandasWriter::BOOL; + break; + case Type::UINT8: + INTEGER_CASE(UINT8); + case Type::INT8: + INTEGER_CASE(INT8); + case Type::UINT16: + INTEGER_CASE(UINT16); + case Type::INT16: + INTEGER_CASE(INT16); + case Type::UINT32: + INTEGER_CASE(UINT32); + case Type::INT32: + INTEGER_CASE(INT32); + case Type::UINT64: + INTEGER_CASE(UINT64); + case Type::INT64: + INTEGER_CASE(INT64); + case Type::HALF_FLOAT: + *output_type = PandasWriter::HALF_FLOAT; + break; + case Type::FLOAT: + *output_type = PandasWriter::FLOAT; + break; + case Type::DOUBLE: + *output_type = PandasWriter::DOUBLE; + break; + case Type::STRING: // fall through + case Type::LARGE_STRING: // fall through + case Type::STRING_VIEW: // fall through + case Type::BINARY: // fall through + case Type::LARGE_BINARY: + case Type::BINARY_VIEW: + case Type::NA: // fall through + case Type::FIXED_SIZE_BINARY: // fall through + case Type::STRUCT: // fall through + case Type::TIME32: // fall through + case Type::TIME64: // fall through + case Type::DECIMAL32: // fall through + case Type::DECIMAL64: // fall through + case Type::DECIMAL128: // fall through + case Type::DECIMAL256: // fall through + case Type::INTERVAL_MONTH_DAY_NANO: // fall through + *output_type = PandasWriter::OBJECT; + break; + case Type::DATE32: + if (options.date_as_object) { + *output_type = PandasWriter::OBJECT; + } else if (options.coerce_temporal_nanoseconds) { + *output_type = PandasWriter::DATETIME_NANO; + } else if (options.to_numpy) { + // Numpy supports Day, but Pandas does not + *output_type = PandasWriter::DATETIME_DAY; + } else { + *output_type = PandasWriter::DATETIME_MILLI; + } + break; + case Type::DATE64: + if (options.date_as_object) { + *output_type = PandasWriter::OBJECT; + } else if (options.coerce_temporal_nanoseconds) { + *output_type = PandasWriter::DATETIME_NANO; + } else { + *output_type = PandasWriter::DATETIME_MILLI; + } + break; + case Type::TIMESTAMP: { + const auto& ts_type = checked_cast(*data.type()); + if (options.timestamp_as_object && ts_type.unit() != TimeUnit::NANO) { + // Nanoseconds are never out of bounds for pandas, so in that case + // we don't convert to object + *output_type = PandasWriter::OBJECT; + } else if (options.coerce_temporal_nanoseconds) { + if (!ts_type.timezone().empty()) { + *output_type = PandasWriter::DATETIME_NANO_TZ; + } else { + *output_type = PandasWriter::DATETIME_NANO; + } + } else { + if (!ts_type.timezone().empty()) { + switch (ts_type.unit()) { + case TimeUnit::SECOND: + *output_type = PandasWriter::DATETIME_SECOND_TZ; + break; + case TimeUnit::MILLI: + *output_type = PandasWriter::DATETIME_MILLI_TZ; + break; + case TimeUnit::MICRO: + *output_type = PandasWriter::DATETIME_MICRO_TZ; + break; + case TimeUnit::NANO: + *output_type = PandasWriter::DATETIME_NANO_TZ; + break; + } + } else { + switch (ts_type.unit()) { + case TimeUnit::SECOND: + *output_type = PandasWriter::DATETIME_SECOND; + break; + case TimeUnit::MILLI: + *output_type = PandasWriter::DATETIME_MILLI; + break; + case TimeUnit::MICRO: + *output_type = PandasWriter::DATETIME_MICRO; + break; + case TimeUnit::NANO: + *output_type = PandasWriter::DATETIME_NANO; + break; + } + } + } + } break; + case Type::DURATION: { + const auto& dur_type = checked_cast(*data.type()); + if (options.coerce_temporal_nanoseconds) { + *output_type = PandasWriter::TIMEDELTA_NANO; + } else { + switch (dur_type.unit()) { + case TimeUnit::SECOND: + *output_type = PandasWriter::TIMEDELTA_SECOND; + break; + case TimeUnit::MILLI: + *output_type = PandasWriter::TIMEDELTA_MILLI; + break; + case TimeUnit::MICRO: + *output_type = PandasWriter::TIMEDELTA_MICRO; + break; + case TimeUnit::NANO: + *output_type = PandasWriter::TIMEDELTA_NANO; + break; + } + } + } break; + case Type::FIXED_SIZE_LIST: + case Type::LIST: + case Type::LARGE_LIST: + case Type::LIST_VIEW: + case Type::LARGE_LIST_VIEW: + case Type::MAP: { + auto list_type = std::static_pointer_cast(data.type()); + if (!ListTypeSupported(*list_type->value_type())) { + return Status::NotImplemented("Not implemented type for Arrow list to pandas: ", + list_type->value_type()->ToString()); + } + *output_type = PandasWriter::OBJECT; + } break; + case Type::DICTIONARY: + *output_type = PandasWriter::CATEGORICAL; + break; + case Type::EXTENSION: + *output_type = PandasWriter::EXTENSION; + break; + default: + return Status::NotImplemented( + "No known equivalent Pandas block for Arrow data of type ", + data.type()->ToString(), " is known."); + } + return Status::OK(); +} + +// Construct the exact pandas "BlockManager" memory layout +// +// * For each column determine the correct output pandas type +// * Allocate 2D blocks (ncols x nrows) for each distinct data type in output +// * Allocate block placement arrays +// * Write Arrow columns out into each slice of memory; populate block +// * placement arrays as we go +class PandasBlockCreator { + public: + using WriterMap = std::unordered_map>; + + explicit PandasBlockCreator(const PandasOptions& options, FieldVector fields, + ChunkedArrayVector arrays) + : options_(options), fields_(std::move(fields)), arrays_(std::move(arrays)) { + num_columns_ = static_cast(arrays_.size()); + if (num_columns_ > 0) { + num_rows_ = arrays_[0]->length(); + } + column_block_placement_.resize(num_columns_); + } + virtual ~PandasBlockCreator() = default; + + virtual Status Convert(PyObject** out) = 0; + + Status AppendBlocks(const WriterMap& blocks, PyObject* list) { + for (const auto& it : blocks) { + PyObject* item; + RETURN_NOT_OK(it.second->GetDataFrameResult(&item)); + if (PyList_Append(list, item) < 0) { + RETURN_IF_PYERROR(); + } + + // ARROW-1017; PyList_Append increments object refcount + Py_DECREF(item); + } + return Status::OK(); + } + + protected: + PandasOptions options_; + + FieldVector fields_; + ChunkedArrayVector arrays_; + int num_columns_; + int64_t num_rows_; + + // column num -> relative placement within internal block + std::vector column_block_placement_; +}; + +// Helper function for extension chunked arrays +// Constructing a storage chunked array of an extension chunked array +std::shared_ptr GetStorageChunkedArray(std::shared_ptr arr) { + auto value_type = checked_cast(*arr->type()).storage_type(); + ArrayVector storage_arrays; + for (int c = 0; c < arr->num_chunks(); c++) { + const auto& arr_ext = checked_cast(*arr->chunk(c)); + storage_arrays.emplace_back(arr_ext.storage()); + } + return std::make_shared(std::move(storage_arrays), value_type); +}; + +// Helper function to decode RunEndEncodedArray +Result> GetDecodedChunkedArray( + std::shared_ptr arr) { + ARROW_ASSIGN_OR_RAISE(Datum decoded, compute::RunEndDecode(arr)); + ARROW_DCHECK(decoded.is_chunked_array()); + return decoded.chunked_array(); +}; + +class ConsolidatedBlockCreator : public PandasBlockCreator { + public: + using PandasBlockCreator::PandasBlockCreator; + + Status Convert(PyObject** out) override { + column_types_.resize(num_columns_); + RETURN_NOT_OK(CreateBlocks()); + RETURN_NOT_OK(WriteTableToBlocks()); + PyAcquireGIL lock; + + PyObject* result = PyList_New(0); + RETURN_IF_PYERROR(); + + RETURN_NOT_OK(AppendBlocks(blocks_, result)); + RETURN_NOT_OK(AppendBlocks(singleton_blocks_, result)); + + *out = result; + return Status::OK(); + } + + Status GetBlockType(int column_index, PandasWriter::type* out) { + if (options_.IsExtensionColumn(fields_[column_index]->name())) { + *out = PandasWriter::EXTENSION; + return Status::OK(); + } else { + // In case of an extension array default to the storage type + if (arrays_[column_index]->type()->id() == Type::EXTENSION) { + arrays_[column_index] = GetStorageChunkedArray(arrays_[column_index]); + } + // In case of a RunEndEncodedArray default to the values type + else if (arrays_[column_index]->type()->id() == Type::RUN_END_ENCODED) { + ARROW_ASSIGN_OR_RAISE(arrays_[column_index], + GetDecodedChunkedArray(arrays_[column_index])); + } + return GetPandasWriterType(*arrays_[column_index], options_, out); + } + } + + Status CreateBlocks() { + for (int i = 0; i < num_columns_; ++i) { + const DataType& type = *arrays_[i]->type(); + PandasWriter::type output_type; + RETURN_NOT_OK(GetBlockType(i, &output_type)); + + int block_placement = 0; + std::shared_ptr writer; + if (output_type == PandasWriter::CATEGORICAL || + output_type == PandasWriter::DATETIME_SECOND_TZ || + output_type == PandasWriter::DATETIME_MILLI_TZ || + output_type == PandasWriter::DATETIME_MICRO_TZ || + output_type == PandasWriter::DATETIME_NANO_TZ || + output_type == PandasWriter::EXTENSION) { + RETURN_NOT_OK(MakeWriter(options_, output_type, type, num_rows_, + /*num_columns=*/1, &writer)); + singleton_blocks_[i] = writer; + } else { + auto it = block_sizes_.find(output_type); + if (it != block_sizes_.end()) { + block_placement = it->second; + // Increment count + ++it->second; + } else { + // Add key to map + block_sizes_[output_type] = 1; + } + } + column_types_[i] = output_type; + column_block_placement_[i] = block_placement; + } + + // Create normal non-categorical blocks + for (const auto& it : this->block_sizes_) { + PandasWriter::type output_type = static_cast(it.first); + std::shared_ptr block; + RETURN_NOT_OK(MakeWriter(this->options_, output_type, /*unused*/ *null(), num_rows_, + it.second, &block)); + this->blocks_[output_type] = block; + } + return Status::OK(); + } + + Status GetWriter(int i, std::shared_ptr* block) { + PandasWriter::type output_type = this->column_types_[i]; + switch (output_type) { + case PandasWriter::CATEGORICAL: + case PandasWriter::DATETIME_SECOND_TZ: + case PandasWriter::DATETIME_MILLI_TZ: + case PandasWriter::DATETIME_MICRO_TZ: + case PandasWriter::DATETIME_NANO_TZ: + case PandasWriter::EXTENSION: { + auto it = this->singleton_blocks_.find(i); + if (it == this->singleton_blocks_.end()) { + return Status::KeyError("No block allocated"); + } + *block = it->second; + } break; + default: + auto it = this->blocks_.find(output_type); + if (it == this->blocks_.end()) { + return Status::KeyError("No block allocated"); + } + *block = it->second; + break; + } + return Status::OK(); + } + + Status WriteTableToBlocks() { + auto WriteColumn = [this](int i) { + std::shared_ptr block; + RETURN_NOT_OK(this->GetWriter(i, &block)); + // ARROW-3789 Use std::move on the array to permit self-destructing + return block->Write(std::move(arrays_[i]), i, this->column_block_placement_[i]); + }; + + return OptionalParallelFor(options_.use_threads, num_columns_, WriteColumn); + } + + private: + // column num -> block type id + std::vector column_types_; + + // block type -> type count + std::unordered_map block_sizes_; + std::unordered_map block_types_; + + // block type -> block + WriterMap blocks_; + + WriterMap singleton_blocks_; +}; + +/// \brief Create blocks for pandas.DataFrame block manager using one block per +/// column strategy. This permits some zero-copy optimizations as well as the +/// ability for the table to "self-destruct" if selected by the user. +class SplitBlockCreator : public PandasBlockCreator { + public: + using PandasBlockCreator::PandasBlockCreator; + + Status GetWriter(int i, std::shared_ptr* writer) { + PandasWriter::type output_type = PandasWriter::OBJECT; + const DataType& type = *arrays_[i]->type(); + if (options_.IsExtensionColumn(fields_[i]->name())) { + output_type = PandasWriter::EXTENSION; + } else { + // Null count needed to determine output type + RETURN_NOT_OK(GetPandasWriterType(*arrays_[i], options_, &output_type)); + } + return MakeWriter(this->options_, output_type, type, num_rows_, 1, writer); + } + + Status Convert(PyObject** out) override { + PyAcquireGIL lock; + + PyObject* result = PyList_New(0); + RETURN_IF_PYERROR(); + + for (int i = 0; i < num_columns_; ++i) { + std::shared_ptr writer; + RETURN_NOT_OK(GetWriter(i, &writer)); + // ARROW-3789 Use std::move on the array to permit self-destructing + RETURN_NOT_OK(writer->Write(std::move(arrays_[i]), i, /*rel_placement=*/0)); + + PyObject* item; + RETURN_NOT_OK(writer->GetDataFrameResult(&item)); + if (PyList_Append(result, item) < 0) { + RETURN_IF_PYERROR(); + } + // PyList_Append increments object refcount + Py_DECREF(item); + } + + *out = result; + return Status::OK(); + } + + private: + std::vector> writers_; +}; + +Status ConvertCategoricals(const PandasOptions& options, ChunkedArrayVector* arrays, + FieldVector* fields) { + std::vector columns_to_encode; + + // For Categorical conversions + auto EncodeColumn = [&](int j) { + int i = columns_to_encode[j]; + if (options.zero_copy_only) { + return Status::Invalid("Need to dictionary encode a column, but ", + "only zero-copy conversions allowed"); + } + compute::ExecContext ctx(options.pool); + ARROW_ASSIGN_OR_RAISE( + Datum out, DictionaryEncode((*arrays)[i], + compute::DictionaryEncodeOptions::Defaults(), &ctx)); + (*arrays)[i] = out.chunked_array(); + (*fields)[i] = (*fields)[i]->WithType((*arrays)[i]->type()); + return Status::OK(); + }; + + if (options.HasCategoricalColumns()) { + for (int i = 0; i < static_cast(arrays->size()); i++) { + if ((*arrays)[i]->type()->id() != Type::DICTIONARY && + options.IsCategoricalColumn((*fields)[i]->name())) { + columns_to_encode.push_back(i); + } + } + } + if (options.strings_to_categorical) { + for (int i = 0; i < static_cast(arrays->size()); i++) { + if (is_base_binary_like((*arrays)[i]->type()->id()) || + is_binary_view_like((*arrays)[i]->type()->id())) { + columns_to_encode.push_back(i); + } + } + } + return OptionalParallelFor(options.use_threads, + static_cast(columns_to_encode.size()), EncodeColumn); +} + +} // namespace + +Status ConvertArrayToPandas(const PandasOptions& options, std::shared_ptr arr, + PyObject* py_ref, PyObject** out) { + return ConvertChunkedArrayToPandas( + options, std::make_shared(std::move(arr)), py_ref, out); +} + +Status ConvertChunkedArrayToPandas(const PandasOptions& options, + std::shared_ptr arr, PyObject* py_ref, + PyObject** out) { + if (options.decode_dictionaries && arr->type()->id() == Type::DICTIONARY) { + // XXX we should return an error as below if options.zero_copy_only + // is true, but that would break compatibility with existing tests. + const auto& dense_type = + checked_cast(*arr->type()).value_type(); + RETURN_NOT_OK(DecodeDictionaries(options.pool, dense_type, &arr)); + ARROW_DCHECK_NE(arr->type()->id(), Type::DICTIONARY); + + // The original Python DictionaryArray won't own the memory anymore + // as we actually built a new array when we decoded the DictionaryArray + // thus let the final resulting numpy array own the memory through a Capsule + py_ref = nullptr; + } + + if (options.strings_to_categorical && (is_base_binary_like(arr->type()->id()) || + is_binary_view_like(arr->type()->id()))) { + if (options.zero_copy_only) { + return Status::Invalid("Need to dictionary encode a column, but ", + "only zero-copy conversions allowed"); + } + compute::ExecContext ctx(options.pool); + ARROW_ASSIGN_OR_RAISE( + Datum out, + DictionaryEncode(arr, compute::DictionaryEncodeOptions::Defaults(), &ctx)); + arr = out.chunked_array(); + } + + PandasOptions modified_options = options; + modified_options.strings_to_categorical = false; + + // ARROW-7596: We permit the hybrid Series/DataFrame code path to do zero copy + // optimizations that we do not allow in the default case when converting + // Table->DataFrame + modified_options.allow_zero_copy_blocks = true; + + // In case of an extension array default to the storage type + if (arr->type()->id() == Type::EXTENSION) { + arr = GetStorageChunkedArray(arr); + } + // In case of a RunEndEncodedArray decode the array + else if (arr->type()->id() == Type::RUN_END_ENCODED) { + if (options.zero_copy_only) { + return Status::Invalid("Need to dencode a RunEndEncodedArray, but ", + "only zero-copy conversions allowed"); + } + ARROW_ASSIGN_OR_RAISE(arr, GetDecodedChunkedArray(arr)); + + // Because we built a new array when we decoded the RunEndEncodedArray + // the final resulting numpy array should own the memory through a Capsule + py_ref = nullptr; + } + + PandasWriter::type output_type; + RETURN_NOT_OK(GetPandasWriterType(*arr, modified_options, &output_type)); + if (options.decode_dictionaries) { + ARROW_DCHECK_NE(output_type, PandasWriter::CATEGORICAL); + } + + std::shared_ptr writer; + RETURN_NOT_OK(MakeWriter(modified_options, output_type, *arr->type(), arr->length(), + /*num_columns=*/1, &writer)); + RETURN_NOT_OK(writer->TransferSingle(std::move(arr), py_ref)); + return writer->GetSeriesResult(out); +} + +Status ConvertTableToPandas(const PandasOptions& options, std::shared_ptr table, + PyObject** out) { + ChunkedArrayVector arrays = table->columns(); + FieldVector fields = table->fields(); + + // ARROW-3789: allow "self-destructing" by releasing references to columns as + // we convert them to pandas + table = nullptr; + + RETURN_NOT_OK(ConvertCategoricals(options, &arrays, &fields)); + + PandasOptions modified_options = options; + modified_options.strings_to_categorical = false; + modified_options.categorical_columns.reset(); + + if (options.split_blocks) { + modified_options.allow_zero_copy_blocks = true; + SplitBlockCreator helper(modified_options, std::move(fields), std::move(arrays)); + return helper.Convert(out); + } else { + ConsolidatedBlockCreator helper(modified_options, std::move(fields), + std::move(arrays)); + return helper.Convert(out); + } +} + +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.h new file mode 100644 index 0000000000000000000000000000000000000000..b4e91e6cf5ab4469236c7eec3c8174d4b0d47427 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_pandas.h @@ -0,0 +1,165 @@ +// 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. + +// Functions for converting between pandas's NumPy-based data representation +// and Arrow data structures + +#pragma once + +#include "arrow/python/platform.h" + +#include +#include +#include + +#include "arrow/memory_pool.h" +#include "arrow/python/visibility.h" + +namespace arrow { + +class Array; +class ChunkedArray; +class Column; +class DataType; +class MemoryPool; +class Status; +class Table; + +namespace py { + +enum class MapConversionType { + DEFAULT, // convert arrow maps to assoc lists (list of kev-value tuples) in Pandas + LOSSY, // report warnings when lossiness is encountered due to duplicate keys + STRICT_, // raise a Python exception when lossiness is encountered due to duplicate + // keys +}; + +struct PandasOptions { + bool HasCategoricalColumns() const { + return categorical_columns && !categorical_columns->empty(); + } + + bool IsCategoricalColumn(const std::string& name) const { + return categorical_columns && categorical_columns->count(name); + } + + bool HasExtensionColumns() const { + return extension_columns && !extension_columns->empty(); + } + + bool IsExtensionColumn(const std::string& name) const { + return extension_columns && extension_columns->count(name); + } + + /// arrow::MemoryPool to use for memory allocations + MemoryPool* pool = default_memory_pool(); + + /// If true, we will convert all string columns to categoricals + bool strings_to_categorical = false; + bool zero_copy_only = false; + bool integer_object_nulls = false; + bool date_as_object = false; + bool timestamp_as_object = false; + bool use_threads = false; + + /// Coerce all date and timestamp to datetime64[ns] + bool coerce_temporal_nanoseconds = false; + + /// Used to maintain backwards compatibility for + /// timezone bugs (see ARROW-9528). Should be removed + /// after Arrow 2.0 release. + bool ignore_timezone = false; + + /// \brief If true, do not create duplicate PyObject versions of equal + /// objects. This only applies to immutable objects like strings or datetime + /// objects + bool deduplicate_objects = false; + + /// \brief For certain data types, a cast is needed in order to store the + /// data in a pandas DataFrame or Series (e.g. timestamps are always stored + /// as nanoseconds in pandas). This option controls whether it is a safe + /// cast or not. + bool safe_cast = true; + + /// \brief If true, create one block per column rather than consolidated + /// blocks (1 per data type). Do zero-copy wrapping when there are no + /// nulls. pandas currently will consolidate the blocks on its own, causing + /// increased memory use, so keep this in mind if you are working on a + /// memory-constrained situation. + bool split_blocks = false; + + /// \brief If true, allow non-writable zero-copy views to be created for + /// single column blocks. This option is also used to provide zero copy for + /// Series data + bool allow_zero_copy_blocks = false; + + /// \brief If true, attempt to deallocate buffers in passed Arrow object if + /// it is the only remaining shared_ptr copy of it. See ARROW-3789 for + /// original context for this feature. Only currently implemented for Table + /// conversions + bool self_destruct = false; + + /// \brief The default behavior (DEFAULT), is to convert Arrow Map arrays to + /// Python association lists (list-of-tuples) in the same order as the Arrow + /// Map, as in [(key1, value1), (key2, value2), ...] + /// If LOSSY or STRICT, convert Arrow Map arrays to native Python dicts. + /// This can change the ordering of (key, value) pairs, and will deduplicate + /// multiple keys, resulting in a possible loss of data. + /// If 'lossy', this key deduplication results in a warning printed + /// when detected. If 'strict', this instead results in an exception + /// being raised when detected. + MapConversionType maps_as_pydicts = MapConversionType::DEFAULT; + + // Used internally for nested arrays. + bool decode_dictionaries = false; + + // Columns that should be casted to categorical + // + // This is wrapped in a shared_ptr because this struct is copied internally for + // each column or nested field (see GH-47861). + std::shared_ptr> categorical_columns; + + // Columns that should be passed through to be converted to + // ExtensionArray/Block + std::shared_ptr> extension_columns; + + // Used internally to decipher between to_numpy() and to_pandas() when + // the expected output differs + bool to_numpy = false; +}; + +ARROW_PYTHON_EXPORT +Status ConvertArrayToPandas(const PandasOptions& options, std::shared_ptr arr, + PyObject* py_ref, PyObject** out); + +ARROW_PYTHON_EXPORT +Status ConvertChunkedArrayToPandas(const PandasOptions& options, + std::shared_ptr col, PyObject* py_ref, + PyObject** out); + +// Convert a whole table as efficiently as possible to a pandas.DataFrame. +// +// The returned Python object is a list of tuples consisting of the exact 2D +// BlockManager structure of the pandas.DataFrame used as of pandas 0.19.x. +// +// tuple item: (indices: ndarray[int32], block: ndarray[TYPE, ndim=2]) +ARROW_PYTHON_EXPORT +Status ConvertTableToPandas(const PandasOptions& options, std::shared_ptr
table, + PyObject** out); + +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_python_internal.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_python_internal.h new file mode 100644 index 0000000000000000000000000000000000000000..514cda320012316b1f9bc04a76c45159dc5bd181 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/arrow_to_python_internal.h @@ -0,0 +1,49 @@ +// 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. + +#pragma once + +#include "arrow/array.h" +#include "arrow/python/platform.h" + +namespace arrow { +namespace py { +namespace internal { +// TODO(ARROW-12976): See if we can refactor Pandas ObjectWriter logic +// to the .cc file and move this there as well if we can. + +// Converts array to a sequency of python objects. +template +inline Status WriteArrayObjects(const ArrayType& arr, WriteValue&& write_func, + Assigner out_values) { + // TODO(ARROW-12976): Use visitor here? + const bool has_nulls = arr.null_count() > 0; + for (int64_t i = 0; i < arr.length(); ++i) { + if (has_nulls && arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values = Py_None; + } else { + RETURN_NOT_OK(write_func(arr.GetView(i), out_values)); + } + ++out_values; + } + return Status::OK(); +} + +} // namespace internal +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/async.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/async.h new file mode 100644 index 0000000000000000000000000000000000000000..1568d21938e6e79e724d957120e68a7576ba9c2a --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/async.h @@ -0,0 +1,60 @@ +// 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. + +#pragma once + +#include + +#include "arrow/python/common.h" +#include "arrow/status.h" +#include "arrow/util/future.h" + +namespace arrow::py { + +/// \brief Bind a Python callback to an arrow::Future. +/// +/// If the Future finishes successfully, py_wrapper is called with its +/// result value and should return a PyObject*. If py_wrapper is successful, +/// py_cb is called with its return value. +/// +/// If either the Future or py_wrapper fails, py_cb is called with the +/// associated Python exception. +/// +/// \param future The future to bind to. +/// \param py_cb The Python callback function. Will be passed the result of +/// py_wrapper, or a Python exception if the future failed or one was +/// raised by py_wrapper. +/// \param py_wrapper A function (likely defined in Cython) to convert the C++ +/// result of the future to a Python object. +template +void BindFuture(Future future, PyObject* py_cb, PyWrapper py_wrapper) { + Py_INCREF(py_cb); + OwnedRefNoGIL cb_ref(py_cb); + + auto future_cb = [cb_ref = std::move(cb_ref), + py_wrapper = std::move(py_wrapper)](Result result) { + SafeCallIntoPythonVoid([&]() { + OwnedRef py_value_or_exc{WrapResult(std::move(result), std::move(py_wrapper))}; + Py_XDECREF( + PyObject_CallFunctionObjArgs(cb_ref.obj(), py_value_or_exc.obj(), NULLPTR)); + ARROW_WARN_NOT_OK(CheckPyError(), "Internal error in async call"); + }); + }; + future.AddCallback(std::move(future_cb)); +} + +} // namespace arrow::py diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.cc b/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.cc new file mode 100644 index 0000000000000000000000000000000000000000..6dcc959ed221247eb93a80179e61a1f40a726e29 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.cc @@ -0,0 +1,38 @@ +// 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. + +#include "arrow/python/benchmark.h" +#include "arrow/python/helpers.h" + +namespace arrow { +namespace py { +namespace benchmark { + +void Benchmark_PandasObjectIsNull(PyObject* list) { + if (!PyList_CheckExact(list)) { + PyErr_SetString(PyExc_TypeError, "expected a list"); + return; + } + Py_ssize_t i, n = PyList_GET_SIZE(list); + for (i = 0; i < n; i++) { + internal::PandasObjectIsNull(PyList_GET_ITEM(list, i)); + } +} + +} // namespace benchmark +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.h new file mode 100644 index 0000000000000000000000000000000000000000..8060dd33722a08eb0935687ea5cb306dbd38a9f0 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/benchmark.h @@ -0,0 +1,36 @@ +// 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. + +#pragma once + +#include "arrow/python/platform.h" + +#include "arrow/python/visibility.h" + +namespace arrow { +namespace py { +namespace benchmark { + +// Micro-benchmark routines for use from ASV + +// Run PandasObjectIsNull() once over every object in *list* +ARROW_PYTHON_EXPORT +void Benchmark_PandasObjectIsNull(PyObject* list); + +} // namespace benchmark +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.cc b/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.cc new file mode 100644 index 0000000000000000000000000000000000000000..a6d963b9b4e43ace79d686f9739baf002bf1a3a7 --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.cc @@ -0,0 +1,246 @@ +// 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. + +#include "arrow/python/common.h" + +#include +#include +#include +#include + +#include "arrow/memory_pool.h" +#include "arrow/status.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/logging.h" + +#include "arrow/python/helpers.h" + +namespace arrow { + +using internal::checked_cast; + +namespace py { + +static std::mutex memory_pool_mutex; +static MemoryPool* default_python_pool = nullptr; + +void set_default_memory_pool(MemoryPool* pool) { + std::lock_guard guard(memory_pool_mutex); + default_python_pool = pool; +} + +MemoryPool* get_memory_pool() { + std::lock_guard guard(memory_pool_mutex); + if (default_python_pool) { + return default_python_pool; + } else { + return default_memory_pool(); + } +} + +// ---------------------------------------------------------------------- +// PythonErrorDetail + +namespace { + +const char kErrorDetailTypeId[] = "arrow::py::PythonErrorDetail"; + +// Try to match the Python exception type with an appropriate Status code +StatusCode MapPyError(PyObject* exc_type) { + StatusCode code; + + if (PyErr_GivenExceptionMatches(exc_type, PyExc_MemoryError)) { + code = StatusCode::OutOfMemory; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_IndexError)) { + code = StatusCode::IndexError; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_KeyError)) { + code = StatusCode::KeyError; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_TypeError)) { + code = StatusCode::TypeError; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_ValueError) || + PyErr_GivenExceptionMatches(exc_type, PyExc_OverflowError)) { + code = StatusCode::Invalid; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_EnvironmentError)) { + code = StatusCode::IOError; + } else if (PyErr_GivenExceptionMatches(exc_type, PyExc_NotImplementedError)) { + code = StatusCode::NotImplemented; + } else { + code = StatusCode::UnknownError; + } + return code; +} + +// PythonErrorDetail indicates a Python exception was raised. +class PythonErrorDetail : public StatusDetail { + public: + const char* type_id() const override { return kErrorDetailTypeId; } + + std::string ToString() const override { + // This is simple enough not to need the GIL + Result result = FormatImpl(); + + if (result.ok()) { + return result.ValueOrDie(); + } else { + // Fallback to just the exception type + const auto ty = reinterpret_cast(exc_type_.obj()); + return std::string("Python exception: ") + ty->tp_name; + } + } + + void RestorePyError() const { + Py_INCREF(exc_type_.obj()); + Py_INCREF(exc_value_.obj()); + Py_INCREF(exc_traceback_.obj()); + PyErr_Restore(exc_type_.obj(), exc_value_.obj(), exc_traceback_.obj()); + } + + PyObject* exc_type() const { return exc_type_.obj(); } + + PyObject* exc_value() const { return exc_value_.obj(); } + + static std::shared_ptr FromPyError() { + PyObject* exc_type = nullptr; + PyObject* exc_value = nullptr; + PyObject* exc_traceback = nullptr; + + PyErr_Fetch(&exc_type, &exc_value, &exc_traceback); + PyErr_NormalizeException(&exc_type, &exc_value, &exc_traceback); + ARROW_CHECK(exc_type) + << "PythonErrorDetail::FromPyError called without a Python error set"; + ARROW_DCHECK(PyType_Check(exc_type)); + ARROW_DCHECK(exc_value); // Ensured by PyErr_NormalizeException, double-check + if (exc_traceback == nullptr) { + // Needed by PyErr_Restore() + Py_INCREF(Py_None); + exc_traceback = Py_None; + } + + std::shared_ptr detail(new PythonErrorDetail); + detail->exc_type_.reset(exc_type); + detail->exc_value_.reset(exc_value); + detail->exc_traceback_.reset(exc_traceback); + return detail; + } + + protected: + Result FormatImpl() const { + PyAcquireGIL lock; + + // Use traceback.format_exception() + OwnedRef traceback_module; + RETURN_NOT_OK(internal::ImportModule("traceback", &traceback_module)); + + OwnedRef fmt_exception; + RETURN_NOT_OK(internal::ImportFromModule(traceback_module.obj(), "format_exception", + &fmt_exception)); + + OwnedRef formatted; + formatted.reset(PyObject_CallFunctionObjArgs(fmt_exception.obj(), exc_type_.obj(), + exc_value_.obj(), exc_traceback_.obj(), + NULL)); + RETURN_IF_PYERROR(); + + std::stringstream ss; + ss << "Python exception: "; + Py_ssize_t num_lines = PySequence_Length(formatted.obj()); + RETURN_IF_PYERROR(); + + for (Py_ssize_t i = 0; i < num_lines; ++i) { + Py_ssize_t line_size; + + PyObject* line = PySequence_GetItem(formatted.obj(), i); + RETURN_IF_PYERROR(); + + const char* data = PyUnicode_AsUTF8AndSize(line, &line_size); + RETURN_IF_PYERROR(); + + ss << std::string_view(data, line_size); + } + return ss.str(); + } + + PythonErrorDetail() = default; + + OwnedRefNoGIL exc_type_, exc_value_, exc_traceback_; +}; + +} // namespace + +// ---------------------------------------------------------------------- +// Python exception <-> Status + +Status ConvertPyError(StatusCode code) { + auto detail = PythonErrorDetail::FromPyError(); + if (code == StatusCode::UnknownError) { + code = MapPyError(detail->exc_type()); + } + + std::string message; + RETURN_NOT_OK(internal::PyObject_StdStringStr(detail->exc_value(), &message)); + return Status(code, message, detail); +} + +bool IsPyError(const Status& status) { + if (status.ok()) { + return false; + } + auto detail = status.detail(); + bool result = detail != nullptr && detail->type_id() == kErrorDetailTypeId; + return result; +} + +void RestorePyError(const Status& status) { + ARROW_CHECK(IsPyError(status)); + const auto& detail = checked_cast(*status.detail()); + detail.RestorePyError(); +} + +// ---------------------------------------------------------------------- +// PyBuffer + +PyBuffer::PyBuffer() : Buffer(nullptr, 0) {} + +Status PyBuffer::Init(PyObject* obj) { + if (!PyObject_GetBuffer(obj, &py_buf_, PyBUF_ANY_CONTIGUOUS)) { + data_ = reinterpret_cast(py_buf_.buf); + ARROW_CHECK_NE(data_, nullptr) << "Null pointer in Py_buffer"; + size_ = py_buf_.len; + capacity_ = py_buf_.len; + is_mutable_ = !py_buf_.readonly; + return Status::OK(); + } else { + return ConvertPyError(StatusCode::Invalid); + } +} + +Result> PyBuffer::FromPyObject(PyObject* obj) { + PyBuffer* buf = new PyBuffer(); + std::shared_ptr res(buf); + RETURN_NOT_OK(buf->Init(obj)); + return res; +} + +PyBuffer::~PyBuffer() { + if (data_ != nullptr) { + PyAcquireGIL lock; + PyBuffer_Release(&py_buf_); + } +} + +} // namespace py +} // namespace arrow diff --git a/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.h b/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.h new file mode 100644 index 0000000000000000000000000000000000000000..affefe2859b65f04f711de1f90d14f642641c5fb --- /dev/null +++ b/lib/python3.12/site-packages/pyarrow/src/arrow/python/common.h @@ -0,0 +1,457 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "arrow/buffer.h" +#include "arrow/python/pyarrow.h" +#include "arrow/python/visibility.h" +#include "arrow/result.h" +#include "arrow/util/macros.h" + +namespace arrow { + +class MemoryPool; +template +class Result; + +namespace py { + +// Convert current Python error to a Status. The Python error state is cleared +// and can be restored with RestorePyError(). +ARROW_PYTHON_EXPORT Status ConvertPyError(StatusCode code = StatusCode::UnknownError); +// Query whether the given Status is a Python error (as wrapped by ConvertPyError()). +ARROW_PYTHON_EXPORT bool IsPyError(const Status& status); +// Restore a Python error wrapped in a Status. +ARROW_PYTHON_EXPORT void RestorePyError(const Status& status); + +// Catch a pending Python exception and return the corresponding Status. +// If no exception is pending, Status::OK() is returned. +inline Status CheckPyError(StatusCode code = StatusCode::UnknownError) { + if (ARROW_PREDICT_TRUE(!PyErr_Occurred())) { + return Status::OK(); + } else { + return ConvertPyError(code); + } +} + +#define RETURN_IF_PYERROR() ARROW_RETURN_NOT_OK(CheckPyError()) + +#define PY_RETURN_IF_ERROR(CODE) ARROW_RETURN_NOT_OK(CheckPyError(CODE)) + +// For Cython, as you can't define template C++ functions in Cython, only use them. +// This function can set a Python exception. It assumes that T has a (cheap) +// default constructor. +template +T GetResultValue(Result result) { + if (ARROW_PREDICT_TRUE(result.ok())) { + return *std::move(result); + } else { + int r = internal::check_status(result.status()); // takes the GIL + assert(r == -1); // should have errored out + ARROW_UNUSED(r); + return {}; + } +} + +/// \brief Wrap a Result and return the corresponding Python object. +/// +/// If the Result is successful, py_wrapper is called with its result value +/// and should return a PyObject*. If py_wrapper is successful (returns +/// a non-NULL value), its return value is returned. +/// +/// If either the Result or py_wrapper fails, the associated Python exception +/// is raised and NULL is returned. +// +/// \param result The Result whose value to wrap in a Python object. +/// \param py_wrapper A function (likely defined in Cython) to convert the C++ +/// value of the Result to a Python object. +/// \return A new Python reference, or NULL if an exception occurred +template +PyObject* WrapResult(Result result, PyWrapper&& py_wrapper) { + static_assert(std::is_same_v()))>, + "PyWrapper argument to WrapResult should return a PyObject* " + "when called with a T*"); + Status st = result.status(); + if (st.ok()) { + PyObject* py_value = py_wrapper(result.MoveValueUnsafe()); + st = CheckPyError(); + if (st.ok()) { + return py_value; + } + Py_XDECREF(py_value); // should be null, but who knows + } + // Status is an error, convert it to an exception. + return internal::convert_status(st); +} + +// A RAII-style helper that ensures the GIL is acquired inside a lexical block. +class ARROW_PYTHON_EXPORT PyAcquireGIL { + public: + PyAcquireGIL() : acquired_gil_(false) { acquire(); } + + ~PyAcquireGIL() { release(); } + + void acquire() { + if (!acquired_gil_) { + state_ = PyGILState_Ensure(); + acquired_gil_ = true; + } + } + + // idempotent + void release() { + if (acquired_gil_) { + PyGILState_Release(state_); + acquired_gil_ = false; + } + } + + private: + bool acquired_gil_; + PyGILState_STATE state_; + ARROW_DISALLOW_COPY_AND_ASSIGN(PyAcquireGIL); +}; + +// A RAII-style helper that releases the GIL until the end of a lexical block +class ARROW_PYTHON_EXPORT PyReleaseGIL { + public: + PyReleaseGIL() : ptr_(PyEval_SaveThread(), &unique_ptr_deleter) {} + + private: + static void unique_ptr_deleter(PyThreadState* state) { + if (state) { + PyEval_RestoreThread(state); + } + } + std::unique_ptr ptr_; +}; + +// A helper to call safely into the Python interpreter from arbitrary C++ code. +// The GIL is acquired, and the current thread's error status is preserved. +template +auto SafeCallIntoPython(Function&& func) -> decltype(func()) { + PyAcquireGIL lock; + PyObject* exc_type; + PyObject* exc_value; + PyObject* exc_traceback; + PyErr_Fetch(&exc_type, &exc_value, &exc_traceback); + auto maybe_status = std::forward(func)(); + // If the return Status is a "Python error", the current Python error status + // describes the error and shouldn't be clobbered. + if (!IsPyError(::arrow::ToStatus(maybe_status)) && exc_type != NULLPTR) { + PyErr_Restore(exc_type, exc_value, exc_traceback); + } + return maybe_status; +} + +template +auto SafeCallIntoPythonVoid(Function&& func) -> decltype(func()) { + PyAcquireGIL lock; + PyObject* exc_type; + PyObject* exc_value; + PyObject* exc_traceback; + PyErr_Fetch(&exc_type, &exc_value, &exc_traceback); + func(); + if (exc_type != NULLPTR) { + PyErr_Restore(exc_type, exc_value, exc_traceback); + } +} + +// A RAII primitive that DECREFs the underlying PyObject* when it +// goes out of scope. +class ARROW_PYTHON_EXPORT OwnedRef { + public: + OwnedRef() : obj_(NULLPTR) {} + OwnedRef(OwnedRef&& other) : OwnedRef(other.detach()) {} + explicit OwnedRef(PyObject* obj) : obj_(obj) {} + + OwnedRef& operator=(OwnedRef&& other) { + obj_ = other.detach(); + return *this; + } + + ~OwnedRef() { + // GH-38626: destructor may be called after the Python interpreter is finalized. + if (Py_IsInitialized()) { + reset(); + } + } + + void reset(PyObject* obj) { + Py_XDECREF(obj_); + obj_ = obj; + } + + void reset() { reset(NULLPTR); } + + PyObject* detach() { + PyObject* result = obj_; + obj_ = NULLPTR; + return result; + } + + PyObject* obj() const { return obj_; } + + PyObject** ref() { return &obj_; } + + operator bool() const { return obj_ != NULLPTR; } + + private: + ARROW_DISALLOW_COPY_AND_ASSIGN(OwnedRef); + + PyObject* obj_; +}; + +// Same as OwnedRef, but ensures the GIL is taken when it goes out of scope. +// This is for situations where the GIL is not always known to be held +// (e.g. if it is released in the middle of a function for performance reasons) +class ARROW_PYTHON_EXPORT OwnedRefNoGIL : public OwnedRef { + public: + OwnedRefNoGIL() : OwnedRef() {} + OwnedRefNoGIL(OwnedRefNoGIL&& other) : OwnedRef(other.detach()) {} + explicit OwnedRefNoGIL(PyObject* obj) : OwnedRef(obj) {} + + ~OwnedRefNoGIL() { + // GH-38626: destructor may be called after the Python interpreter is finalized. + if (Py_IsInitialized() && obj() != NULLPTR) { + PyAcquireGIL lock; + reset(); + } + } +}; + +template