&`iNSPddlZddlZddlmZmZmZmZmZmZm Z m Z m Z m Z m Z ddlZddlmZddlmZddlmZddlmZddlmZmZmZmZddlmZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%dd l&m'Z'er dd l(m)Z)dd l*m+Z+e d Z,ed e'Z-GddeZ.GddeZ/dS)N) TYPE_CHECKINGAnyDictIteratorListMappingOptionalSequenceTupleTypeVarUnion) env_integer)TENSOR_COLUMN_NAME) BlockBuilder) SizeEstimator) NULL_SENTINELfind_partition_indexis_nan keys_equal)Block BlockAccessorBlockColumnAccessorBlockExecStatsBlockMetadataWithSchema BlockTypeKeyTypeU)DEFAULT_TARGET_MAX_BLOCK_SIZE)SortKey) AggregateFnT#RAY_DATA_MAX_UNCOMPACTED_SIZE_BYTESc(eZdZdZdeeeejfddfdZ de ddfdZ e de eee fdefd Ze d eedefd Ze de fd Ze defd ZdefdZdefdZdefdZdefdZddZdS)TableBlockBuilderctjt|_g|_d|_d|_t|_d|_ d|_ d|_ ||_ dSNr) collections defaultdictlist_columns_tables_tables_size_cursor_tables_size_bytesr_uncompacted_size _num_rows_num_uncompacted_rows_num_compactions _block_type)self block_types r/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/ray/data/_internal/table_block.py__init__zTableBlockBuilder.__init__7s`#/55 "$ $% #$!.%&" !%itemreturnNct|dr|}n#t|tjr t |i}t|t jjs0td |t||D]}||j vrdg|j z|j |<|j D]7}||}|j ||8|xjdz c_|xj dz c_ ||j|dS)N as_pydictzLReturned elements of an TableBlock must be of type `dict`, got {} (type {}).)hasattrr; isinstancenpndarrayrr'abcr ValueErrorformattyper*r0getappendr/_compact_if_neededr.add)r3r8 column_namevalues r5rHzTableBlockBuilder.addLsP 4 % % .>>##DD bj ) ) .&-D$  788 $$*F4d$<$<    Q QK$-//.2Vd6P-P k*= 5 5KHH[))E M+ & - -e 4 4 4 4 ! ""a'"" !!! ""4(((((r7blockc$t||js*tdt|d|jd|t j|}|j||xj| z c_dS)NzGot a block of type z , expected z[.If you are mapping a function, ensure it returns an object with the expected type. Block: ) r>r2 TypeErrorrDr for_blockr+rFr/num_rows)r3rKaccessors r5 add_blockzTableBlockBuilder.add_blockes%!122 tE{{t?O  !*511 E""" (++---r7columnsctNNotImplementedError)rRs r5_table_from_pydictz$TableBlockBuilder._table_from_pydictq!!r7tablesctrTrU)rYs r5_combine_tablesz!TableBlockBuilder._combine_tablesurXr7ctrTrUr7r5 _empty_tablezTableBlockBuilder._empty_tableyrXr7ctrTrUr]r7r5_concat_would_copyz$TableBlockBuilder._concat_would_copy}rXr7cl|jrdS|ot|jdkS)NTr<)r*r`lenr+r3s r5will_build_yield_copyz'TableBlockBuilder.will_build_yield_copys7 = 4&&((BS->->-BBr7c|jr||jg}ng}||jt |dkr|S||Sr&)r*rWextendr+rbr^r[)r3rYs r5buildzTableBlockBuilder.buildsv = --dm<<=FFF dl### v;;!  $$&& &''// /r7c|jSrT)r/rcs r5rOzTableBlockBuilder.num_rowss ~r7c&|jdkrdS|j|jdD]6}|xjt j|z c_7t|j|_|j|jzSr&) r/r+r,r-rrN size_bytesrbr.r3tables r5get_estimated_memory_usagez,TableBlockBuilder.get_estimated_memory_usages >Q  1\$":"<"<= S SE  # #}'>u'E'E'P'P'R'R R # # ##&t|#4#4 &)?)J)J)L)LLLr7cD|jsJ|jtkrdS||j}||t |_|j|xjdz c_d|_ dS)Nr<r) r*r.rjMAX_UNCOMPACTED_SIZE_BYTESrWrQrclearr1r0)r3rKs r5rGz$TableBlockBuilder._compact_if_neededs}}  ! , , . .1K K K F'' 66 u!.  "%&"""r7)r9N)__name__ __module__ __qualname__r6r dictrr?r@rHrrQ staticmethodrstrrrrWr[r^boolr`rdrgintrOrmrGr]r7r5r$r$6s&&&*)dGRZ78)T))))2 .s .t . . . ."Dd3i$8"U"""\""U """"\""#"""\"""""\"CtCCCC 0u 0 0 0 0#MCMMMM ' ' ' ' ' 'r7r$ceZdZdefdZedZededede j fdZ de fdZ deefd Zd ed ede fd Zde fd ZdeddfdZd7dZedefdZdedddefdZdedddefdZd8dededeefdZdededeefdZdededeefdZdededeefdZdededeefd Z d9deded"eedeefd#Z!d$efd%Z"ddd&e#d'de fd(Z$e% d:d*ee ddd&e#d'd+ede#e d,ff d-Z&d.ee#eddfd/Z'e% d9d*ee d0ee(dee fd1Z)e%d2e d3e(fd4Z*d5e de fd6Z+d!S);TableBlockAccessorrlc||_dSrT_tablerks r5r6zTableBlockAccessor.__init__s  r7c|d|dzS)N_r<r])namecounts r5_munge_conflictz"TableBlockAccessor._munge_conflicts$$$$$r7rowrow_idxr9ctrTrU)rrs r5_build_tensor_rowz$TableBlockAccessor._build_tensor_rowrXr7c.|}|SrT) to_pandas)r3defaults r5 to_defaultzTableBlockAccessor.to_defaults..""r7ctrTrUrcs r5 column_nameszTableBlockAccessor.column_names!!r7rrJctrTrU)r3rrJs r5 fill_columnzTableBlockAccessor.fill_columnrr7c|jSrTr|rcs r5to_blockzTableBlockAccessor.to_blocks {r7accrctrTrU)r3rs r5_zipzTableBlockAccessor._ziprr7otherctj|}t|t|st|tr`t|trKt|j|g\}}tj||Std t|t|| | krGtd | | | |S)Nz#Cannot zip {} with block of type {}z3Cannot zip self (length {}) with block of length {}) rrNr>rDrznormalize_block_typesr}ziprBrCrOr)r3rr self_norm other_norms r5rzTableBlockAccessor.zips:%e,,#tDzz** $ 233  '99  );(P(P[%())% :%.y99==jIII 9@@T DKK <<>>T]]__ , ,ELLMMOOS\\^^  yy~~r7ctrTrUr]r7r5r^zTableBlockAccessor._empty_tablerXr7 n_samplessort_keyrctrTrU)r3rrs r5_samplezTableBlockAccessor._samplerr7c|t|rtd||dkr|St ||}|||S)Nz+Table sort key must be a column name, was: r)callablerVrOr^minr)r3rrks r5samplezTableBlockAccessor.samples  x11 %HhHH  ==??a  $$&& &  4==?? + +||Ax(((r7Fon ignore_nullscltj|j|}||SN)r)r for_columnr}rr3rrrPs r5rzTableBlockAccessor.counts,&1$+b/BB~~<~888r7c||tj|j|}||Sr)_validate_columnrrr}sumrs r5rzTableBlockAccessor.sum@ b!!!&1$+b/BB|||666r7c||tj|j|}||Sr)rrrr}rrs r5rzTableBlockAccessor.minrr7c||tj|j|}||Sr)rrrr}maxrs r5rzTableBlockAccessor.maxrr7c||tj|j|}||Sr)rrrr}meanrs r5rzTableBlockAccessor.mean s@ b!!!&1$+b/BB}},}777r7Nrc||tj|j|}||Sr)rrrr}sum_of_squared_diffs_from_mean)r3rrrrPs r5rz1TableBlockAccessor.sum_of_squared_diffs_from_meansE b!!!&1$+b/BB66L6QQQr7colc|td|d||vr'td|d|dS)Nz-Provided `on` value has to be non-null (got 'z')zReferencing column 'z' not present in the schema: )rBrschema)r3rs r5rz#TableBlockAccessor._validate_columnsm ;TSTTTUU U ))++ + +XsXXXX , +r7aggsr c|dttttt fffd }}|D]H\}}|t|dkr|dfd|D}tt|D]S}tj |} | dkr%|| |||||<Ti} rt|D] \} } | | | < tjt } t||D]I\}}|j}| |dkr|| |}| |xxdz cc<|| |<J|| J|S)aYApplies provided aggregations to groups of rows with the same key. This assumes the block is already sorted by key in ascending order. Args: sort_key: A column name or list of column names. If this is ``None``, place all rows in a single group. aggs: The aggregations to do. Returns: A sorted block of [k, v_1, ..., v_n] columns where k is the groupby key and v_i is the partially combined accumulator for the ith given aggregation. If key is None then the k column is omitted. r9c3Ks&tfVdSdx}}d}d} |t|}|}t ||r>|dz } t|}n#t $rd}YnwxYwt ||>|||fV|}n#t $rYdSwxYw)z/Creates an iterator over zero-copy group views.NrFpublic_row_formatTr<)tupler iter_rowsnextr StopIterationslice)startenditernext_row next_keyskeysr3s r5 iter_groupsz2TableBlockAccessor._aggregate..iter_groups9s3 ggt}}....OEC>>E>::DH '#':: (I$Xd^Y??"q"'+DzzHH,"""'+H!E" %Xd^Y??"$TZZs%;%;;;;;EE$EE s6 4C?BC BCB5C C"!C"r<rc:g|]}|Sr])init).0agg init_valss r5 z1TableBlockAccessor._aggregate..[s%@@@CCHHY//@@@r7) get_columnsrr r rrbuilderrbrangerrNrOaccumulate_blockrr'r(rxrrrHrg)r3rrrr group_keys group_view accumulatorsirPrrgkrr accumulatorrrrs` @@r5 _aggregatezTableBlockAccessor._aggregate&s"#..00 XeHW,=u,D&EF       6,,..&1kmm   "J "I:!##&qM @@@@4@@@L3t99%%  (2:>>$$&&**&*1g&>&>$Q''LO C  z22  EArCFF+C00E$'l$;$; ( ( [x;??//eDkBBDd q 'D KK    }}r7Tblocksfinalizerc6t|}tj}|fdfd}t jd|Dd|idtj|d} t}fd} d} dgt|z} dgt|z} | D]} | rtj t}tt|D]\}||j}||dkr!t|||}||xxd z cc<|| |<| || |<]d } tt|D]3}||| || | || |<4i}rt#|D] \}}|||< t#|| | D]&\}}}|r||||<!|||<'||n#t($rYnwxYw|}|t-j|| fS) aXCombine previously aggregated blocks. This assumes blocks are already sorted by key in ascending order, so we can do merge sort to get all the rows with the same key. Args: blocks: A list of partially combined and sorted blocks. sort_key: The column name of key or None for global aggregation. aggs: The aggregations to do. finalize: Whether to finalize the aggregation. This is used as an optimization for cases where we repeatedly combine partially aggregated groups. Returns: A block of [k, v_1, ..., v_n] columns and its metadata where k is the groupby key and v_i is the corresponding aggregation result for the ith given aggregation. If key is None then the k column is omitted. c6rt|SdS)N)rr)rrs r5_key_fnz>TableBlockAccessor._combine_aggregated_blocks.._key_fns  QtW~~%tr7cL|}td|DS)NcBg|]}|t|rtn|SrT)rr)rvs r5rzVTableBlockAccessor._combine_aggregated_blocks..safe_key_fn..s*PPPA!)vayy)aPPPr7r)rvaluesrs r5 safe_key_fnzBTableBlockAccessor._combine_aggregated_blocks..safe_key_fns3WQZZFPPPPP r7c^g|]*}tj|d+S)Fr)rrNr)rrKs r5rzATableBlockAccessor._combine_aggregated_blocks..sD'..8858QQr7keyNrTc3KtrCV tn#t$rdYdSwxYwtAdSdSrT)rrr)rrrrsr5genz:TableBlockAccessor._combine_aggregated_blocks..gens%WWX%6%6 BB"&"'+DzzHH,"""'+H!EE" %WWX%6%6 BB"""""s2 AAr<F)stats)rzrrrrheapqmergerrNrrbr'r(rxrrrrrrHrrgr from_block)clsrrrrrrrnext_key_columnsrfirstrresolved_agg_namesrrrrrcol_namenext_keyragg_namerretrrrrrs @@@@@r5_combine_aggregated_blocksz-TableBlockAccessor._combine_aggregated_blocksws[:$99&AA&((##%%           {#    )&)44<<>>; : ##DzzH#GH-- #'  " " " " " " " " $vD 1 &*Vc$ii%7"A + 7 < <!&s4yy!1!1 6 6A#'7TableBlockAccessor._find_partitions_sorted..s7    !h A A   r7r)rFr})r3rr partitionsboundslast_idxidxs` ` r5_find_partitions_sortedz*TableBlockAccessor._find_partitions_sorteds       &      C   dk(3,7 8 8 8HH$+hii0111r7target_block_typectjt}|D]m}tj|}t |t s tdt|d|| xxdz cc<nt|dkr*&gt| kr|S.t|ddd\}fd |Dtfd Drtd d Dd S)aNormalize input blocks to the specified `normalize_type`. If the blocks are already all of the same type, returns original blocks. Args: blocks: A list of TableBlocks to be normalized. target_block_type: The type to normalize the blocks to. If None, Ray Data chooses a type to minimize the amount of data conversions. Returns: A list of blocks of the same type. zWBlock type normalization is only supported for TableBlock, but received block of type: .r<Nc|dS)Nr<r])xs r5z:TableBlockAccessor.normalize_block_types..2s adr7T)rreverserc<g|]}|Sr])try_convert_block_type)rrKrrs r5rz.6s7   EJC & &u.? @ @   r7c3`K|](}t|td V)dS)rN)r>rD)rrKresultss r5 z;TableBlockAccessor.normalize_block_types..:s<LL5:eT'!*%5%5666LLLLLLr7zYExpected all blocks to be of the same type after normalization, but got different types: c,g|]}t|Sr])rD)rbs r5rz.=s(B(B(BQa(B(B(Br7zP. Try using blocks of the same type to avoid the issue with block normalization.)r'r(rxrrNr>rzrBrDr4rbr)rsorteditemsany)rrr seen_typesrKblock_accessorrr s` ` @r5rz(TableBlockAccessor.normalize_block_types s",7+B3+G+G  9 9E*4U;;Nn.@AA  B37;;BBB ~0022 3 3 3q 8 3 3 3 3 z??a    %*;)<Z__EVEV@W@W)W)WM  $#)  """N$$$ $ A      NT    LLLLGLLL L L ,(B(B'(B(B(B,,,  r7rKr4c&|tjkr&tj|S|tjkr&tj|Stj|SrT)rARROWrrNto_arrowPANDASrr)rrKr4s r5r z)TableBlockAccessor.try_convert_block_typeCst  ( ( *511::<< < 9+ + + *511;;== = *511<<>> >r7 other_blockct)aCombine this table with another table horizontally (column-wise). This will append the columns. Args: other_block: The table to hstack side-by-side with. Returns: A new table with columns from both tables combined. rU)r3rs r5hstackzTableBlockAccessor.hstackLs "!r7)rrr9r)FrT)T),rqrrrsrr6rurrrxr?r@rrrrrvrrrrrrr^rrrwr rrrrrrrrr r classmethodrrrrr rr]r7r5rzrzsc%%\%"w"""""\"E"d3i""""""C"E""""%" "'""""4"#"""\""" "c"""" ) )y )S ) ) ) )99949HQK99997c77(1+7777 7c77(1+7777 7c77(1+7777 8s8$88A;8888! R R  R Rqk R ! R R R RCO9OE-4HOUOOOOb  xQxQU xQxQM" xQ  xQ u// 0 xQxQxQ[xQts$22655U 5$I.5 e 555[5n?5?i???[? "% "E " " " " " "r7rz)0r'rtypingrrrrrrr r r r r numpyr?ray._private.ray_constantsrray.air.constantsr ray.data._internal.block_builderr!ray.data._internal.size_estimatorrray.data._internal.utilrrrrray.data.blockrrrrrrrrray.data.contextr2ray.data._internal.planner.exchange.sort_task_specrray.data.aggregater r!ror$rzr]r7r5r'sj                           222222000000999999;;;;;;                    ;:::::/JJJJJJ...... GCLL)[)+H r'r'r'r'r' r'r'r'jk"k"k"k"k"k"k"k"k"k"r7