
    bib!                       d dl mZ d dlZd dlZd dlZd dlZd dlZd dlm	Z	 d dl
Z
d dlmZ d dlmZ d dlmZmZ  ej&                  g dg dd      Z ej*                  ed	
      Zej.                  j1                  dg d      d        Zd Zd Zd Zd Zej.                  j1                  dg d      d        Zej.                  j1                  dg d      d        Zej.                  j1                  dddg      d        Z ej.                  j1                  ddejB                  g      d        Z"ej.                  j1                  dg d      d        Z#ej.                  j1                  dg d      d        Z$d Z%d  Z&ej.                  j1                  d!g d      d"        Z'ej.                  j1                  d#g d$      d%        Z(d& Z)ej.                  jU                   e	ejV                         e	d'      k(  d()      d*        Z,y)+    )annotationsN)Version)	assert_eq)tmpdirtmpfileabcd            xyr   npartitionsorient)splitrecordsindexcolumnsvaluesc                   t        d      5 }t        j                  || d       t        j                  || dd      }t        j                  || d      }t        j                  |j                  t        j                  d      ft        |      z  d      |d	<   |j                  j                  dk(  sJ t        ||       d d d        y # 1 sw Y   y xY w)
NjsonFr   linesTr   r    include_path_column/categorydtypepath)r   dfto_jsondd	read_jsonpdSeriesreplaceosseplenr'   r&   r   r   factual	actual_pds       \/home/cdr/jupyterlab/.venv/lib/python3.12/site-packages/dask/dataframe/io/tests/test_json.pytest_read_json_with_path_columnr7      s    	 %A


1V5
1aeQUVLL6?	 IIYYrvvs#%I6j
	& {{  J...&)$% % %s   B6CCc                    t        d      5 } t        j                  | dd       t        j                  t
        d      5  t        j                  | ddd       d d d        d d d        y # 1 sw Y   xY w# 1 sw Y   y xY w)	Nr   r   Fr   zFiles already containmatchr   r!   )r   r(   r)   pytestraises
ValueErrorr*   r+   r3   s    r6   7test_read_json_path_column_with_duplicate_name_is_errorr?   $   sv    	 TA


1Ye
4]]:-DE 	TLL9EsS	TT T	T 	TT Ts#   4A7 A+A7+A4	0A77B c                 T   d} d }t        d      5 }t        j                  |dd       t        j                  |dd| |      }t        j                  |dd      }t        j                   ||      ft        |      z  d	      || <   t        ||       d d d        y # 1 sw Y   y xY w)
N	filenamesc                     y)Nz	asdf.json r   s    r6   path_converterz:test_read_json_with_path_converter.<locals>.path_converter.   s        r   r   Fr   )r   r    r"   rE   r$   r%   )	r   r(   r)   r*   r+   r,   r-   r1   r   )path_column_namerE   r3   r4   r5   s        r6   "test_read_json_with_path_converterrH   +   s    " 
 %A


1Ye
4 0)
 LL9EB	&(iiA 3y>1'
	"# 	&)$% % %s   BBB'c                     t        j                  t        d      5  t        j                  ddd       d d d        y # 1 sw Y   y xY w)NLine-delimited JSONr9   znofile.jsonr   Tr   )r;   r<   r=   r*   r+   rC   rF   r6   &test_read_orient_not_records_and_linesrK   A   s:    	z)>	? @
]7$?@ @ @s	   >Ac                     t        d      5 } t        j                  t        d      5  t	        j
                  t        | dd       d d d        d d d        y # 1 sw Y   xY w# 1 sw Y   y xY w)Nr   rJ   r9   r   Tr   )r   r;   r<   r=   r*   r)   ddfr>   s    r6   'test_write_orient_not_records_and_linesrN   F   s[    	 ;A]]:-BC 	;JJsAgT:	;; ;	; 	;; ;s"   A#AA#A 	A##A,	blocksize)      !      i_ c                >   t         |j                  d            j                  t        j                  d      }t         |j                  d            j                  t        j                  d      }t        j                  t        d      g dd      }|j                  |j                  dz         }d	}d
}|j                  |||       |j                  |||       t        j                  ||f      }t        j                  |ft        |      z  |      |d<   t        j                  |ft        |      z  |      |d<   t        j                  ||g      }	t        j                   t         |j                  d            ||d
|       }
t#        |
|	d       y )N	fil1.jsonr#   	fil2.jsonrP   r	   r
   r   r   er         ?rD   r   Tr   r%   r'   	fil*.json)r   r    r"   rO   Fcheck_index)strjoinr.   r/   r0   r,   	DataFramerangeassignr   r)   CategoricalDtyper-   r1   concatr*   r+   r   )rO   r   fil1fil2r(   df2r   r    
path_dtypesolress              r6   .test_read_json_multiple_files_with_path_columnrj   L   sK   {v{{;'(00=D{v{{;'(00=D	E!H+DE	FB
))bddSj)
!CFEJJtF%J0KKV5K1$$dD\2JD7SW,J?BvJ))TGc#h.jACK
))RI
C
,,KFKK$% C c3E*rF   c                R   t        d      5 }t        j                  || d       t        j                  || d      }t        j                  || d      }t        ||       | dk(  rt        t        j                        |_        t        |t               d d d        y # 1 sw Y   y xY w)Nr   Fr   r   )	r   r(   r)   r*   r+   r,   r   listr   r2   s       r6   test_read_json_basicrm   d   s    	 A


1V5
1ae<LL6?	&)$X!"**-FN&"  s   BBB&fkeywordpandasr   c                      fd}t        d      5 }t        j                  |dd       t        j                  |dd|      }t        j                  |dd      }t        ||       d d d        y # 1 sw Y   y xY w)Nc                     dk(  r/t         j                  j                  t        j                  |        S t        j
                  |  S )Nr   )r,   r_   	from_dictr   loadr+   )argskwargsrn   s     r6   _my_json_readerz0test_read_json_fkeyword.<locals>._my_json_readers   s8    v<<))$))T*:;;||T""rF   r   r   Fr   )r   r    engine)r   r(   r)   r*   r+   r,   r   )rn   rv   r3   r4   r5   s   `    r6   test_read_json_fkeywordrx   q   sj    #
 
 %A


1Ye
4a	WLL9EB	&)$	% % %s   AA11A:rw   ujsonc                    t        d      5 }t        j                  |d       t        j                  || d      }t        |t               d d d        y # 1 sw Y   y xY w)Nr   F)r    )rw   r    )r   r(   r)   r*   r+   r   )rw   r3   gots      r6   test_read_json_engine_strr|      sL    	 A


1E
"ll1V59#r  s   A AAc                   t        j                  t        d      g dd      }|j                  |j                  dz         }| dk(  }|j                  t         |j                  d            | |       |j                  t         |j                  d	            | |       t        j                  ||g      }|j                  d d
 }| dk(  rd
dgx|_
        |_
        t        j                  t         |j                  d            | ||      }t        ||       | dk(  rBt        j                  t         |j                  d            | |dd      }t        ||d       y y )NrP   rW   r   rY   rD   r   rU   r   rV   r   r   r   rZ   )r   metar    T2   )r   r~   r    rO   Fr[   )r,   r_   r`   ra   r   r)   r]   r^   rc   ilocr   r*   r+   r   )r   r   r(   rf   r    rh   r~   ri   s           r6   test_read_json_metar      s6   	E!H+DE	FB
))bddSj)
!CiEJJs;6;;{+,V5JIKKKFKK,-fEKJ
))RI
C88BQ<D&'V+dl
,,KFKK$%f4uC c3llK()
 	#s. rF   c                H   t               5 }t        j                  j                  |d      }t        j                  || d       t        j                  || d      }| dk(  rt        t        j                        |_	        t        |t               d d d        y # 1 sw Y   y xY w)N1.jsonFr   r   )r   r/   r'   r^   r(   r)   r*   r+   rl   r   r   )r   r'   fnr4   s       r6   test_write_json_basicr      sx    	 TWW\\$)


2fE
2bu=X!"**-FN&"  s   BBB!c                    ddl m dgfd} t        j                  g dg dd      }t	        j
                  |d	      }t               5 }|j                  |d
| i       d   sJ t	        j                  t        j                  j                  |d            }t        ||d       d d d        y # 1 sw Y   y xY w)Nr   )getFc                     dd<    | i |S )NTr   rC   )rt   ru   flagmp_gets     r6   my_getz%test_to_json_with_get.<locals>.my_get   s    Qt&v&&rF   r   r   r   r   r   	scheduler)compute_kwargs*r[   )dask.multiprocessingr   r,   r_   r*   from_pandasr   r)   r+   r/   r'   r^   r   )r   r(   rM   dnresultr   r   s        @@r6   test_to_json_with_getr      s    27D' 
0|D	EB
..
+C	 1RBV'<=Awwbggll2s34&"%0	1 1 1s   AB66B?c                    t        d      5 } t        j                  t              5  t        j                  | dd       d d d        t        j                  | dd       t        j                  t              5  t        j                  | dd       d d d        d d d        y # 1 sw Y   dxY w# 1 sw Y   xY w# 1 sw Y   y xY w)Nr   r   Tr   Fr   )r   rO   )r   r;   r<   r=   r(   r)   r*   r+   r>   s    r6   test_read_json_errorr      s    	 9A]]:& 	6JJqJ5	6


1WE
2]]:& 	9LL7a8	9	9 9	6 	6	9 	9	9 9s9   B:B"9B:8B.B:"B+	'B:.B7	3B::Cblockc                2   t               5 }t        j                  j                  |d      }t        j                  |dd       t        j                  || d      }|j                  dkD  s| dkD  sJ t        |t        d	
       d d d        y # 1 sw Y   y xY w)Nr   r   Tr   
   )rO   sampler      Fr[   )
r   r/   r'   r^   r(   r)   r*   r+   r   r   )r   r'   r   r   s       r6   test_read_chunkedr      sz    	 ,TWW\\$)


2it
4LLuR8!urz22!RU+, , ,s   A9BBcompression)Ngzipxzc                   t               5 }t        j                  t        ||        t        j                  t
        j                  j                  |d      |       }t        t        |d       d d d        y # 1 sw Y   y xY w)Nr   r   Fr[   )
r   r*   r)   rM   r+   r/   r'   r^   r   r(   )r   r'   r4   s      r6   test_json_compressedr      sY    	 1T


3+6bggll45;O"f%01 1 1   A$A88Bc                    t               5 } t        j                  j                  | d      }t	        j
                  t        |d       t	        j                  |      }t        t        |d       d d d        y # 1 sw Y   y xY w)Nz	*.json.gzr   r   Fr[   )
r   r/   r'   r^   r*   r)   rM   r+   r   r(   )r'   r   r4   s      r6   #test_read_json_inferred_compressionr      s\    	 1TWW\\$,


3/b!"f%0	1 1 1r   z2023.9.1z)https://github.com/dask/dask/issues/10515)reasonc            
     
   t        d      5 t        j                        } | t        t        j                        D cg c]%  }t
        j                  j                  | d      ' c}k(  sJ 	 d d d        t        d      5 t        j                  d      }t        j                  | } | t        fdt        t        j                        D              k(  sJ 	 d d d        y c c}w # 1 sw Y   }xY w# 1 sw Y   y xY w)Nr   .partF)computec              3  d   K   | ]'  }t         j                  j                  | d        ) yw)r   N)r/   r'   r^   ).0nr3   s     r6   	<genexpr>z'test_to_json_results.<locals>.<genexpr>   s+      
-.BGGLLqcK(
s   -0)r   rM   r)   r`   r   r/   r'   r^   daskr   tuple)pathsr   list_of_delayedr3   s      @r6   test_to_json_resultsr      s    
 
 WAAuS__?UV!aA3e5VVVVW 
 
A++a+7o. 
272H
 
 
 	
 
	
 
 WW W
 
s)   2C-*C()C-AC9(C--C69D)-
__future__r   r   r/   fsspecro   r,   r;   packaging.versionr   r   dask.dataframe	dataframer*   dask.dataframe.utilsr   
dask.utilsr   r   r_   r(   r   rM   markparametrizer7   r?   rH   rK   rN   rj   rm   rx   r+   r|   r   r   r   r   r   r   r   skipif__version__r   rC   rF   r6   <module>r      s   "  	    %   * &R\\,<@AbnnRQ' #UV% W%T%,@
; &=>+ ?+. #UV	 W	 h%78
% 9
% GR\\#:; < #UV/ W/< #UV W1&9 "9:, ;, (<=1 >11 F7:#666  
	
rF   