infinite.py - tropo - Project Hosting on Google Code

‹r71

r529

Source path: svn/ trunk/ Python/ tr_mapreduce/ infinite.py

#!/usr/local/bin/python2.5

"""The point of this is to offer InfiniteSortedObjectSequence which lets
   you build up an "infinitely large" list, and then sort it.

   The use case is that you first call Append() any number of times
   with a pickelable object, and then you call Sort().

   The implementation of Append() uses an in-memory list that spills over to disk
   when the specified limits are exceeded.

   The implementation of Sort() does an in-memory sort for the easy case, or
   it uses the n-way mergesort from http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/511509
   for case when the data set is too large to fit into memory.

   More advanced options involve how the object are stored to disk (which pickle protocol
   and should the files be compressed).
   """

import itertools
import sys
import os
import pickle_io
import mergesort
import zlib
import gzip
import cPickle
import cStringIO
import types

from tr_persist import block

def _EstimateSize(obj):
  t = type(obj)
  if t == types.TupleType:
    res = 8 # overhead est
    for inner in obj:
      t = type(inner)
      if t == types.StringType:
        res += len(inner)
      elif t == types.IntType:
        res += 4
      elif t == types.LongType:
        res += 8
      else:
        print "what? ", t
        res += 4 * len(inner) # total guess
      return res
  elif t == types.StringType:
    return 8 + len(obj) # assume some overhead
  else:
    return 8 + 4 * len(obj) # total guess



class InfiniteSortedObjectSequence:
  """Used to sort any number of objects.
  We try to store them in memory and when our limit
  is reached we spill over to disk.

  The usage pattern is you call Append() repeatedly
  and then when done call Sort(). If the data is in memory
  then Sort() just sorts it, else it performs a mergesort on
  the objects that were written to disk.

  When this object is destroyed it returns memory and deletes
  temporary disk files.
  """

  def __init__(self,
               limit = 10000,
	       tmp_base = '/tmp',
               maxopen = 100,
	       maxrecords = 128,
	       maxdata = 1024 * 1024 * 10,
               protocol = 2,
               level = 1):
    """limit is the max number of entries to store in memory before flushing to disk
       tmp_base is the directory where objects are written when memory is exceeded
       maxrecords - max records to put in 1 batch when writing a temp file
       maxdata - max bytes to buffer in memory, very rough
    """
    if tmp_base is None:
      tmp_base = '/tmp'
    self.limit = limit
    self.tmp_base = tmp_base
    self.ar = []
    self.files = []
    self.maxopen = maxopen
    self.maxrecords = maxrecords
    self.maxdata = maxdata
    self.protocol = protocol
    self.level = level
    self.running_data = 0



  def Append(self, obj):
    """Add one object to the dataset, and spill to disk if memory limits
    are exceeded.
    """
    self.ar.append(obj)
    if self.limit is None or self.maxdata is not None:
      self.running_data += _EstimateSize(obj)
      if self.running_data >= self.maxdata:
        print "NEW FLUSH ", len(self.ar), self.running_data, self.maxdata
        self._Flush()
    else:
      if len(self.ar) >= self.limit:
        self._Flush()

  def _Flush(self):
    """Write a batch to disk."""
    if len(self.ar) == 0:
      return
    fn = "%s/mr.%d" % (self.tmp_base, len(self.files))
    self.files.append(fn)
    w = block.ObjectBlockWriter(fn, maxrecords = self.maxrecords,
				maxdata = self.maxdata,
                                protocol = self.protocol)

    for ent in sorted(self.ar):
      w.write(ent)
    w = None
    print "INFINITE FILE ", fn, os.stat(fn).st_size
    self.ar = []
    self.running_data = 0

  def Sort(self):
    """Sort the data set. If the data fit in memory then we just call sorted() and return
    the list, while if we had to spill to disk we return a generator that returns all items.
    """
    if len(self.files) == 0: # special case, small data set
      return sorted(self.ar)
    self._Flush()
    # pairs is generator of (obj, file)
    xfiles = self.files
    while len(xfiles) > 0:
      todo = xfiles[0 : self.maxopen]
      rest = xfiles[self.maxopen :]
      print "INFINITE TODO ", len(todo), todo
      print "INFINITE REST ", len(rest), rest
      pairs = mergesort.mergesort([block.ObjectBlockReader(fn) for fn in todo])
      if len(rest) == 0:
        print "INFINITE  A"
        # firsts is generator of objects
        return (p[0] for p in pairs)
      else:
        fn = "%s/mr.%d" % (self.tmp_base, len(self.files))
        print "INFINITE B ", fn
        self.files.append(fn)
        rest.append(fn)
        xfiles = rest
        #self.writer(fn, (p[0] for p in pairs))
	w = block.ObjectBlockWriter(fn, maxrecords = self.maxrecords,
				    maxdata = self.maxdata)
	for p in pairs:
	  w.write(p[0])
        print "FILE ", fn, os.stat(fn).st_size
    print "BUG"
    return None

  def __del__(self):
    """Remove all temporary files on destruction."""
    for fn in self.files:
      os.unlink(fn)

  def stats(self):
    size = sum([os.stat(f).st_size for f in self.files])
    return { 'nfiles' : len(self.files),
             'size' : size}



if __name__ == '__main__':
  import random
  import time

#  print _EstimateSize("yes")
#  print _EstimateSize( ("alpha", "beta"))
#  print _EstimateSize( ("alpha", 3))
#  print _EstimateSize( (4L, "beta"))
#  sys.exit(1)

  random.seed(0)
  if True:
    ins = InfiniteSortedObjectSequence(limit = 97, maxdata = 1024)
    for x in xrange(200):
      ins.Append("%04d" % x)
    for ents in ins.Sort():
      print "ent=", ents
    sys.exit(1)

  t1 = time.time()
#  ins = InfiniteSortedObjectSequence()
  ins = InfiniteSortedObjectSequence(reader = ZipBatchPickleReader,
                                     writer = ZipBatchPickleWriter)
  for x in xrange(1024000/4):
    s = ""
    for i in xrange(10):
      s += chr(random.randrange(97, 97 + 26))
    ins.Append("%06d %s" % (x % 10379, s))
  t2 = time.time()
  print "sorting.."
  ne = 0
  mod = 1
  for ents in ins.Sort():
    ne += 1
    if ne % mod == 0:
      print "sample: %8d: '%s'" % (ne, ents)
      mod *= 2
  stats = ins.stats()
  ins = None
  t3 = time.time()
  print "Ne: ", ne
  print "Add: %.2f" % (t2-t1)
  print "Sort/Print: %.2f" % (t3-t2)
  print "Stats: ", stats