Skip to content

Commit

Permalink
BUG: DateRange / subclass handling in Index.join. Handle name field i…
Browse files Browse the repository at this point in the history
…n Index.union, GH #500, #501
  • Loading branch information
wesm committed Dec 17, 2011
1 parent 4c1547f commit 2856c64
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 35 deletions.
72 changes: 47 additions & 25 deletions pandas/core/daterange.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,30 @@ def union(self, other):
if not isinstance(other, DateRange) or other.offset != self.offset:
return Index.union(self.view(Index), other)

if self._can_fast_union(other):
return self._fast_union(other)
else:
return Index.union(self, other)

def _wrap_union_result(self, other, result):
name = self.name if self.name == other.name else None
if isinstance(other, DateRange) and self._can_fast_union(other):
result = self._view_like(result)
result.name = name
return result
else:
return Index(result, name=name)

def _wrap_joined_index(self, joined, other):
name = self.name if self.name == other.name else None
if isinstance(other, DateRange) and self._can_fast_union(other):
joined = self._view_like(joined)
joined.name = name
return joined
else:
return Index(joined, name=name)

def _can_fast_union(self, other):
offset = self.offset

# to make our life easier, "sort" the two ranges
Expand All @@ -298,10 +322,30 @@ def union(self, other):
right_start = right[0]

# Only need to "adjoin", not overlap
if (left_end + offset) >= right_start:
return left._fast_union(right)
return (left_end + offset) >= right_start

def _fast_union(self, other):
# to make our life easier, "sort" the two ranges
if self[0] <= other[0]:
left, right = self, other
else:
return Index.union(self, other)
left, right = other, self

left_start, left_end = left[0], left[-1]
right_end = right[-1]

if not _will_use_cache(self.offset):
# concatenate dates
if left_end < right_end:
loc = right.searchsorted(left_end, side='right')
right_chunk = right.values[loc:]
dates = np.concatenate((left.values, right_chunk))
return self._view_like(dates)
else:
return left
else:
return DateRange(left_start, max(left_end, right_end),
offset=left.offset)

def intersection(self, other):
"""
Expand Down Expand Up @@ -335,35 +379,13 @@ def intersection(self, other):
left_chunk = left.values[lslice]
return self._view_like(left_chunk)

def _fast_union(self, other):
left, right = self, other

left_start, left_end = left[0], left[-1]
right_end = right[-1]

if not _will_use_cache(self.offset):
# concatenate dates
if left_end < right_end:
loc = right.searchsorted(left_end, side='right')
right_chunk = right.values[loc:]
dates = np.concatenate((left.values, right_chunk))
return self._view_like(dates)
else:
return left
else:
return DateRange(left_start, max(left_end, right_end),
offset=left.offset)

def _view_like(self, ndarray):
result = ndarray.view(DateRange)
result.offset = self.offset
result.tzinfo = self.tzinfo
result.name = self.name
return result

def _wrap_union_result(self, other, result):
return Index(result)

def tz_normalize(self, tz):
"""
Convert DateRange from one time zone to another (using pytz)
Expand Down
17 changes: 7 additions & 10 deletions pandas/core/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -378,10 +378,11 @@ def union(self, other):
return self._wrap_union_result(other, result)

def _wrap_union_result(self, other, result):
name = self.name if self.name == other.name else None
if type(self) == type(other):
return type(self)(result)
return type(self)(result, name=name)
else:
return Index(result)
return Index(result, name=name)

def intersection(self, other):
"""
Expand Down Expand Up @@ -586,12 +587,6 @@ def join(self, other, how='left', return_indexers=False):
return join_index

def _join_monotonic(self, other, how='left', return_indexers=False):
this_vals = self.values

# if self.dtype != other.dtype:
# other = Index(other, dtype=object)
other_vals = other.values

if how == 'left':
join_index = self
lidx = None
Expand All @@ -601,10 +596,12 @@ def _join_monotonic(self, other, how='left', return_indexers=False):
lidx = self._left_indexer(other, self)
ridx = None
elif how == 'inner':
join_index, lidx, ridx = self._inner_indexer(this_vals, other_vals)
join_index, lidx, ridx = self._inner_indexer(self.values,
other.values)
join_index = self._wrap_joined_index(join_index, other)
elif how == 'outer':
join_index, lidx, ridx = self._outer_indexer(this_vals, other_vals)
join_index, lidx, ridx = self._outer_indexer(self.values,
other.values)
join_index = self._wrap_joined_index(join_index, other)
else: # pragma: no cover
raise Exception('do not recognize join method %s' % how)
Expand Down
10 changes: 10 additions & 0 deletions pandas/tests/test_index.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,16 @@ def test_union(self):
# non-iterable input
self.assertRaises(Exception, first.union, 0.5)

# preserve names
first.name = 'A'
second.name = 'A'
union = first.union(second)
self.assert_(union.name == 'A')

second.name = 'B'
union = first.union(second)
self.assert_(union.name is None)

def test_add(self):
firstCat = self.strIndex + self.dateIndex
secondCat = self.strIndex + self.strIndex
Expand Down

0 comments on commit 2856c64

Please sign in to comment.