Source code for pyarrow.orc

# 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 itertools import count
from numbers import Integral
import warnings

from pyarrow import types
from pyarrow.lib import Schema, Table
import pyarrow._orc as _orc


def _is_map(typ):
    return (types.is_list(typ) and
            types.is_struct(typ.value_type) and
            typ.value_type.num_fields == 2 and
            typ.value_type[0].name == 'key' and
            typ.value_type[1].name == 'value')


def _traverse(typ, counter):
    if isinstance(typ, Schema) or types.is_struct(typ):
        for field in typ:
            path = (field.name,)
            yield path, next(counter)
            for sub, c in _traverse(field.type, counter):
                yield path + sub, c
    elif _is_map(typ):
        yield from _traverse(typ.value_type, counter)
    elif types.is_list(typ):
        # Skip one index for list type, since this can never be selected
        # directly
        next(counter)
        yield from _traverse(typ.value_type, counter)
    elif types.is_union(typ):
        # Union types not supported, just skip the indexes
        for dtype in typ:
            next(counter)
            for sub_c in _traverse(dtype, counter):
                pass


def _schema_to_indices(schema):
    return {'.'.join(i): c for i, c in _traverse(schema, count(1))}


[docs]class ORCFile: """ Reader interface for a single ORC file Parameters ---------- source : str or pyarrow.io.NativeFile Readable source. For passing Python file objects or byte buffers, see pyarrow.io.PythonFileInterface or pyarrow.io.BufferReader. """
[docs] def __init__(self, source): self.reader = _orc.ORCReader() self.reader.open(source) self._column_index_lookup = _schema_to_indices(self.schema)
@property def metadata(self): """The file metadata, as an arrow KeyValueMetadata""" return self.reader.metadata() @property def schema(self): """The file schema, as an arrow schema""" return self.reader.schema() @property def nrows(self): """The number of rows in the file""" return self.reader.nrows() @property def nstripes(self): """The number of stripes in the file""" return self.reader.nstripes() def _select_indices(self, columns=None): if columns is None: return None schema = self.schema indices = [] for col in columns: if isinstance(col, Integral): col = int(col) if 0 <= col < len(schema): col = schema[col].name else: raise ValueError("Column indices must be in 0 <= ind < %d," " got %d" % (len(schema), col)) if col in self._column_index_lookup: indices.append(self._column_index_lookup[col]) else: raise ValueError("Unknown column name %r" % col) return indices
[docs] def read_stripe(self, n, columns=None): """Read a single stripe from the file. Parameters ---------- n : int The stripe index columns : list If not None, only these columns will be read from the stripe. A column name may be a prefix of a nested field, e.g. 'a' will select 'a.b', 'a.c', and 'a.d.e' Returns ------- pyarrow.lib.RecordBatch Content of the stripe as a RecordBatch. """ include_indices = self._select_indices(columns) return self.reader.read_stripe(n, include_indices=include_indices)
[docs] def read(self, columns=None): """Read the whole file. 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' Returns ------- pyarrow.lib.Table Content of the file as a Table. """ include_indices = self._select_indices(columns) return self.reader.read(include_indices=include_indices)
class ORCWriter: """ Writer interface for a single ORC file Parameters ---------- where : str or pyarrow.io.NativeFile Writable target. For passing Python file objects or byte buffers, see pyarrow.io.PythonFileInterface, pyarrow.io.BufferOutputStream or pyarrow.io.FixedSizeBufferWriter. """ def __init__(self, where): self.writer = _orc.ORCWriter() self.writer.open(where) def write(self, table): """ Write the table into an ORC file. The schema of the table must be equal to the schema used when opening the ORC file. Parameters ---------- schema : pyarrow.lib.Table The table to be written into the ORC file """ self.writer.write(table) def close(self): """ Close the ORC file """ self.writer.close() def write_table(table, where): """ Write a table into an ORC file Parameters ---------- table : pyarrow.lib.Table The table to be written into the ORC file where : str or pyarrow.io.NativeFile Writable target. For passing Python file objects or byte buffers, see pyarrow.io.PythonFileInterface, pyarrow.io.BufferOutputStream or pyarrow.io.FixedSizeBufferWriter. """ if isinstance(where, Table): warnings.warn( "The order of the arguments has changed. Pass as " "'write_table(table, where)' instead. The old order will raise " "an error in the future.", FutureWarning, stacklevel=2 ) table, where = where, table writer = ORCWriter(where) writer.write(table) writer.close()