fix(scalar): validate dictionary indices before conversion

arrow/scalar/nested.go

@@ -435,43 +435,61 @@ func (s *Dictionary) ValidateFull() (err error) {
 		return nil
 	}
 
-	max := s.Value.Dict.Len() - 1
+	_, err = s.physicalIndex()
+	return
+}
+
+func (s *Dictionary) physicalIndex() (int, error) {
+	dt := s.Type.(*arrow.DictionaryType)
+	length := s.Value.Dict.Len()
+	outOfBounds := func(idx interface{}) error {
+		return fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+	}
+
 	switch idx := s.Value.Index.value().(type) {
 	case int8:
-		if idx < 0 || int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if idx < 0 || int64(idx) >= int64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case uint8:
-		if int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if uint64(idx) >= uint64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case int16:
-		if idx < 0 || int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if idx < 0 || int64(idx) >= int64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case uint16:
-		if int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if uint64(idx) >= uint64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case int32:
-		if idx < 0 || int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if idx < 0 || int64(idx) >= int64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case uint32:
-		if int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if uint64(idx) >= uint64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case int64:
-		if idx < 0 || int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if idx < 0 || idx >= int64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
 	case uint64:
-		if int(idx) > max {
-			err = fmt.Errorf("%s scalar index value out of bounds: %d", s.DataType(), idx)
+		if idx >= uint64(length) {
+			return 0, outOfBounds(idx)
 		}
+		return int(idx), nil
+	default:
+		return 0, fmt.Errorf("unimplemented dictionary type %s", dt.IndexType)
 	}
-
-	return
 }
 
 func (s *Dictionary) String() string {
@@ -493,30 +511,16 @@ func (s *Dictionary) CastTo(arrow.DataType) (Scalar, error) {
 
 func (s *Dictionary) GetEncodedValue() (Scalar, error) {
 	dt := s.Type.(*arrow.DictionaryType)
+	if err := s.Validate(); err != nil {
+		return nil, err
+	}
 	if !s.IsValid() {
 		return MakeNullScalar(dt.ValueType), nil
 	}
 
-	var idxValue int
-	switch dt.IndexType.ID() {
-	case arrow.INT8:
-		idxValue = int(s.Value.Index.value().(int8))
-	case arrow.UINT8:
-		idxValue = int(s.Value.Index.value().(uint8))
-	case arrow.INT16:
-		idxValue = int(s.Value.Index.value().(int16))
-	case arrow.UINT16:
-		idxValue = int(s.Value.Index.value().(uint16))
-	case arrow.INT32:
-		idxValue = int(s.Value.Index.value().(int32))
-	case arrow.UINT32:
-		idxValue = int(s.Value.Index.value().(uint32))
-	case arrow.INT64:
-		idxValue = int(s.Value.Index.value().(int64))
-	case arrow.UINT64:
-		idxValue = int(s.Value.Index.value().(uint64))
-	default:
-		return nil, fmt.Errorf("unimplemented dictionary type %s", dt.IndexType)
+	idxValue, err := s.physicalIndex()
+	if err != nil {
+		return nil, err
 	}
 	return GetScalar(s.Value.Dict, idxValue)
 }

arrow/scalar/scalar_test.go

@@ -1166,6 +1166,92 @@ func TestDictionaryScalarValidateErrors(t *testing.T) {
 	}
 }
 
+func TestDictionaryScalarIndexBounds(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	valueType := arrow.BinaryTypes.String
+	dict, _, _ := array.FromJSON(mem, valueType, strings.NewReader(`["alpha", "beta", "gamma"]`))
+	defer dict.Release()
+	emptyDict, _, _ := array.FromJSON(mem, valueType, strings.NewReader(`[]`))
+	defer emptyDict.Release()
+
+	for _, indexType := range dictIndexTypes {
+		t.Run(fmt.Sprint(indexType), func(t *testing.T) {
+			bits := indexType.(arrow.FixedWidthDataType).BitWidth()
+			newIndex := func(signed int64, unsigned uint64) scalar.Scalar {
+				if arrow.IsUnsignedInteger(indexType.ID()) {
+					idx, err := scalar.MakeUnsignedIntegerScalar(unsigned, bits)
+					require.NoError(t, err)
+					return idx
+				}
+
+				idx, err := scalar.MakeIntegerScalar(signed, bits)
+				require.NoError(t, err)
+				return idx
+			}
+			checkValid := func(signed int64, unsigned uint64) {
+				idx := newIndex(signed, unsigned)
+				value := scalar.NewDictScalar(idx, dict)
+				defer value.Release()
+
+				assert.NoError(t, value.ValidateFull())
+				encoded, err := value.GetEncodedValue()
+				assert.NoError(t, err)
+				assert.NoError(t, encoded.ValidateFull())
+			}
+			checkInvalid := func(signed int64, unsigned uint64) {
+				idx := newIndex(signed, unsigned)
+				value := scalar.NewDictScalar(idx, dict)
+				defer value.Release()
+
+				assert.Error(t, value.ValidateFull())
+				_, err := value.GetEncodedValue()
+				assert.Error(t, err)
+			}
+
+			checkValid(0, 0)
+			checkValid(2, 2)
+			checkInvalid(3, 3)
+
+			if arrow.IsUnsignedInteger(indexType.ID()) {
+				var max uint64
+				switch bits {
+				case 8:
+					max = uint64(^uint8(0))
+				case 16:
+					max = uint64(^uint16(0))
+				case 32:
+					max = uint64(^uint32(0))
+				case 64:
+					max = ^uint64(0)
+				}
+				checkInvalid(0, max)
+			} else {
+				checkInvalid(-1, 0)
+				var max int64
+				switch bits {
+				case 8:
+					max = int64(^uint8(0) >> 1)
+				case 16:
+					max = int64(^uint16(0) >> 1)
+				case 32:
+					max = int64(^uint32(0) >> 1)
+				case 64:
+					max = int64(^uint64(0) >> 1)
+				}
+				checkInvalid(max, 0)
+			}
+
+			empty := scalar.NewDictScalar(newIndex(0, 0), emptyDict)
+			assert.Error(t, empty.ValidateFull())
+			_, err := empty.GetEncodedValue()
+			assert.Error(t, err)
+			empty.Release()
+		})
+	}
+}
+
 func checkGetValidUnionScalar(t *testing.T, arr arrow.Array, idx int, expected, expectedValue scalar.Scalar) {
 	s, err := scalar.GetScalar(arr, idx)
 	assert.NoError(t, err)