--- /dev/null
+# 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._compute import ( # noqa
+ Function,
+ FunctionOptions,
+ FunctionRegistry,
+ HashAggregateFunction,
+ HashAggregateKernel,
+ Kernel,
+ ScalarAggregateFunction,
+ ScalarAggregateKernel,
+ ScalarFunction,
+ ScalarKernel,
+ VectorFunction,
+ VectorKernel,
+ # Option classes
+ ArraySortOptions,
+ AssumeTimezoneOptions,
+ CastOptions,
+ CountOptions,
+ DayOfWeekOptions,
+ DictionaryEncodeOptions,
+ ElementWiseAggregateOptions,
+ ExtractRegexOptions,
+ FilterOptions,
+ IndexOptions,
+ JoinOptions,
+ MakeStructOptions,
+ MatchSubstringOptions,
+ ModeOptions,
+ NullOptions,
+ PadOptions,
+ PartitionNthOptions,
+ QuantileOptions,
+ ReplaceSliceOptions,
+ ReplaceSubstringOptions,
+ RoundOptions,
+ RoundToMultipleOptions,
+ ScalarAggregateOptions,
+ SelectKOptions,
+ SetLookupOptions,
+ SliceOptions,
+ SortOptions,
+ SplitOptions,
+ SplitPatternOptions,
+ StrftimeOptions,
+ StrptimeOptions,
+ TakeOptions,
+ TDigestOptions,
+ TrimOptions,
+ VarianceOptions,
+ WeekOptions,
+ # Functions
+ call_function,
+ function_registry,
+ get_function,
+ list_functions,
+)
+
+import inspect
+from textwrap import dedent
+import warnings
+
+import pyarrow as pa
+
+
+def _get_arg_names(func):
+ return func._doc.arg_names
+
+
+def _decorate_compute_function(wrapper, exposed_name, func, option_class):
+ # Decorate the given compute function wrapper with useful metadata
+ # and documentation.
+ wrapper.__arrow_compute_function__ = dict(name=func.name,
+ arity=func.arity)
+ wrapper.__name__ = exposed_name
+ wrapper.__qualname__ = exposed_name
+
+ doc_pieces = []
+
+ cpp_doc = func._doc
+ summary = cpp_doc.summary
+ if not summary:
+ arg_str = "arguments" if func.arity > 1 else "argument"
+ summary = ("Call compute function {!r} with the given {}"
+ .format(func.name, arg_str))
+
+ description = cpp_doc.description
+ arg_names = _get_arg_names(func)
+
+ doc_pieces.append("""\
+ {}.
+
+ """.format(summary))
+
+ if description:
+ doc_pieces.append("{}\n\n".format(description))
+
+ doc_pieces.append("""\
+ Parameters
+ ----------
+ """)
+
+ for arg_name in arg_names:
+ if func.kind in ('vector', 'scalar_aggregate'):
+ arg_type = 'Array-like'
+ else:
+ arg_type = 'Array-like or scalar-like'
+ doc_pieces.append("""\
+ {} : {}
+ Argument to compute function
+ """.format(arg_name, arg_type))
+
+ doc_pieces.append("""\
+ memory_pool : pyarrow.MemoryPool, optional
+ If not passed, will allocate memory from the default memory pool.
+ """)
+ if option_class is not None:
+ doc_pieces.append("""\
+ options : pyarrow.compute.{0}, optional
+ Parameters altering compute function semantics.
+ """.format(option_class.__name__))
+ options_sig = inspect.signature(option_class)
+ for p in options_sig.parameters.values():
+ doc_pieces.append("""\
+ {0} : optional
+ Parameter for {1} constructor. Either `options`
+ or `{0}` can be passed, but not both at the same time.
+ """.format(p.name, option_class.__name__))
+
+ wrapper.__doc__ = "".join(dedent(s) for s in doc_pieces)
+ return wrapper
+
+
+def _get_options_class(func):
+ class_name = func._doc.options_class
+ if not class_name:
+ return None
+ try:
+ return globals()[class_name]
+ except KeyError:
+ warnings.warn("Python binding for {} not exposed"
+ .format(class_name), RuntimeWarning)
+ return None
+
+
+def _handle_options(name, option_class, options, kwargs):
+ if kwargs:
+ if options is None:
+ return option_class(**kwargs)
+ raise TypeError(
+ "Function {!r} called with both an 'options' argument "
+ "and additional named arguments"
+ .format(name))
+
+ if options is not None:
+ if isinstance(options, dict):
+ return option_class(**options)
+ elif isinstance(options, option_class):
+ return options
+ raise TypeError(
+ "Function {!r} expected a {} parameter, got {}"
+ .format(name, option_class, type(options)))
+
+ return options
+
+
+def _make_generic_wrapper(func_name, func, option_class):
+ if option_class is None:
+ def wrapper(*args, memory_pool=None):
+ return func.call(args, None, memory_pool)
+ else:
+ def wrapper(*args, memory_pool=None, options=None, **kwargs):
+ options = _handle_options(func_name, option_class, options,
+ kwargs)
+ return func.call(args, options, memory_pool)
+ return wrapper
+
+
+def _make_signature(arg_names, var_arg_names, option_class):
+ from inspect import Parameter
+ params = []
+ for name in arg_names:
+ params.append(Parameter(name, Parameter.POSITIONAL_OR_KEYWORD))
+ for name in var_arg_names:
+ params.append(Parameter(name, Parameter.VAR_POSITIONAL))
+ params.append(Parameter("memory_pool", Parameter.KEYWORD_ONLY,
+ default=None))
+ if option_class is not None:
+ params.append(Parameter("options", Parameter.KEYWORD_ONLY,
+ default=None))
+ options_sig = inspect.signature(option_class)
+ for p in options_sig.parameters.values():
+ # XXX for now, our generic wrappers don't allow positional
+ # option arguments
+ params.append(p.replace(kind=Parameter.KEYWORD_ONLY))
+ return inspect.Signature(params)
+
+
+def _wrap_function(name, func):
+ option_class = _get_options_class(func)
+ arg_names = _get_arg_names(func)
+ has_vararg = arg_names and arg_names[-1].startswith('*')
+ if has_vararg:
+ var_arg_names = [arg_names.pop().lstrip('*')]
+ else:
+ var_arg_names = []
+
+ wrapper = _make_generic_wrapper(name, func, option_class)
+ wrapper.__signature__ = _make_signature(arg_names, var_arg_names,
+ option_class)
+ return _decorate_compute_function(wrapper, name, func, option_class)
+
+
+def _make_global_functions():
+ """
+ Make global functions wrapping each compute function.
+
+ Note that some of the automatically-generated wrappers may be overriden
+ by custom versions below.
+ """
+ g = globals()
+ reg = function_registry()
+
+ # Avoid clashes with Python keywords
+ rewrites = {'and': 'and_',
+ 'or': 'or_'}
+
+ for cpp_name in reg.list_functions():
+ name = rewrites.get(cpp_name, cpp_name)
+ func = reg.get_function(cpp_name)
+ assert name not in g, name
+ g[cpp_name] = g[name] = _wrap_function(name, func)
+
+
+_make_global_functions()
+
+
+def cast(arr, target_type, safe=True):
+ """
+ Cast array values to another data type. Can also be invoked as an array
+ instance method.
+
+ Parameters
+ ----------
+ arr : Array or ChunkedArray
+ target_type : DataType or type string alias
+ Type to cast to
+ safe : bool, default True
+ Check for overflows or other unsafe conversions
+
+ Examples
+ --------
+ >>> from datetime import datetime
+ >>> import pyarrow as pa
+ >>> arr = pa.array([datetime(2010, 1, 1), datetime(2015, 1, 1)])
+ >>> arr.type
+ TimestampType(timestamp[us])
+
+ You can use ``pyarrow.DataType`` objects to specify the target type:
+
+ >>> cast(arr, pa.timestamp('ms'))
+ <pyarrow.lib.TimestampArray object at 0x7fe93c0f6910>
+ [
+ 2010-01-01 00:00:00.000,
+ 2015-01-01 00:00:00.000
+ ]
+
+ >>> cast(arr, pa.timestamp('ms')).type
+ TimestampType(timestamp[ms])
+
+ Alternatively, it is also supported to use the string aliases for these
+ types:
+
+ >>> arr.cast('timestamp[ms]')
+ <pyarrow.lib.TimestampArray object at 0x10420eb88>
+ [
+ 1262304000000,
+ 1420070400000
+ ]
+ >>> arr.cast('timestamp[ms]').type
+ TimestampType(timestamp[ms])
+
+ Returns
+ -------
+ casted : Array
+ """
+ if target_type is None:
+ raise ValueError("Cast target type must not be None")
+ if safe:
+ options = CastOptions.safe(target_type)
+ else:
+ options = CastOptions.unsafe(target_type)
+ return call_function("cast", [arr], options)
+
+
+def count_substring(array, pattern, *, ignore_case=False):
+ """
+ Count the occurrences of substring *pattern* in each value of a
+ string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ pattern to search for exact matches
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("count_substring", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def count_substring_regex(array, pattern, *, ignore_case=False):
+ """
+ Count the non-overlapping matches of regex *pattern* in each value
+ of a string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ pattern to search for exact matches
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("count_substring_regex", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def find_substring(array, pattern, *, ignore_case=False):
+ """
+ Find the index of the first occurrence of substring *pattern* in each
+ value of a string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ pattern to search for exact matches
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("find_substring", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def find_substring_regex(array, pattern, *, ignore_case=False):
+ """
+ Find the index of the first match of regex *pattern* in each
+ value of a string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ regex pattern to search for
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("find_substring_regex", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def match_like(array, pattern, *, ignore_case=False):
+ """
+ Test if the SQL-style LIKE pattern *pattern* matches a value of a
+ string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ SQL-style LIKE pattern. '%' will match any number of
+ characters, '_' will match exactly one character, and all
+ other characters match themselves. To match a literal percent
+ sign or underscore, precede the character with a backslash.
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+
+ """
+ return call_function("match_like", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def match_substring(array, pattern, *, ignore_case=False):
+ """
+ Test if substring *pattern* is contained within a value of a string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ pattern to search for exact matches
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("match_substring", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def match_substring_regex(array, pattern, *, ignore_case=False):
+ """
+ Test if regex *pattern* matches at any position a value of a string array.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ pattern : str
+ regex pattern to search
+ ignore_case : bool, default False
+ Ignore case while searching.
+
+ Returns
+ -------
+ result : pyarrow.Array or pyarrow.ChunkedArray
+ """
+ return call_function("match_substring_regex", [array],
+ MatchSubstringOptions(pattern,
+ ignore_case=ignore_case))
+
+
+def mode(array, n=1, *, skip_nulls=True, min_count=0):
+ """
+ Return top-n most common values and number of times they occur in a passed
+ numerical (chunked) array, in descending order of occurrence. If there are
+ multiple values with same count, the smaller one is returned first.
+
+ Parameters
+ ----------
+ array : pyarrow.Array or pyarrow.ChunkedArray
+ n : int, default 1
+ Specify the top-n values.
+ skip_nulls : bool, default True
+ If True, ignore nulls in the input. Else return an empty array
+ if any input is null.
+ min_count : int, default 0
+ If there are fewer than this many values in the input, return
+ an empty array.
+
+ Returns
+ -------
+ An array of <input type "Mode", int64_t "Count"> structs
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pyarrow.compute as pc
+ >>> arr = pa.array([1, 1, 2, 2, 3, 2, 2, 2])
+ >>> modes = pc.mode(arr, 2)
+ >>> modes[0]
+ <pyarrow.StructScalar: {'mode': 2, 'count': 5}>
+ >>> modes[1]
+ <pyarrow.StructScalar: {'mode': 1, 'count': 2}>
+ """
+ options = ModeOptions(n, skip_nulls=skip_nulls, min_count=min_count)
+ return call_function("mode", [array], options)
+
+
+def filter(data, mask, null_selection_behavior='drop'):
+ """
+ Select values (or records) from array- or table-like data given boolean
+ filter, where true values are selected.
+
+ Parameters
+ ----------
+ data : Array, ChunkedArray, RecordBatch, or Table
+ mask : Array, ChunkedArray
+ Must be of boolean type
+ null_selection_behavior : str, default 'drop'
+ Configure the behavior on encountering a null slot in the mask.
+ Allowed values are 'drop' and 'emit_null'.
+
+ - 'drop': nulls will be treated as equivalent to False.
+ - 'emit_null': nulls will result in a null in the output.
+
+ Returns
+ -------
+ result : depends on inputs
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> arr = pa.array(["a", "b", "c", None, "e"])
+ >>> mask = pa.array([True, False, None, False, True])
+ >>> arr.filter(mask)
+ <pyarrow.lib.StringArray object at 0x7fa826df9200>
+ [
+ "a",
+ "e"
+ ]
+ >>> arr.filter(mask, null_selection_behavior='emit_null')
+ <pyarrow.lib.StringArray object at 0x7fa826df9200>
+ [
+ "a",
+ null,
+ "e"
+ ]
+ """
+ options = FilterOptions(null_selection_behavior)
+ return call_function('filter', [data, mask], options)
+
+
+def index(data, value, start=None, end=None, *, memory_pool=None):
+ """
+ Find the index of the first occurrence of a given value.
+
+ Parameters
+ ----------
+ data : Array or ChunkedArray
+ value : Scalar-like object
+ start : int, optional
+ end : int, optional
+ memory_pool : MemoryPool, optional
+ If not passed, will allocate memory from the default memory pool.
+
+ Returns
+ -------
+ index : the index, or -1 if not found
+ """
+ if start is not None:
+ if end is not None:
+ data = data.slice(start, end - start)
+ else:
+ data = data.slice(start)
+ elif end is not None:
+ data = data.slice(0, end)
+
+ if not isinstance(value, pa.Scalar):
+ value = pa.scalar(value, type=data.type)
+ elif data.type != value.type:
+ value = pa.scalar(value.as_py(), type=data.type)
+ options = IndexOptions(value=value)
+ result = call_function('index', [data], options, memory_pool)
+ if start is not None and result.as_py() >= 0:
+ result = pa.scalar(result.as_py() + start, type=pa.int64())
+ return result
+
+
+def take(data, indices, *, boundscheck=True, memory_pool=None):
+ """
+ Select values (or records) from array- or table-like data given integer
+ selection indices.
+
+ The result will be of the same type(s) as the input, with elements taken
+ from the input array (or record batch / table fields) at the given
+ indices. If an index is null then the corresponding value in the output
+ will be null.
+
+ Parameters
+ ----------
+ data : Array, ChunkedArray, RecordBatch, or Table
+ indices : Array, ChunkedArray
+ Must be of integer type
+ boundscheck : boolean, default True
+ Whether to boundscheck the indices. If False and there is an out of
+ bounds index, will likely cause the process to crash.
+ memory_pool : MemoryPool, optional
+ If not passed, will allocate memory from the default memory pool.
+
+ Returns
+ -------
+ result : depends on inputs
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> arr = pa.array(["a", "b", "c", None, "e", "f"])
+ >>> indices = pa.array([0, None, 4, 3])
+ >>> arr.take(indices)
+ <pyarrow.lib.StringArray object at 0x7ffa4fc7d368>
+ [
+ "a",
+ null,
+ "e",
+ null
+ ]
+ """
+ options = TakeOptions(boundscheck=boundscheck)
+ return call_function('take', [data, indices], options, memory_pool)
+
+
+def fill_null(values, fill_value):
+ """
+ Replace each null element in values with fill_value. The fill_value must be
+ the same type as values or able to be implicitly casted to the array's
+ type.
+
+ This is an alias for :func:`coalesce`.
+
+ Parameters
+ ----------
+ values : Array, ChunkedArray, or Scalar-like object
+ Each null element is replaced with the corresponding value
+ from fill_value.
+ fill_value : Array, ChunkedArray, or Scalar-like object
+ If not same type as data will attempt to cast.
+
+ Returns
+ -------
+ result : depends on inputs
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> arr = pa.array([1, 2, None, 3], type=pa.int8())
+ >>> fill_value = pa.scalar(5, type=pa.int8())
+ >>> arr.fill_null(fill_value)
+ pyarrow.lib.Int8Array object at 0x7f95437f01a0>
+ [
+ 1,
+ 2,
+ 5,
+ 3
+ ]
+ """
+ if not isinstance(fill_value, (pa.Array, pa.ChunkedArray, pa.Scalar)):
+ fill_value = pa.scalar(fill_value, type=values.type)
+ elif values.type != fill_value.type:
+ fill_value = pa.scalar(fill_value.as_py(), type=values.type)
+
+ return call_function("coalesce", [values, fill_value])
+
+
+def top_k_unstable(values, k, sort_keys=None, *, memory_pool=None):
+ """
+ Select the indices of the top-k ordered elements from array- or table-like
+ data.
+
+ This is a specialization for :func:`select_k_unstable`. Output is not
+ guaranteed to be stable.
+
+ Parameters
+ ----------
+ values : Array, ChunkedArray, RecordBatch, or Table
+ Data to sort and get top indices from.
+ k : int
+ The number of `k` elements to keep.
+ sort_keys : List-like
+ Column key names to order by when input is table-like data.
+ memory_pool : MemoryPool, optional
+ If not passed, will allocate memory from the default memory pool.
+
+ Returns
+ -------
+ result : Array of indices
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pyarrow.compute as pc
+ >>> arr = pa.array(["a", "b", "c", None, "e", "f"])
+ >>> pc.top_k_unstable(arr, k=3)
+ <pyarrow.lib.UInt64Array object at 0x7fdcb19d7f30>
+ [
+ 5,
+ 4,
+ 2
+ ]
+ """
+ if sort_keys is None:
+ sort_keys = []
+ if isinstance(values, (pa.Array, pa.ChunkedArray)):
+ sort_keys.append(("dummy", "descending"))
+ else:
+ sort_keys = map(lambda key_name: (key_name, "descending"), sort_keys)
+ options = SelectKOptions(k, sort_keys)
+ return call_function("select_k_unstable", [values], options, memory_pool)
+
+
+def bottom_k_unstable(values, k, sort_keys=None, *, memory_pool=None):
+ """
+ Select the indices of the bottom-k ordered elements from
+ array- or table-like data.
+
+ This is a specialization for :func:`select_k_unstable`. Output is not
+ guaranteed to be stable.
+
+ Parameters
+ ----------
+ values : Array, ChunkedArray, RecordBatch, or Table
+ Data to sort and get bottom indices from.
+ k : int
+ The number of `k` elements to keep.
+ sort_keys : List-like
+ Column key names to order by when input is table-like data.
+ memory_pool : MemoryPool, optional
+ If not passed, will allocate memory from the default memory pool.
+
+ Returns
+ -------
+ result : Array of indices
+
+ Examples
+ --------
+ >>> import pyarrow as pa
+ >>> import pyarrow.compute as pc
+ >>> arr = pa.array(["a", "b", "c", None, "e", "f"])
+ >>> pc.bottom_k_unstable(arr, k=3)
+ <pyarrow.lib.UInt64Array object at 0x7fdcb19d7fa0>
+ [
+ 0,
+ 1,
+ 2
+ ]
+ """
+ if sort_keys is None:
+ sort_keys = []
+ if isinstance(values, (pa.Array, pa.ChunkedArray)):
+ sort_keys.append(("dummy", "ascending"))
+ else:
+ sort_keys = map(lambda key_name: (key_name, "ascending"), sort_keys)
+ options = SelectKOptions(k, sort_keys)
+ return call_function("select_k_unstable", [values], options, memory_pool)
projects / ceph.git / blobdiff