From 5fc454080feddd7ae248a854be6315011dfe71ea Mon Sep 17 00:00:00 2001 From: Pete Gadomski Date: Thu, 17 Oct 2024 07:05:38 -0600 Subject: [PATCH] feat: include duckdb in python bindings --- python/Cargo.toml | 4 +++ python/data/extended-item.parquet | Bin 0 -> 34991 bytes python/src/error.rs | 6 ++++ python/src/lib.rs | 1 + python/src/search.rs | 52 ++++++++++++++++++++++-------- python/tests/test_search.py | 5 +++ 6 files changed, 55 insertions(+), 13 deletions(-) create mode 100644 python/data/extended-item.parquet diff --git a/python/Cargo.toml b/python/Cargo.toml index a1f926db..77ef6f67 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -13,6 +13,9 @@ name = "stacrs" crate-type = ["cdylib"] [dependencies] +duckdb = { workspace = true, features = [ + "bundled", +] } # we don't use it directly, but we need to ensure it's bundled geojson = { workspace = true } pyo3 = { workspace = true, features = ["extension-module"] } pythonize = { workspace = true } @@ -25,4 +28,5 @@ stac = { workspace = true, features = [ "validate-blocking", ] } stac-api = { workspace = true, features = ["client"] } +stac-duckdb = { workspace = true } tokio = { workspace = true, features = ["rt"] } diff --git a/python/data/extended-item.parquet b/python/data/extended-item.parquet new file mode 100644 index 0000000000000000000000000000000000000000..5ef4689e47affbd604a041b2766b9d15b0624be5 GIT binary patch literal 34991 zcmd^o3v^sXc3?kAHf;~%!N7ZYViiX58D8cLwpwa^+%gQ~*J}NYWJ|)9WeLN2iPv!)$6*|YupA$j5aO^LVi@8OLWmFRWr^cueK;&TnPphk>m^>6 zWqtPETd!Wd?)SPS>oYT)^l04G^|`;gb?a8ut)fpKXppO9MOJ#`3fXnN%XN**-_R8L5{k4{-%yyVQNO@&8U=XD=RA>>2Ptq^nHy~MD%P` zBA&}v#q+Tw=*wE=+%@p09bbhC^K&8|)Ur;lmw?jlM?jNZVar z+vvE#p_1y9Y09TyxK_lYG+xe0$~j5fhU_~vE(sBczFoeZ6MYj9-B4fGT;I~*iDqX* z(Yl;>%GXfm^>{be*=fJ2bX~PfnKLKXz!2FIO~;9yeHF^S3T+tKDFbLeB)}ZlK+aqh zFxNBKsN1@r3UgKZ*O|2Ive7MI1EQ0W{751;6-y9}4=4|48H61bHQB>2;&NDUky0>ysW&o$vXq$ra~#G#XMGQk-Lsmy~}^Eh_v8nZ~f zuDGR*;<9~YOzUpteF^@wdywM!pab;}$PaMptAWzH;>k=RR@*)kqVG0Ssc(kDuARL- zwZShckR^MYP-oy7pZJ@Nnq;DlfN3+8SuWRq?vYDsOWdf`&7a&vo6_vtq^sNr% za_ zzEr?ww=We?S;2o5Ewjl{c|~3-A+)5?ybc@YwUE?R*>oai@x;h4Jtir~B&`Xl zHX#?{04h&P@W<&KGG;HffKF#aW3jribb2fi3uWRtn#j_bSZY%wJ(lAuY;pi{PHyT* zC>2V~=Hrn%Snt~UY6xeu6+on}6kyRQ4zrz`^D1n_r&Y?QRoc5W*{u{!E6OZrhbC)W zH@*hC(X`1QN=0jRqhD=(`n948YH`NGWN0xQuS&|RlIBMpE{5d@9B7`AXE@D1pgEod zE%ty0j^0*l@tnUs7E7~CSkR5e6R~ec66wk4rgSDBPsZnBu$sm-qh*}>77IyTbjyW_ zOYGF0tW-``Y6nQ=7L9aipw|nRhpVtflS*F=CaK11c(rwHIBr;GlUR2wy#-ra z>~d8qlw?bgHznmwNsFRl+fV@6cgS~evT;EhPE5v-#<6THma2qAHWtONR6J{=7?z|_ zwF+A-ncT0y(j}QrjhIxwBq=XRS`eu&L<1y~uzk?9a@)WlL_`EgqNospC?61mN)16H3y9#T@sUk78u z)n|MMjqgLox1B#;qPlt8CG&T~=wu=>G98cR$4ACPiP4c{XvW5&9k0)A&oekYF3HDD zUtl}%vV2o?-9%;#q*7>6w=XYzb(!IrJ>=vnu7Tr(WH=RyC&Y|@s!BOkrJW}8yUet# z1aF|zH_12gxf1uEc5^5MYWR+R0)QZzg5epfzTZ zDv}E%OH7+qs!9R0Q)Mau2&-^%6S|>_q@b$gG$$%M&T~&o%F|qUCh{T-K;}x0Vkbv_ zIW-;NFRP{lJRx@nw|(_gJU1Eo84TeSVhp+)&g=LP{wfUNM&_aq5?{7q{MCYDweaUc zBxM=PFG$J@l6D7<6|wJmE@DJ=4V{wFgMzK^3~{N zt_Gj{8|HZp_Y)2DWjylL%8L^GX#+^D4FS+vFV}Ng4F}wG$hn7ZBec**CK*;PipnBK ze6?~yf!b}x#n5>9s^Uv{iMw>{6`>y#v)GM(=4f(9?2Q*3U)jZ8+9 zj+=1bT|s2(N&**_$zfKn1a^g2{@5zzu~pgy^2%3|q!s5E2Ej&oBe&vUGhMqS-^A1*p{m_4Dng z*?6g)uZ22qjd}SLtAeL0l&31#4$Zt^7ls%_P;@AhqhoZ~Pv<$9#PL}VVHOG%NVaU}{ovX?;VIbLu~F&fC90V!q{)~zR5xtU#bAe(MHW6TDUY+A zJJC+0d?qZBhfsm0BB$3pxH94k3WvZ~3yRgkpM_-d;75zLTix9G00CvsNbtwUXg_GV zUYlbv5@rAd#6yJgweeIWF&T~JYD2YCp=>;spT$s#7z?P#)sAMW|<;W0Cm z(I));Z*q%pz4J3x&mpsymYi|_$QqAkY&L@GN}-in%D1kedVp^kn1 zp`Z4@c=Lb#j=e*9iR$JlK^}JKS|hx>ax7=>Ox{E^_L*Qro?X8h>GQuF$l)cbo2Q(~ zSqEn&bO1LJNkinAP4jtc0mo-T6nS>ae)HF$d*Ndj6faTTJXJ^;*ZegY^t~mP$-(Yp z(QxDq1qh!BOyn6v!4j4S(4CYSQA4+|Zy=6>IB~9}5PdfHyT|kY|v4uY7Nhd@U`_8IkW}h`jl^_SIW-_T&9m z@H=yVh4K3$Y5}-88{tv_X9sIw3l#$);Pd_~#15{bS6z{ipX)0tS3Tm@U`fq2)VGe9 z%4b8V91iaS#;#XBli*ML9x;X>3&HycQDL6Jn8hrTQswU4YF7?kVyG6>G38HD*L)OR zkHPO-a+9eMIAWTjU{0Hs*kKu}3s4K;7mo2!`mA%Hj+dxzo^qyc9Z{Da9fbqDD4rA+ z=`U|B;P_04BF`>2YIphGQj|f7(Hc_DlwD`aEi@NTPSVj_;Z(n;vL0P~q(MV`I8esrSm`2TICaakx4H!o4mlrwoUkQWPO z;h+&!Jqi~tLaB(|g|~zAnScvXfc#d%7gz&yfcvY`H&2x# zYSxJ;-cW$>nZQJzov01pyLjQTqfV0IC90dJ$`LigPOS?EKW`{N_)K6T&rZ|}vCKn9 z|K3SbyhL^LR5_w%orvNM1qh!BOyn6vy(Y=8najYtVSbWuvcs?l+%@K!z=hBGE&k0> z^7_TUd3~F$yL&~^sVfU!^o(^k!wmsj5Gk3RCA)ZNm2zm61~)FidtO;;eg@1|udk>! zZ)1v-Tv4sKTvW)eK2{rdCoKhwy<}IdF8*?y1-;W1vuejXW|l@}E^<|uyUAq%K{u`7 zjw)T_5@NF7U#eDb!!7Q`-mVg$#RVZiix+$>)YxKztKn20Hf-2c9sy&X+bq>;>ANTv zY+DV6a8bIb`z6)tQbF8nA*bx!gX;2^4=i9fNsEr67hEkCMA$-(JxaaWAxz!Uz(EIW zS(fT%Kz6l=$hH)dG?o$^ZAfDYAgp~{9C}-Ai@mkbzUaWy>cvN%TF_HWDU4B!1hr(d z08X%=6whfXph(GT6;Dwuxn>w5^be0sv=;LtM=0jube7^@u{^OqjM`FpdiYQmWO$-xkJx&dO>WpI7sM3kzih3fpKYm|tP}S-z!_ z8FmT^%a5Wh#lK+P=Wwi8?R2ha;b7h)y6XAQ=q%vuXKY~qZHH!-2Ig562$&y+S&9n} zoj@TAj*=`Dw=VTC2Q2>*LHvq^9P@wx*fLKM#IQH>oQuWYJ(vx1JKh4mXyd!Oc<0+v z)!c{%RD7S>TFiHdp_nJF#eA<8ih0sn%(q#gnD4z>i($E`F;?z<4H?9bLRz$`NoW;Sae5;t}(YJ zpgmuITH2eN1W;@)LoLMdB`p;5q_s6)$wDztT8g(Wb@>W|V2L2I201m#7Wp1nhPkS1 zloJ*5ppiVVMt-5fNdIJ=4BI@c%KPi#_}55(a-%$8BrpGx{1cviu0~E6*~h;w@8j9; zZGbo+Bmcy&$`?5Fi<{)S3k>3BXWK#0sBa#Vn|6?M3cE2P)%h^15NRr_`5bKO*e+PGA zj2{PoPZ9$A;D?eNHVTeFz%-}ny+4FX9M*|vMWwUPiS)id76cys6ARQ&|5Q|Z;w6!O z{biBf`-&uI49rJim*{&&dhcI|IxqjFNMHObk$&ba$yC7oXGH#)w?+E!e-OkS|BoX7 z%-@RiXMZQs$KDg^eec6%P~skx|BG;)D%| zV($dG4aE3di4`iW0TR~V6bod{U|}HwmvQ4uzSae>#e)ZORS5*e?F%pQ%BAQkXlCmK zK+U3Wop@%oEfR*><^?r*?Lz2W288?~oD+q%Uct*|R-s#CUwM8ND55>U6e5El-zh|s z`bVr=enx^ey!(G#V#eT0dG(?ci!{v=wPx*YGjQyd>w6wsUxU`ius`8uDcKuxtA=R0CR zlF5=`<)UM9#|DU(H6RikPOS~v3J;7~rtPej^({Xul<65)ZM=+|Z1fi%IG+eoO#|-PJYKUu@M5Nj3b%D%+!mTQqlxM-7#ezgr>S zsOEhk zJFrH(Sow2%%T9S8h*y=5f%q$uOS}9vlX|uA#krrX(~hpY$~Whf&j|QO|4ou7T-u{I z2-tkiC-Jg3;ND-aJ+gjDufQq$k0LMce1tBWOM4Zfy!gSt&7UB|_^{h2H)^llXnWdC zc^W_|%K4v4a-B=NP;El8JrKs@(D|Eamw!onR#}lp(Ueof-XkC5iykiR^e>kt3LPEb zkiS}Xu15Rdp8$^z@0Ka=5~HX7yCnNu+EcZqVe-If{;=5buWPU0Jpa>T%Ev_YWkAs9 z(hh&4Gy-PliEQ?M*LxeZ*M6mDxb&gPLY@y(9{fF2iq}7pq6gz(5ez;ex{UAlc9e0(%kz{5rOcG1P#v=g^2 z`1K;?3&tHFy9cJM@0iH0`t_mncW6)DQS$vE<-n6@M&G+E$#=N4AGe!$_}DX-Dnd8d zvpR41wM$!ThQE~W>L||u*Q#>(e>2_rWLs&rEo#GF!nxS3J=?v&w{Vo#NV+fnA4v|n zw6}JYMkBO^T~>XjS3B5KGrZ`-k;Nb6(bzgi@_FKO=5rjo%OsqupQ#p(*G}1cFsBE! zLjy3W;W?mqXhyk6GTHwHP26AXGC}j3q9*7vcSYUsbPN1x2X|@D?{s*-MS1v#sMec~ zX|W_MEb@3<`((V#GaSlglHkd`6|&!@y_hHiTbK_9`s^O<$Q}S~o+Miim5J`hj^X(D zjDfzGGC}k4Ejq_>fU}1%t|YZLlIURa<7mstG+{YLSYG;dV&I*;iG{0xu{~4V|L5=1 z4rZ3O|5xbs2jxZJP*uJF4$;SZdbSjcmd%V`?ep)K zLLfANO(5s)*G|v<{3ehelUzOqs&UbJ?$^Z*0M{_1?;X{PJzD@r654BfVL@+Hf zUk#R`S%g0r0Ou*%>4yLrhP#MI!jxwT(^aT?+G;o;0iJM@SI zrklqkSSJ#D9trn#BIXJZ15V+-A6yMFCFT)M4jdb4r+)*<;V}@_nZ%w)~1)m!f9uI3%GX( zT3|O!LNn%pd+iR6IjeTmdSMg6TN& zq##4J>`a4i!tYbVxuW*=QJ}yc9^w!&LO4znjw3&)z@^;jV}OERHZym*2?4u20B4oj z$seN9^0*R_@;T%&r3DnKas^Pp7<%FlEL>PRi_Haa{;6H~ePEz4CZ#Aug`i$0s4qWG zv%v??SfEozho-^_`kLE4nIGzbPk*mYkQx? z&f(E9BIR?)BQZV=C{*RdzoO20;g75s?AL7Jo zTqNm#1b7BvB6<<<@GDY97gph{U3=qCQ7k+dNThrYc?A8@r%>>({w4|j)0ZsJi%w-M zJq!o=+K>MoBI031BIR?)BSgmt(G`a1=|2NRe1waW9t#c#5YP_4j5v5~l1TX+@(9OE zgyZD{#LUH41de&HgF*~~cKQ^eRG*A`f5vu12)g^{%@819_e&E4R?+61szk=uz?Tf#Laca9Lsk}*Wp8G9wL@%AO zknD8gXQ3LRRJ0EPuZU0+DW4g&Ch~~t4+z@_fDPuZ=l;C~oAW861s_DhXvf}0d_3e$ zq~9eVkJb|@pFfa$}Mmh0oN`>ZzuNxtpaie;z zJ}|IjXbTJPQcrC1&}Xb)*LD0qe(u!uX}xVs--cQ6t?PcuhoA6U*Sql-WbtqK(F8ed zk#_iZ2pm)MQX+L0`(su!*=?*CG$KM_J zdq+QjYU@T&gGAOe{0{5_vh;}F3%Sz)gs%gW&e;}fXg46;H*dS^5d_BJGGcIGV0C8F%MYBfBd@=+xiCQ;6Jg{HIN9W zwr9fLrj~HBBOkA~XV#m2_h{O+Ndsp#qbZ;)A2lb?Yo(v z1^OHshOtPT>N^lzq-)z$Xd+V|Np?&|yuqnRvL3oN0E4I)KN~cp+H7#`0j2ls>VrYU ziilnS>DR|e-@0#wB8Cbv(MU1$IZ+=R0zB{`d)(!p4SUlq5%2hTq!ZL&8{wy>om30{ z>F*}}nC+YJcvw>`q2ZzHr0>&_Wb1u9l7rI`Zv=`!K0u%2c7 zXKO>Sqje(O7)(T9$U*xd!~*>eHlIl9Hn0-Eq?AZWe=p}BNCta3hNhW5T>jLR!N%=- z^uR>DUB1(7z!M5U4EDPf0oCE^m#wvu=CS4F6>tnXz)xE|IiSmk3}7WyTHVesCPGR1DiCE2%>zz9+&1v zk6?Klw(lHBM7-N)L&N_1&`|57-Z9Z&=f`4i+NqpCM>`YsVbn>K55Qyn&FO8wlg(dH zOva!53?Q25ca#Uow{u_|`U+@iNKnh(o#pB2p>ptJz~At;=D=1=09{jIK;!Eh_IPHA z9&RuFf!G+5*@^Xs!NUOOfYLZvT5*o*G4vRI-%9OwdO;pUw5bpMUaCa$r`9@+I-ADkk@;uoNT?|uYL@>iKg;C8sUb21z^B~f~M4e^+=D0+85j7U1Sc- z@fw&$e!KcWNoFb2+8;;?(Km^yZX}R7X(Uho95?~3-7sky_KGn8P1;Dl-2pIVt=N)% zSkLUlAHmdMF5J=D7#bSDc{&YqE6hpgv$VCBgx53Y1aAQ0?e-2%f^NWxk^R2oCvL zcEem1>|)~W#ra@{g-UN1)r#H|p=pCO84zcNnTyyYGB8`LN zVd%BCzSIKr5$y+a>PcSuL^IG^e=DeRCzM$H^ImcqSl3X0+==iTc4Ga~C;m#)+e_t)USUtj)JHqV8^S}| zK(7&N=c%b@-U z>}&!39mD#?>m}h7GRuaZkls%4seg7O{H8f*&(=$_8^iHIN<#-z?l291Rb3H0iw)2#*wwkRZJ_E zbk5i}kaG>Lb!N@^2^bH&a4VogswZv=io|dF9)L@H6w!n-KI` z{NZBxP9VNA@}=@3`Rcl-tbGbhBYkYX5c0jV52n5{^6kWOn|yIV`tOT)TPJ&9{)9F6 z9`Jd`QFK9qSLE*v(I_L;-Hw%5|3r}9_&oSygrDRi{H>ka|0H_IKB1Apk3X=3=%=AcOzfB-{J4HC)~88oLi)t) z8~%o^Gx*T&jy@PV^#NROm+s#Hl^4yo{xsq%Gym+yaw5{=&+b8dW%_pv%csen!W=!V z=jW-Pq58;=-!k6lcdB?*I&kg<`-{@k^)#5T zGW_eiv7D{vwoin1X7+@=d-Pp6f0W_BAIqn)Yy25lvw{DQ^^5wy*k9`eAIfv?`*`&)>#0x z9~TRt^xy*d*t`Mbj0OR7olJf9lP3)=-~;qQ6%svF74};HwhM05#nS3uOZRUu{zLN{ zwrpsqYpAQ=uyI2s3z3J}S$Ik~F_}z30lyOlWh^hZVaxp@7t*m*Bpr>X#-Q@hoqnjt zn~dOlpiuq08~W0T*|BtL!`&M<;2V#S)v(!H*V@pqxv{0Cwbk3W+2`HZP}kVh=<_u< zv^F=iG&I-OH*Ta_%}p(>Ee%br&5fHIny^+&Ykkw^mKJY4RBPSbeE0Xi|9zKQoe5>{ mn~dc*WplL#TeZ#JI&WP)Tx5eQk6-FP*RDHUu0DOB;r|0dpv8m$ literal 0 HcmV?d00001 diff --git a/python/src/error.rs b/python/src/error.rs index 395648be..7c27ea4a 100644 --- a/python/src/error.rs +++ b/python/src/error.rs @@ -16,6 +16,12 @@ impl From for Error { } } +impl From for Error { + fn from(value: stac_duckdb::Error) -> Self { + Error(value.to_string()) + } +} + impl From for Error { fn from(value: geojson::Error) -> Self { Error(value.to_string()) diff --git a/python/src/lib.rs b/python/src/lib.rs index 524eb996..12f69490 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -7,6 +7,7 @@ mod search; mod validate; mod write; +use duckdb as _; use error::Error; use pyo3::prelude::*; diff --git a/python/src/search.rs b/python/src/search.rs index 4e016788..d5bf630f 100644 --- a/python/src/search.rs +++ b/python/src/search.rs @@ -6,6 +6,7 @@ use pyo3::{ use serde::de::DeserializeOwned; use stac::Format; use stac_api::{BlockingClient, Fields, Item, ItemCollection, Items, Search}; +use stac_duckdb::Client; use std::str::FromStr; use tokio::runtime::Builder; @@ -40,6 +41,9 @@ use tokio::runtime::Builder; /// will be interpreted as cql2-text, dictionaries as cql2-json. /// query (dict[str, Any] | None): Additional filtering based on properties. /// It is recommended to use filter instead, if possible. +/// use_duckdb (bool | None): Query with DuckDB. If None and the href has a +/// 'parquet' or 'geoparquet' extension, will be set to True. Defaults +/// to None. /// /// Returns: /// list[dict[str, Any]]: A list of the returned STAC items. @@ -53,7 +57,7 @@ use tokio::runtime::Builder; /// ... max_items=1, /// ... ) #[pyfunction] -#[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None))] +#[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, use_duckdb=None))] pub fn search<'py>( py: Python<'py>, href: String, @@ -69,6 +73,7 @@ pub fn search<'py>( sortby: Option, filter: Option, query: Option>, + use_duckdb: Option, ) -> PyResult> { let items = search_items( href, @@ -84,6 +89,7 @@ pub fn search<'py>( sortby, filter, query, + use_duckdb, )?; pythonize::pythonize(py, &items) .map_err(PyErr::from) @@ -126,6 +132,9 @@ pub fn search<'py>( /// format (str | None): The output format. If none, will be inferred from /// the outfile extension, and if that fails will fall back to compact JSON. /// options (list[tuple[str, str]] | None): Configuration values to pass to the object store backend. +/// use_duckdb (bool | None): Query with DuckDB. If None and the href has a +/// 'parquet' or 'geoparquet' extension, will be set to True. Defaults +/// to None. /// /// Returns: /// list[dict[str, Any]]: A list of the returned STAC items. @@ -139,7 +148,7 @@ pub fn search<'py>( /// ... max_items=1, /// ... ) #[pyfunction] -#[pyo3(signature = (outfile, href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, format=None, options=None))] +#[pyo3(signature = (outfile, href, *, intersects=None, ids=None, collections=None, max_items=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, format=None, options=None, use_duckdb=None))] pub fn search_to( outfile: String, href: String, @@ -157,6 +166,7 @@ pub fn search_to( query: Option>, format: Option, options: Option>, + use_duckdb: Option, ) -> PyResult { let items = search_items( href, @@ -172,6 +182,7 @@ pub fn search_to( sortby, filter, query, + use_duckdb, )?; let format = format .map(|s| s.parse()) @@ -206,8 +217,8 @@ fn search_items( sortby: Option, filter: Option, query: Option>, + use_duckdb: Option, ) -> PyResult> { - let client = BlockingClient::new(&href).map_err(Error::from)?; let mut fields = Fields::default(); if let Some(include) = include { fields.include = include.into(); @@ -225,7 +236,7 @@ fn search_items( .map(|q| pythonize::depythonize(&q.into_bound(py))) .transpose() })?; - let search = Search { + let mut search = Search { intersects: intersects.map(|i| i.into()).transpose()?, ids: ids.map(|ids| ids.into()), collections: collections.map(|c| c.into()), @@ -243,18 +254,33 @@ fn search_items( ..Default::default() }, }; - let items = client.search(search).map_err(Error::from)?; - if let Some(max_items) = max_items { - items - .take(max_items) - .collect::>() + if use_duckdb + .unwrap_or_else(|| matches!(Format::infer_from_href(&href), Some(Format::Geoparquet(_)))) + { + if let Some(max_items) = max_items { + search.items.limit = Some(max_items.try_into()?); + } + let client = Client::from_href(href).map_err(Error::from)?; + client + .search_to_json(search) + .map(|item_collection| item_collection.items) .map_err(Error::from) .map_err(PyErr::from) } else { - items - .collect::>() - .map_err(Error::from) - .map_err(PyErr::from) + let client = BlockingClient::new(&href).map_err(Error::from)?; + let items = client.search(search).map_err(Error::from)?; + if let Some(max_items) = max_items { + items + .take(max_items) + .collect::>() + .map_err(Error::from) + .map_err(PyErr::from) + } else { + items + .collect::>() + .map_err(Error::from) + .map_err(PyErr::from) + } } } diff --git a/python/tests/test_search.py b/python/tests/test_search.py index b4d57599..ba2f07a8 100644 --- a/python/tests/test_search.py +++ b/python/tests/test_search.py @@ -44,3 +44,8 @@ def test_search_to_geoparquet(tmp_path: Path) -> None: table = pyarrow.parquet.read_table(tmp_path / "out.parquet") items = list(stac_geoparquet.arrow.stac_table_to_items(table)) assert len(items) == 1 + + +def test_search_geoparquet(data: Path) -> None: + items = stacrs.search(str(data / "extended-item.parquet")) + assert len(items) == 1