add key type checking to Data get methods

Check that requested key is of appropriate arity for RDD keys,
based on a call to rdd.first() to get example key.

Author: industrial-sloth

python/test/test_data.py

@@ -21,6 +21,9 @@ def test_getMissing(self):
     def test_get(self):
         assert_true(array_equal(self.ary2, self.images.get(1)))
 
+        # keys are integers, ask for sequence
+        assert_raises(ValueError, self.images.get, (1, 2))
+
     def test_getMany(self):
         vals = self.images.getMany([0, -1, 1, 0])
         assert_equals(4, len(vals))
@@ -29,6 +32,10 @@ def test_getMany(self):
         assert_true(array_equal(self.ary2, vals[2]))
         assert_true(array_equal(self.ary1, vals[3]))
 
+        # keys are integers, ask for sequences:
+        assert_raises(ValueError, self.images.get, [(0, 0)])
+        assert_raises(ValueError, self.images.get, [0, (0, 0), 1, 0])
+
     def test_getRanges(self):
         vals = self.images.getRange(slice(None))
         assert_equals(2, len(vals))
@@ -55,6 +62,9 @@ def test_getRanges(self):
         vals = self.images.getRange(slice(2, 3))
         assert_equals(0, len(vals))
 
+        # keys are integers, ask for sequence
+        assert_raises(ValueError, self.images.getRange, [slice(1), slice(1)])
+
         # raise exception if 'step' specified:
         assert_raises(ValueError, self.images.getRange, slice(1, 2, 2))
 
@@ -108,12 +118,15 @@ def setUp(self):
         self.series = Data(self.sc.parallelize(self.dataLocal), dtype='float32')
 
     def test_getMissing(self):
-        assert_is_none(self.series.get(-1))
+        assert_is_none(self.series.get((-1, -1)))
 
     def test_get(self):
         expected = self.dataLocal[1][1]
         assert_true(array_equal(expected, self.series.get((0, 1))))
 
+        assert_raises(ValueError, self.series.get, 1)  # keys are sequences, ask for integer
+        assert_raises(ValueError, self.series.get, (1, 2, 3))  # key length mismatch
+
     def test_getMany(self):
         vals = self.series.getMany([(0, 0), (17, 256), (1, 0), (0, 0)])
         assert_equals(4, len(vals))
@@ -122,6 +135,9 @@ def test_getMany(self):
         assert_true(array_equal(self.dataLocal[2][1], vals[2]))
         assert_true(array_equal(self.dataLocal[0][1], vals[3]))
 
+        assert_raises(ValueError, self.series.getMany, [1])  # keys are sequences, ask for integer
+        assert_raises(ValueError, self.series.getMany, [(0, 0), 1, (1, 0), (0, 0)])  # asking for integer again
+
     def test_getRanges(self):
         vals = self.series.getRange([slice(2), slice(2)])
         assert_equals(4, len(vals))
@@ -174,6 +190,12 @@ def test_getRanges(self):
         vals = self.series.getRange([slice(2, 3), slice(None)])
         assert_equals(0, len(vals))
 
+        # keys are sequences, ask for single slice
+        assert_raises(ValueError, self.series.getRange, slice(2, 3))
+
+        # ask for wrong number of slices
+        assert_raises(ValueError, self.series.getRange, [slice(2, 3), slice(2, 3), slice(2, 3)])
+
         # raise exception if 'step' specified:
         assert_raises(ValueError, self.series.getRange, [slice(0, 4, 2), slice(2, 3)])
 

python/thunder/rdds/data.py

@@ -111,13 +111,32 @@ def take(self, *args, **kwargs):
         """
         return self.rdd.take(*args, **kwargs)
 
+    @staticmethod
+    def __getKeyTypeCheck(actualKey, keySpec):
+        if hasattr(actualKey, "__iter__"):
+            try:
+                specLen = len(keySpec) if hasattr(keySpec, "__len__") else \
+                    reduce(lambda x, y: x + y, [1 for item in keySpec], initial=0)
+                if specLen != len(actualKey):
+                    raise ValueError("Length of key specifier '%s' does not match length of first key '%s'" %
+                                     (str(keySpec), str(actualKey)))
+            except TypeError:
+                raise ValueError("Key specifier '%s' appears not to be a sequence type, but actual keys are " %
+                                 str(keySpec) + "sequences (first key: '%s')" % str(actualKey))
+        else:
+            if hasattr(keySpec, "__iter__"):
+                raise ValueError("Key specifier '%s' appears to be a sequence type, " % str(keySpec) +
+                                 "but actual keys are not (first key: '%s')" % str(actualKey))
+
     def get(self, key):
         """Returns a single value matching the passed key, or None if no matching keys found
 
         If multiple records are found with keys matching the passed key, a sequence of all matching
         values will be returned. (This is not expected as a normal occurance, but could happen with
         some user-created rdds.)
         """
+        firstKey = self.first()[0]
+        Data.__getKeyTypeCheck(firstKey, key)
         filteredVals = self.rdd.filter(lambda (k, v): k == key).values().collect()
         if len(filteredVals) == 1:
             return filteredVals[0]
@@ -135,6 +154,9 @@ def getMany(self, keys):
         If multiple values are found, the corresponding sequence element will be a sequence containing all
         matching values.
         """
+        firstKey = self.first()[0]
+        for key in keys:
+            Data.__getKeyTypeCheck(firstKey, key)
         keySet = frozenset(keys)
         filteredRecs = self.rdd.filter(lambda (k, _): k in keySet).collect()
         sortingDict = {}
@@ -197,7 +219,9 @@ def multiSlicesPredicate(kv):
                         return False
             return True
 
-        if not hasattr(sliceOrSlices, '__len__'):
+        firstKey = self.first()[0]
+        Data.__getKeyTypeCheck(firstKey, sliceOrSlices)
+        if not hasattr(sliceOrSlices, '__iter__'):
             # make my func the pFunc; http://en.wikipedia.org/wiki/P._Funk_%28Wants_to_Get_Funked_Up%29
             pFunc = singleSlicePredicate
             if hasattr(sliceOrSlices, 'step') and sliceOrSlices.step is not None: