
    sOi                         d dl mZ d dlmZmZ d dlmZmZmZ d dl	Z
d dlmZ d dlmZ d dlmZ d dlmZ d	eez  d
edefdZ G d dee      Zy)    )ABC)countislice)Any	GeneratorIterableN)common_types)Record)ExtendedPointId)Workeriterablesizereturnc              #   ~   K   t        |       }|r,t        t        ||            }t        |      dk(  ry| |r+yyw)zF
    >>> list(iter_batch([1,2,3,4,5], 3))
    [[1, 2, 3], [4, 5]]
    r   N)iterlistr   len)r   r   source_iterbs       f/volume1/homes/robertsu/coba/app/.venv/lib/python3.12/site-packages/qdrant_client/uploader/uploader.py
iter_batchr      s=     
 x.K
T*+q6Q;	 s   8==c                   j   e Zd Zedeeej                  z     dedefd       Z	ede
eej                  f   ej                  z  eej                     z  dee
   dz  dee   dz  dedef
d	       Zedej                  dedee   fd
       Zede
eej                  f   dedee
eee   f      fd       Zy)BaseUploaderrecords
batch_sizer   c              #      K   t        ||      }|D ]g  }g g g }}}|D ]S  }|j                  |j                         |j                  |j                         |j                  |j                         U |||f i y wN)r   appendidvectorpayload)	clsr   r   record_batchesrecord_batch	ids_batchvectors_batchpayload_batchrecords	            r   iterate_records_batchesz$BaseUploader.iterate_records_batches   sz      $GZ8*L68"bm}I&  +$$V]]3$$V^^4 '
 ]M99 +s   A:A<vectorsr!   Nidsc              #     K   |d t               D        }nt        ||      }|d t               D        }nt        ||      }t        |t        j                        r| j                  ||      }nOt        |t              r3t        d |j                         D              r| j                  ||      }nt        ||      }t        |||      E d {    y 7 w)Nc              3       K   | ]  }d   y wr    .0_s     r   	<genexpr>z/BaseUploader.iterate_batches.<locals>.<genexpr>5   s     $;7aT7   c              3       K   | ]  }d   y wr   r.   r/   s     r   r2   z/BaseUploader.iterate_batches.<locals>.<genexpr>:   s     (?w!wr3   c              3   P   K   | ]  }t        |t        j                           y wr   )
isinstancenpndarray)r0   values     r   r2   z/BaseUploader.iterate_batches.<locals>.<genexpr>@   s       /
7GeJubjj)7Gs   $&)r   r   r6   r7   r8   _vector_batches_from_numpydictanyvalues(_vector_batches_from_numpy_named_vectorszip)r"   r*   r!   r+   r   ids_batchespayload_batchesvector_batchess           r   iterate_batcheszBaseUploader.iterate_batches,   s      ;$;57$;K$S*5K?(?uw(?O(*=Ogrzz*,/,J,J7T^,_N&3 /
7>~~7G/
 ,
 !II'S]^N'<N{NODDDs   CCCCc              #   z   K   t        d| j                  d   |      D ]  }| |||z    j                           y w)Nr   )rangeshapetolist)r*   r   is      r   r:   z'BaseUploader._vector_batches_from_numpyI   s<     q'--*J7A!a*n-4466 8s   9;c           	   #   X   K   t        t         j                         D cg c]  }|j                  d    c}            dk(  sJ d       t	        t         j                                     j                  d   } fdt        |      D        }t        ||      E d {    y c c}w 7 
w)Nr      z8Each named vector should have the same number of vectorsc              3      K   | ]7  }j                         D ci c]  }||   |   j                          c} 9 y c c}w wr   )keysrG   )r0   rH   namer*   s      r   r2   zHBaseUploader._vector_batches_from_numpy_named_vectors.<locals>.<genexpr>X   sL      
' :AHT74=#**,,H' Is   AA A)r   setr=   rF   nextr   rE   r   )r*   r   arrnum_vectorsrB   s   `    r   r>   z5BaseUploader._vector_batches_from_numpy_named_vectorsN   s     
 W^^-=>-=cSYYq\-=>?@AE	FE	FE 4 01288;
;'
 nj999 ? 	:s   B*B#A'B*B(B*)__name__
__module____qualname__classmethodr   r
   typesPointStructintr)   r;   str
NumpyArrayVectorStructr   rC   staticmethodfloatr:   r   r>   r.       r   r   r      sN   :&5#4#445: : 
	: :  Ec5+++,u/?/??(5K]K]B^^E $$&E o&-	E
 E 
E E8 7E,<,< 7# 7RZ[`Ra 7 7 :c5+++,::=:	$sDK'(	): :r^   r   )abcr   	itertoolsr   r   typingr   r   r   numpyr7   qdrant_client.conversionsr	   rV   &qdrant_client.conversions.common_typesr
   qdrant_client.http.modelsr    qdrant_client.parallel_processorr   rX   r   r   r.   r^   r   <module>rg      sO     # + +  ; 9 5 3
I- 
S 
X 
B:63 B:r^   