feat: add file_sort_order to register_parquet

Author: dlovell

python/letsql/expr/__init__.py

@@ -3,3 +3,11 @@
 
 def __getattr__(name):
     return getattr(expr, name)
+
+
+def column(name):
+    return expr.Expr.column(name)
+
+
+def col(name):
+    return expr.Expr.column(name)

src/context.rs

@@ -7,6 +7,7 @@ use crate::catalog::{PyCatalog, PyTable};
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
 use crate::errors::DataFusionError;
+use crate::expr::sort_expr::PySortExpr;
 use crate::functions::greatest::GreatestFunc;
 use crate::functions::least::LeastFunc;
 use crate::ibis_table::IbisTable;
@@ -209,7 +210,8 @@ impl PySessionContext {
                         file_extension=".parquet",
                         skip_metadata=true,
                         schema=None,
-                        storage_options=None
+                        storage_options=None,
+                        file_sort_order=None,
     ))]
     fn register_parquet(
         &mut self,
@@ -221,6 +223,7 @@ impl PySessionContext {
         skip_metadata: bool,
         schema: Option<PyArrowType<Schema>>,
         storage_options: Option<HashMap<String, String>>,
+        file_sort_order: Option<Vec<Vec<PySortExpr>>>,
         py: Python,
     ) -> PyResult<()> {
         let mut options = ParquetReadOptions::default()
@@ -229,6 +232,12 @@ impl PySessionContext {
             .skip_metadata(skip_metadata);
         options.file_extension = file_extension;
         options.schema = schema.as_ref().map(|x| &x.0);
+        options.file_sort_order = file_sort_order
+            .unwrap_or_default()
+            .into_iter()
+            .map(|e| e.into_iter().map(|f| f.into()).collect())
+            .collect();
+
         let storage_options = storage_options.unwrap_or_default();
 
         let result = async {

src/expr.rs

@@ -95,6 +95,8 @@ pub mod union;
 mod wildcard;
 pub mod window;
 
+use sort_expr::PySortExpr;
+
 /// A PyExpr that can be used on a DataFrame
 #[pyclass(name = "Expr", module = "datafusion.expr", subclass)]
 #[derive(Debug, Clone)]
@@ -247,7 +249,7 @@ impl PyExpr {
 
     /// Create a sort PyExpr from an existing PyExpr.
     #[pyo3(signature = (ascending=true, nulls_first=true))]
-    pub fn sort(&self, ascending: bool, nulls_first: bool) -> PyOrdered {
+    pub fn sort(&self, ascending: bool, nulls_first: bool) -> PySortExpr {
         self.expr.clone().sort(ascending, nulls_first).into()
     }
 
@@ -709,6 +711,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<create_memory_table::PyCreateMemoryTable>()?;
     m.add_class::<create_view::PyCreateView>()?;
     m.add_class::<distinct::PyDistinct>()?;
+    m.add_class::<sort_expr::PySortExpr>()?;
     m.add_class::<subquery_alias::PySubqueryAlias>()?;
     m.add_class::<drop_table::PyDropTable>()?;
     m.add_class::<repartition::PyPartitioning>()?;