#!/usr/bin/env python # # Copyright 2007 Google Inc. # # Licensed 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. # """Abstract Datastore Input Reader implementation for the map_job API.""" import random from google.appengine.api import datastore from google.appengine.ext import key_range from google.appengine.ext.mapreduce import datastore_range_iterators as db_iters from google.appengine.ext.mapreduce import errors from google.appengine.ext.mapreduce import key_ranges from google.appengine.ext.mapreduce import model from google.appengine.ext.mapreduce import namespace_range from google.appengine.ext.mapreduce.api.map_job import input_reader class AbstractDatastoreInputReader(input_reader.InputReader): """Implementation of an abstract base class for a Datastore input reader.""" _BATCH_SIZE = 50 _MAX_SHARD_COUNT = 256 MAX_NAMESPACES_FOR_KEY_SHARD = 10 _APP_PARAM = "_app" NAMESPACE_PARAM = "namespace" ENTITY_KIND_PARAM = "entity_kind" KEYS_ONLY_PARAM = "keys_only" BATCH_SIZE_PARAM = "batch_size" KEY_RANGE_PARAM = "key_range" FILTERS_PARAM = "filters" _KEY_RANGE_ITER_CLS = db_iters.AbstractKeyRangeIterator def __init__(self, iterator): """Create new AbstractDatastoreInputReader object. This is internal constructor. Use split_input to create readers instead. Args: iterator: an iterator that generates objects for this input reader. """ self._iter = iterator def __iter__(self): """Yields whatever the internal iterator yields.""" for o in self._iter: yield o def __str__(self): """Returns the string representation of this InputReader.""" return repr(self._iter) def to_json(self): """Inherit doc.""" return self._iter.to_json() @classmethod def from_json(cls, state): """Inherit doc.""" return cls(db_iters.RangeIteratorFactory.from_json(state)) @classmethod def _get_query_spec(cls, params): """Construct a model.QuerySpec from model.MapperSpec.""" entity_kind = params[cls.ENTITY_KIND_PARAM] filters = params.get(cls.FILTERS_PARAM) app = params.get(cls._APP_PARAM) ns = params.get(cls.NAMESPACE_PARAM) return model.QuerySpec( entity_kind=cls._get_raw_entity_kind(entity_kind), keys_only=bool(params.get(cls.KEYS_ONLY_PARAM, False)), filters=filters, batch_size=int(params.get(cls.BATCH_SIZE_PARAM, cls._BATCH_SIZE)), model_class_path=entity_kind, app=app, ns=ns) @classmethod def split_input(cls, job_config): """Inherit doc.""" shard_count = job_config.shard_count params = job_config.input_reader_params query_spec = cls._get_query_spec(params) namespaces = None if query_spec.ns is not None: k_ranges = cls._to_key_ranges_by_shard( query_spec.app, [query_spec.ns], shard_count, query_spec) else: ns_keys = namespace_range.get_namespace_keys( query_spec.app, cls.MAX_NAMESPACES_FOR_KEY_SHARD+1) if not ns_keys: return elif len(ns_keys) <= cls.MAX_NAMESPACES_FOR_KEY_SHARD: namespaces = [ns_key.name() or "" for ns_key in ns_keys] k_ranges = cls._to_key_ranges_by_shard( query_spec.app, namespaces, shard_count, query_spec) else: ns_ranges = namespace_range.NamespaceRange.split(n=shard_count, contiguous=False, can_query=lambda: True, _app=query_spec.app) k_ranges = [key_ranges.KeyRangesFactory.create_from_ns_range(ns_range) for ns_range in ns_ranges] iters = [db_iters.RangeIteratorFactory.create_key_ranges_iterator( r, query_spec, cls._KEY_RANGE_ITER_CLS) for r in k_ranges] return [cls(i) for i in iters] @classmethod def _to_key_ranges_by_shard(cls, app, namespaces, shard_count, query_spec): """Get a list of key_ranges.KeyRanges objects, one for each shard. This method uses scatter index to split each namespace into pieces and assign those pieces to shards. Args: app: app_id in str. namespaces: a list of namespaces in str. shard_count: number of shards to split. query_spec: model.QuerySpec. Returns: a list of key_ranges.KeyRanges objects. """ key_ranges_by_ns = [] for namespace in namespaces: ranges = cls._split_ns_by_scatter( shard_count, namespace, query_spec.entity_kind, app) random.shuffle(ranges) key_ranges_by_ns.append(ranges) ranges_by_shard = [[] for _ in range(shard_count)] for ranges in key_ranges_by_ns: for i, k_range in enumerate(ranges): if k_range: ranges_by_shard[i].append(k_range) key_ranges_by_shard = [] for ranges in ranges_by_shard: if ranges: key_ranges_by_shard.append(key_ranges.KeyRangesFactory.create_from_list( ranges)) return key_ranges_by_shard @classmethod def _split_ns_by_scatter(cls, shard_count, namespace, raw_entity_kind, app): """Split a namespace by scatter index into key_range.KeyRange. TODO: Power this with key_range.KeyRange.compute_split_points. Args: shard_count: number of shards. namespace: namespace name to split. str. raw_entity_kind: low level datastore API entity kind. app: app id in str. Returns: A list of key_range.KeyRange objects. If there are not enough entities to splits into requested shards, the returned list will contain KeyRanges ordered lexicographically with any Nones appearing at the end. """ if shard_count == 1: return [key_range.KeyRange(namespace=namespace, _app=app)] ds_query = datastore.Query(kind=raw_entity_kind, namespace=namespace, _app=app, keys_only=True) ds_query.Order("__scatter__") oversampling_factor = 32 random_keys = ds_query.Get(shard_count * oversampling_factor) if not random_keys: return ([key_range.KeyRange(namespace=namespace, _app=app)] + [None] * (shard_count - 1)) random_keys.sort() if len(random_keys) >= shard_count: random_keys = cls._choose_split_points(random_keys, shard_count) k_ranges = [] k_ranges.append(key_range.KeyRange( key_start=None, key_end=random_keys[0], direction=key_range.KeyRange.ASC, include_start=False, include_end=False, namespace=namespace, _app=app)) for i in range(0, len(random_keys) - 1): k_ranges.append(key_range.KeyRange( key_start=random_keys[i], key_end=random_keys[i+1], direction=key_range.KeyRange.ASC, include_start=True, include_end=False, namespace=namespace, _app=app)) k_ranges.append(key_range.KeyRange( key_start=random_keys[-1], key_end=None, direction=key_range.KeyRange.ASC, include_start=True, include_end=False, namespace=namespace, _app=app)) if len(k_ranges) < shard_count: k_ranges += [None] * (shard_count - len(k_ranges)) return k_ranges @classmethod def _choose_split_points(cls, sorted_keys, shard_count): """Returns the best split points given a random set of datastore.Keys.""" assert len(sorted_keys) >= shard_count index_stride = len(sorted_keys) / float(shard_count) return [sorted_keys[int(round(index_stride * i))] for i in range(1, shard_count)] @classmethod def validate(cls, job_config): """Inherit docs.""" super(AbstractDatastoreInputReader, cls).validate(job_config) params = job_config.input_reader_params if cls.ENTITY_KIND_PARAM not in params: raise errors.BadReaderParamsError("Missing input reader parameter " "'entity_kind'") if cls.BATCH_SIZE_PARAM in params: try: batch_size = int(params[cls.BATCH_SIZE_PARAM]) if batch_size < 1: raise errors.BadReaderParamsError("Bad batch size: %s" % batch_size) except ValueError, e: raise errors.BadReaderParamsError("Bad batch size: %s" % e) try: bool(params.get(cls.KEYS_ONLY_PARAM, False)) except: raise errors.BadReaderParamsError("keys_only expects a boolean value but " "got %s", params[cls.KEYS_ONLY_PARAM]) if cls.NAMESPACE_PARAM in params: if not isinstance(params[cls.NAMESPACE_PARAM], (str, unicode, type(None))): raise errors.BadReaderParamsError("Expected a single namespace string") if cls.FILTERS_PARAM in params: filters = params[cls.FILTERS_PARAM] if not isinstance(filters, list): raise errors.BadReaderParamsError("Expected list for filters parameter") for f in filters: if not isinstance(f, (tuple, list)): raise errors.BadReaderParamsError("Filter should be a tuple or list: " "%s", f) if len(f) != 3: raise errors.BadReaderParamsError("Filter should be a 3-tuple: %s", f) prop, op, _ = f if not isinstance(prop, basestring): raise errors.BadReaderParamsError("Property should be string: %s", prop) if not isinstance(op, basestring): raise errors.BadReaderParamsError("Operator should be string: %s", op) @classmethod def _get_raw_entity_kind(cls, entity_kind_or_model_classpath): """Returns the entity kind to use with low level datastore calls. Args: entity_kind_or_model_classpath: user specified entity kind or model classpath. Returns: the entity kind in str to use with low level datastore calls. """ return entity_kind_or_model_classpath