=[g0dZddlmZmZddlmZmZmZmZm Z m Z m Z m Z m Z ddlZddlZddlmZddlmZmZmZmZddlmZdd lmZd d lmZd ej>d ej@fdZ!eddZ"e"ddddddZ#de e ej@d e ej@fdZ$deejJdeejJe&e'gdfd dfdZ(GddeZ)dejJd efdZ*d ee&e ej@fd!e e+d"ee&efd#e ed$ee&efd ef d%Z,deejJd&e e e&d!e e+d'e'd(ee&efd)e'd*e'd e ee effd+Z- d0d,ed ed-e ed.e'd ej@f d/Z.y)1z*Utilities for processing spark partitions.) defaultdict namedtuple) AnyCallableDictIteratorListOptionalSequenceTupleUnionN) csr_matrix)DataIterDMatrixQuantileDMatrixXGBModelconcat) ArrayLike) get_loggerseriesreturncT|jd}tj|}|S)zStack a series of arrays.F)copy)to_numpynpstack)rarrays M/var/www/html/bid-api/venv/lib/python3.12/site-packages/xgboost/spark/data.py stack_seriesr"s$ OOO 'E HHUOE LAlias)datalabelweightmarginvalidqidvaluesr&r' baseMarginvalidationIndicatorr*seqc|r t|Sy)z&Concatenate the data if it's not None.Nr)r.s r!concat_or_noner0s c{ r#iteratorappendcdtjdtddffd }d}|D]}|tj|j v}|durtj|j vsJ|rJ|j |tjddf}|j |tjddf}n|d}}||d|||dy)znExtract partitions from pyspark iterator. `append` is a user defined function for accepting new partition.partis_validrNc|tj||tj||tj||tj||tj |y)N)aliasr%r&r'r(r*)r4r5r2s r! make_blobz#cache_partitions..make_blob+sStUZZ*tU[[(+tU\\8,tU\\8,tUYY)r#TF)pd DataFrameboolr7r)columnsloc)r1r2r8has_validationr4trainr)s ` r!cache_partitionsr@%s * ***&*N  !"[[DLL8N T !;;$,,. .. HHd5;;//23EHHT%++.12E5E%   eT "r#ceZdZdZdeeefdeede ddffd Z dee e jdee jfdZd edefd Zd d ZxZS) PartIterz7Iterator for creating Quantile DMatrix from partitions.r% device_idkwargsrNc ^d|_||_||_||_t|dy)NrT) release_data)_iter _device_id_data_kwargssuper__init__)selfr%rCrD __class__s r!rLzPartIter.__init__Hs1 #   d+r#c|sy|jUddl}ddl}|jjj |j|j ||jS||jSNr)rHcudfcupycudaruntime setDevicer:rG)rMr%rQcps r!_fetchzPartIter._fetchRs[ ?? &   GGOO % %doo 6>>$tzz"23 3DJJr# input_datac |jt|jtjk(ry|d|j |jtj|j |jj tjd|j |jj tjd|j |jj tjd|j |jj tjdd|j|xjdz c_y)Nrr%r&r' base_marginr*r) rGlenrIr7r%rWgetr&r'r(r*rJ)rMrXs r!nextz PartIter.nextas ::TZZ 34 4 TZZ 34++djjnnU[[$?@;;tzz~~ellDAB DJJNN5<<$FG DJJNN599d;<  ll   a r#cd|_yrP)rG)rMs r!resetzPartIter.resetos  r#)rN)__name__ __module__ __qualname____doc__rstrr r intrrLr r9r:rWrr_ra __classcell__)rNs@r!rBrBEsyA,dO,08 ,IL, ,  8HR\\$:;  @V   x C r#rBr4cvgdgg}}}d}t|j|j|j|jD]\}}}}|dk(rt |} |} |} n2t |} tj| tj} |} |dk(r| }|| k(sJ|j| |j|dt | z|j| tj|} tj|} tj|}t|| | ft ||fS)Nr)dtype)shape)zipfeatureVectorTypefeatureVectorSizefeatureVectorIndicesfeatureVectorValuesrgr]rarangeint32r2r concatenater)r4csr_indices_listcsr_indptr_listcsr_values_list n_featuresvec_type vec_size_ vec_indices vec_valuesvec_size csr_indices csr_valuescsr_indptr_arrcsr_indices_arrcsr_values_arrs r!)_read_csr_matrix_from_unwrapped_spark_vecrss59;aS"oJ8;   !!   94)[* q=9~H%K#J :H))HBHH=K#J ?!JX%%% ,r2S5EEFz*598XXo.Nnn%56O^^O4N  .9#d)ZAX r#r% dev_ordinalmetarefparamsc|s ttjd|St||fi|}t|fi|d|i}|S)z+Handle empty partition for QuantileDMatrix.rr)rr)rremptyrB)r%rrrritms r!make_qdmrsG rxx/S99 $ ,t ,B.f.#.A Hr# feature_colsuse_qdmrDenable_sparse_data_optimhas_validation_colc ttttddtjdtdt ddffd }dtjdtdt ddffd }d t tttjfd t ttfdtfd } |r|} d vrd d k(sJ|} dtt ttft ttttt fffffd } | \} } |rt!|| t#|| d| }nO|st!|| | }n5|rt!|| t#|| d| }nt!|| | }|r |rt#|| || }n|r | nd}nd}|#|j%|j%k(sJ||fS)a~Create DMatrix from spark data partitions. Parameters ---------- iterator : Pyspark partition iterator. feature_cols: A sequence of feature names, used only when rapids plugin is enabled. dev_ordinal: Device ordinal, used when GPU is enabled. use_qdm : Whether QuantileDMatrix should be used instead of DMatrix. kwargs : Metainfo for DMatrix. enable_sparse_data_optim : Whether sparse data should be unwrapped has_validation: Whether there's validation data. Returns ------- Training DMatrix and an optional validation DMatrix. rr4namer5rNc|tjk(s||jvr|tjk(r|jddkDr|}n;||jddkDr$||}|tjk(rt |}nd}|tjk(r*|(dk(r|jd|jdk(sJ|y|r|j |y|j |yyNrr)r7r%r<rlr"r2)r4rr5r rrx train_data valid_datas r!append_mz0create_dmatrix_from_partitions..append_ms 5:: !5 " ,&,,Q/!3.2<.@d!!!$q(T 5::%(/Euzz!e&7?!&QJ!U[[^333}4 ''.4 ''.5"6r#c6|tjk(s||jvrw|tjk(r3t|}dk(r|jd|jdk(sJ||}|r|j |y|j |yyr)r7r%r<rrlr2)r4rr5r rxrrs r!append_m_sparsez7create_dmatrix_from_partitions..append_m_sparses 5:: !5uzz!A$G?!&QJ!U[[^333T 4 ''.4 ''."6r#r+rDc 4t|dk(r:tdjdtddt j di|St |tj}t |jtjd}t |jtjd}t |jtjd}t |jtjd}td|||||d|S)NrXGBoostPySparkz_Detected an empty partition in the training data. Consider to enable repartition_random_shuffler%rrZr\)r]rwarningrrrr0r7r%r^r&r'r(r*)r+rDr%r&r'r(r*s r!makez,create_dmatrix_from_partitions..makes v;!  ' ( 0 0.  ; 0;F; ;fUZZ01vzz%++t<= 5<< >? 5<< >?VZZ 489 U6v3 RX  r#missinggcfd}i}i}jD]\}}||vr|||<|||<||fS)N)max_binrsilentnthreadenable_categorical)items) non_data_keysnon_data_paramsrkvrDs r! split_paramsz4create_dmatrix_from_partitions..split_paramssR  LLNDAqM!%&"Q # _$$r#)rlistr9r:rfr;rr rndarrayrrr r rgfloatr@rnum_col)r1rrrrDrrrrr append_fnrrrdtraindvalidrxrrs ` ` @@@r!create_dmatrix_from_partitionsrsB/:$.?J.9$.?JJ/r||/3/$/4//</bll/#//$/" T#tBJJ//0 $sCx. W $ # F"vi'8C'??? %%S#XS%UD@P:Q5Q0R RS%* >LD&G9-":{D$O  !'9-j&)  '9-*k4vF9-j&) (0Kvv)F2DT*f-F ~~6>>#3333 6>r#modelr[ strict_shapec|jd}t|||j|j|j|j }|j j|dd||S)z4Predict contributions with data with the full model.N)r[rr feature_typesrTF) pred_contribsvalidate_featuresiteration_ranger)_get_iteration_rangerrn_jobsrr get_boosterpredict)rr%r[rr data_dmatrixs r!rrTsv006O   )) 33 L     & &'! ' r#)NF)/re collectionsrrtypingrrrrr r r r r numpyrpandasr9 scipy.sparserxgboostrrrrxgboost.compatr_typingrutilsrSeriesrr"r$r7r0r:rfr;r@rBrrgrrrr\r#r!rs-0/XXX#@@!rzz 7QR h<9NPUV"**!568BJJ;O#r||$#.6 c47PRV7V.W# #@+x+\(BLL(Z(V  sD$$ %  #   sCx.   '   cN    dr||$d8C=)d#d d cN d # dd 7HW% %&dT(,   )$  ZZ r#