func MakeArrayFromScalar()

in arrow/scalar/scalar.go [777:969]
171 lines of code
31 McCabe index (conditional complexity)

func MakeArrayFromScalar(sc Scalar, length int, mem memory.Allocator) (arrow.Array, error) {
	if !sc.IsValid() {
		return MakeArrayOfNull(sc.DataType(), length, mem), nil
	}

	createOffsets := func(valLength int32) *memory.Buffer {
		buffer := memory.NewResizableBuffer(mem)
		buffer.Resize(arrow.Int32Traits.BytesRequired(length + 1))

		out := arrow.Int32Traits.CastFromBytes(buffer.Bytes())
		for i, offset := 0, int32(0); i < length+1; i, offset = i+1, offset+valLength {
			out[i] = offset
		}
		return buffer
	}

	createBuffer := func(data []byte) *memory.Buffer {
		buffer := memory.NewResizableBuffer(mem)
		buffer.Resize(len(data) * length)

		out := buffer.Bytes()
		copy(out, data)
		for j := len(data); j < len(out); j *= 2 {
			copy(out[j:], out[:j])
		}
		return buffer
	}

	finishFixedWidth := func(data []byte) arrow.ArrayData {
		buffer := createBuffer(data)
		defer buffer.Release()
		return array.NewData(sc.DataType(), length, []*memory.Buffer{nil, buffer}, nil, 0, 0)
	}

	switch s := sc.(type) {
	case *Boolean:
		data := memory.NewResizableBuffer(mem)
		defer data.Release()
		data.Resize(int(bitutil.BytesForBits(int64(length))))
		c := byte(0x00)
		if s.Value {
			c = 0xFF
		}
		memory.Set(data.Bytes(), c)
		return array.NewBoolean(length, data, nil, 0), nil
	case BinaryScalar:
		if s.DataType().ID() == arrow.FIXED_SIZE_BINARY {
			data := finishFixedWidth(s.Data())
			defer data.Release()
			return array.MakeFromData(data), nil
		}

		valuesBuf := createBuffer(s.Data())
		offsetsBuf := createOffsets(int32(len(s.Data())))
		data := array.NewData(sc.DataType(), length, []*memory.Buffer{nil, offsetsBuf, valuesBuf}, nil, 0, 0)
		defer func() {
			valuesBuf.Release()
			offsetsBuf.Release()
			data.Release()
		}()
		return array.MakeFromData(data), nil
	case *Decimal128:
		data := finishFixedWidth(arrow.Decimal128Traits.CastToBytes([]decimal128.Num{s.Value}))
		defer data.Release()
		return array.MakeFromData(data), nil
	case *Decimal256:
		data := finishFixedWidth(arrow.Decimal256Traits.CastToBytes([]decimal256.Num{s.Value}))
		defer data.Release()
		return array.MakeFromData(data), nil
	case PrimitiveScalar:
		data := finishFixedWidth(s.Data())
		defer data.Release()
		return array.MakeFromData(data), nil
	case *List:
		values := make([]arrow.Array, length)
		for i := range values {
			values[i] = s.Value
		}

		valueArray, err := array.Concatenate(values, mem)
		if err != nil {
			return nil, err
		}
		defer valueArray.Release()

		offsetsBuf := createOffsets(int32(s.Value.Len()))
		defer offsetsBuf.Release()
		data := array.NewData(s.DataType(), length, []*memory.Buffer{nil, offsetsBuf}, []arrow.ArrayData{valueArray.Data()}, 0, 0)
		defer data.Release()
		return array.MakeFromData(data), nil
	case *FixedSizeList:
		values := make([]arrow.Array, length)
		for i := range values {
			values[i] = s.Value
		}

		valueArray, err := array.Concatenate(values, mem)
		if err != nil {
			return nil, err
		}
		defer valueArray.Release()

		data := array.NewData(s.DataType(), length, []*memory.Buffer{nil}, []arrow.ArrayData{valueArray.Data()}, 0, 0)
		defer data.Release()
		return array.MakeFromData(data), nil
	case *Struct:
		fields := make([]arrow.ArrayData, 0)
		for _, v := range s.Value {
			arr, err := MakeArrayFromScalar(v, length, mem)
			if err != nil {
				return nil, err
			}
			defer arr.Release()
			fields = append(fields, arr.Data())
		}

		data := array.NewData(s.DataType(), length, []*memory.Buffer{nil}, fields, 0, 0)
		defer data.Release()
		return array.NewStructData(data), nil
	case *Map:
		structArr := s.GetList().(*array.Struct)
		keys := make([]arrow.Array, length)
		values := make([]arrow.Array, length)
		for i := 0; i < length; i++ {
			keys[i] = structArr.Field(0)
			values[i] = structArr.Field(1)
		}

		keyArr, err := array.Concatenate(keys, mem)
		if err != nil {
			return nil, err
		}
		defer keyArr.Release()

		valueArr, err := array.Concatenate(values, mem)
		if err != nil {
			return nil, err
		}
		defer valueArr.Release()

		offsetsBuf := createOffsets(int32(structArr.Len()))
		outStructArr := array.NewData(structArr.DataType(), keyArr.Len(), []*memory.Buffer{nil}, []arrow.ArrayData{keyArr.Data(), valueArr.Data()}, 0, 0)
		data := array.NewData(s.DataType(), length, []*memory.Buffer{nil, offsetsBuf}, []arrow.ArrayData{outStructArr}, 0, 0)
		defer func() {
			offsetsBuf.Release()
			outStructArr.Release()
			data.Release()
		}()
		return array.MakeFromData(data), nil
	case *RunEndEncoded:
		dt := s.DataType().(*arrow.RunEndEncodedType)

		var endBytes []byte
		switch dt.RunEnds().ID() {
		case arrow.INT16:
			if length > math.MaxInt16 {
				return nil, fmt.Errorf("%w: length overflows int16 run ends", arrow.ErrInvalid)
			}

			v := int16(length)
			endBytes = (*[2]byte)(unsafe.Pointer(&v))[:]
		case arrow.INT32:
			if length > math.MaxInt32 {
				return nil, fmt.Errorf("%w: final length overflows int32 run ends", arrow.ErrInvalid)
			}

			v := int32(length)
			endBytes = (*[4]byte)(unsafe.Pointer(&v))[:]
		case arrow.INT64:
			v := int64(length)
			endBytes = (*[8]byte)(unsafe.Pointer(&v))[:]
		}

		endBuf := createBuffer(endBytes)
		defer endBuf.Release()

		valueArr, err := MakeArrayFromScalar(s.Value, 1, mem)
		if err != nil {
			return nil, err
		}
		defer valueArr.Release()

		runEndsData := array.NewData(dt.RunEnds(), 1, []*memory.Buffer{nil, endBuf}, nil, 0, 0)
		defer runEndsData.Release()

		finalData := array.NewData(s.DataType(), length, []*memory.Buffer{nil},
			[]arrow.ArrayData{runEndsData, valueArr.Data()}, 0, 0)
		defer finalData.Release()
		return array.NewRunEndEncodedData(finalData), nil
	default:
		return nil, fmt.Errorf("array from scalar not yet implemented for type %s", sc.DataType())
	}
}