IAP GITLAB

Skip to content
Snippets Groups Projects

Resolve "make friendly example scripts for the release"

Closed Alan Coleman requested to merge 626-make-friendly-example-scripts-for-the-release into master
2 unresolved threads
Compare and
33 files
+ 388
122
Compare changes
  • Side-by-side
  • Inline
Files
33
+ 55
0
"""
Converter for the pyarrow data types to numpy ones
(c) Copyright 2020 CORSIKA Project, corsika-project@lists.kit.edu
This software is distributed under the terms of the GNU General Public
Licence version 3 (GPL Version 3). See file LICENSE for a full version of
the license.
"""
import pyarrow
import numpy as np
def convert_to_numpy(pyarrow_table: pyarrow.lib.Table) -> np.ndarray:
"""
Converts a pyarrow Table to a numpy structured array
Parameters
----------
pyarrow_table: pyarrow.lib.Table
PyArrow table of any dimension to be sliced
Returns
-------
np.ndarray:
converted table with the same column labels and data types
"""
# Type conversions for pyarrow data types to numpy ones
# pyarrow: https://arrow.apache.org/docs/python/data.html
# numpy: https://numpy.org/doc/stable/reference/arrays.dtypes.html#arrays-dtypes-constructing
type_conversions = {
pyarrow.int8(): "int8",
pyarrow.int16(): "int16",
pyarrow.int32(): "int32",
pyarrow.int64(): "int64",
pyarrow.uint8(): "uint8",
pyarrow.uint16(): "uint16",
pyarrow.uint32(): "uint32",
pyarrow.uint64(): "uint64",
pyarrow.float16(): "float16",
pyarrow.float32(): "float32",
pyarrow.float64(): "float64",
}
# Perform type conversion of all fields
column_types = [type_conversions[pyarrow_table[key].type] for key in pyarrow_table.column_names]
# Make an empty array and then fill the values
np_table = np.zeros(pyarrow_table.num_rows, dtype={"names": pyarrow_table.column_names, "formats": column_types})
for key in pyarrow_table.column_names:
np_table[key] = pyarrow_table[key]
return np_table
Loading