This repository was archived by the owner on Feb 2, 2024. It is now read-only.
Repository navigation
Expand file tree
/
Copy pathpd_dataframe_type.py
More file actions
111 lines (95 loc) · 4.57 KB
/
Copy pathpd_dataframe_type.py
File metadata and controls
111 lines (95 loc) · 4.57 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# *****************************************************************************
# Copyright (c) 2020, Intel Corporation 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.
#
# 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 numba
from numba import types
from numba.core import cgutils
from numba.extending import (models, register_model, make_attribute_wrapper)
from sdc.str_ext import string_type
class DataFrameType(types.Type): # TODO: IterableType over column names
"""Temporary type class for DataFrame objects.
"""
def __init__(self, data=None, index=None, columns=None, has_parent=False, column_loc=None):
self.data = data
if index is None:
index = types.none
self.index = index
self.columns = columns
# keeping whether it is unboxed from Python to enable reflection of new
# columns
self.has_parent = has_parent
self.column_loc = column_loc
super(DataFrameType, self).__init__(
name="dataframe({}, {}, {}, {})".format(data, index, columns, has_parent))
def copy(self, index=None, has_parent=None):
# XXX is copy necessary?
if index is None:
index = types.none if self.index == types.none else self.index.copy()
data = tuple(a.copy() for a in self.data)
if has_parent is None:
has_parent = self.has_parent
return DataFrameType(data, index, self.columns, has_parent)
@property
def key(self):
# needed?
return self.data, self.index, self.columns, self.has_parent
def unify(self, typingctx, other):
if (isinstance(other, DataFrameType)
and len(other.data) == len(self.data)
and other.columns == self.columns
and other.has_parent == self.has_parent):
new_index = types.none
if self.index != types.none and other.index != types.none:
new_index = self.index.unify(typingctx, other.index)
elif other.index != types.none:
new_index = other.index
elif self.index != types.none:
new_index = self.index
data = tuple(a.unify(typingctx, b) for a, b in zip(self.data, other.data))
return DataFrameType(data, new_index, self.columns, self.has_parent)
def is_precise(self):
return all(a.is_precise() for a in self.data) and self.index.is_precise()
@register_model(DataFrameType)
class DataFrameModel(models.StructModel):
def __init__(self, dmm, fe_type):
types_unique = set()
df_types = []
for col_type in fe_type.data:
if col_type in types_unique:
continue
types_unique.add(col_type)
df_types.append(col_type)
members = [
('data', types.Tuple([types.List(typ) for typ in df_types])),
('index', fe_type.index),
('columns', types.List(string_type)),
('parent', types.pyobject),
]
super(DataFrameModel, self).__init__(dmm, fe_type, members)
make_attribute_wrapper(DataFrameType, 'data', '_data')
make_attribute_wrapper(DataFrameType, 'index', '_index')
make_attribute_wrapper(DataFrameType, 'columns', '_columns')
make_attribute_wrapper(DataFrameType, 'unboxed', '_unboxed')
make_attribute_wrapper(DataFrameType, 'parent', '_parent')