diff --git a/arrow/scalar/nested.go b/arrow/scalar/nested.go index 3a9f7aef..b6adf7f0 100644 --- a/arrow/scalar/nested.go +++ b/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) } diff --git a/arrow/scalar/scalar_test.go b/arrow/scalar/scalar_test.go index bfd00af1..453e115b 100644 --- a/arrow/scalar/scalar_test.go +++ b/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)