feat(nodejs): add order_by method to Query (#3123)

nodejs/src/query.rs

@@ -3,6 +3,12 @@
 
 use std::sync::Arc;
 
+use crate::error::NapiErrorExt;
+use crate::error::convert_error;
+use crate::iterator::RecordBatchIterator;
+use crate::rerankers::RerankHybridCallbackArgs;
+use crate::rerankers::Reranker;
+use crate::util::{parse_distance_type, schema_to_buffer};
 use arrow_array::{
     Array, Float16Array as ArrowFloat16Array, Float32Array as ArrowFloat32Array,
     Float64Array as ArrowFloat64Array, UInt8Array as ArrowUInt8Array,
@@ -19,16 +25,27 @@ use lancedb::query::QueryBase;
 use lancedb::query::QueryExecutionOptions;
 use lancedb::query::Select;
 use lancedb::query::TakeQuery as LanceDbTakeQuery;
-use lancedb::query::VectorQuery as LanceDbVectorQuery;
+use lancedb::query::{ColumnOrdering as LanceDbColumnOrdering, VectorQuery as LanceDbVectorQuery};
 use napi::bindgen_prelude::*;
 use napi_derive::napi;
 
-use crate::error::NapiErrorExt;
-use crate::error::convert_error;
-use crate::iterator::RecordBatchIterator;
-use crate::rerankers::RerankHybridCallbackArgs;
-use crate::rerankers::Reranker;
-use crate::util::{parse_distance_type, schema_to_buffer};
+#[napi(object)]
+pub struct ColumnOrdering {
+    pub ascending: bool,
+    pub nulls_first: bool,
+    pub column_name: String,
+}
+
+impl From<ColumnOrdering> for LanceDbColumnOrdering {
+    fn from(value: ColumnOrdering) -> Self {
+        match (value.ascending, value.nulls_first) {
+            (true, true) => Self::asc_nulls_first(value.column_name),
+            (true, false) => Self::asc_nulls_last(value.column_name),
+            (false, true) => Self::desc_nulls_first(value.column_name),
+            (false, false) => Self::desc_nulls_last(value.column_name),
+        }
+    }
+}
 
 fn bytes_to_arrow_array(data: Uint8Array, dtype: String) -> napi::Result<Arc<dyn Array>> {
     let buf = arrow_buffer::Buffer::from(data.to_vec());
@@ -128,6 +145,18 @@ impl Query {
         self.inner = self.inner.clone().with_row_id();
     }
 
+    #[napi]
+    pub fn order_by(&mut self, ordering: Option<Vec<ColumnOrdering>>) -> napi::Result<()> {
+        let ordering = ordering.map(|ordering| {
+            ordering
+                .into_iter()
+                .map(LanceDbColumnOrdering::from)
+                .collect()
+        });
+        self.inner = self.inner.clone().order_by(ordering);
+        Ok(())
+    }
+
     #[napi(catch_unwind)]
     pub async fn output_schema(&self) -> napi::Result<Buffer> {
         let schema = self.inner.output_schema().await.default_error()?;
@@ -328,6 +357,18 @@ impl VectorQuery {
         Ok(())
     }
 
+    #[napi]
+    pub fn order_by(&mut self, ordering: Option<Vec<ColumnOrdering>>) -> napi::Result<()> {
+        let ordering = ordering.map(|ordering| {
+            ordering
+                .into_iter()
+                .map(LanceDbColumnOrdering::from)
+                .collect()
+        });
+        self.inner = self.inner.clone().order_by(ordering);
+        Ok(())
+    }
+
     #[napi(catch_unwind)]
     pub async fn output_schema(&self) -> napi::Result<Buffer> {
         let schema = self.inner.output_schema().await.default_error()?;