feat: add file_sort_order to register_parquet

Author: dlovell

python/letsql/backends/let/__init__.py

@@ -8,7 +8,7 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow_hotfix  # noqa: F401
-from letsql.internal import WindowUDF
+import toolz
 from ibis import BaseBackend
 from ibis.expr import types as ir, schema as sch
 from sqlglot import exp, parse_one
@@ -17,11 +17,13 @@
 from letsql.backends.let.datafusion import Backend as DataFusionBackend
 from letsql.common.collections import SourceDict
 from letsql.common.utils.graph_utils import replace_fix
+from letsql.common.utils.parquet_utils import get_df_sort_order
 from letsql.expr.relations import (
     CachedNode,
     replace_cache_table,
     RemoteTableReplacer,
 )
+from letsql.internal import WindowUDF
 
 
 def _get_datafusion_table(con, table_name, database="public"):
@@ -90,6 +92,13 @@ def register(
     def read_parquet(
         self, path: str | Path, table_name: str | None = None, **kwargs: Any
     ) -> ir.Table:
+        if "file_sort_order" not in kwargs:
+            try:
+                kwargs = toolz.merge(
+                    kwargs, {"file_sort_order": get_df_sort_order(path)}
+                )
+            except Exception:
+                pass
         registered_table = super().read_parquet(path, table_name=table_name, **kwargs)
         self._sources[registered_table.op()] = registered_table.op()
         return registered_table

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>()?;