refactor: address PR review comments for SQL dialect support

- Use CustomDialectBuilder for Spark dialect instead of manual Dialect impl
- Move SqlDialect enum from spark_dialect.rs to query.rs as general-purpose type
- Expose SqlDialect as a Python enum instead of error-prone string parameter

Co-authored-by: Isaac

Author: Yu Chen

crates/lance-graph-python/src/graph.rs

@@ -27,7 +27,8 @@ use lance_graph::{
     ast::{DistanceMetric as RustDistanceMetric, GraphPattern, ReadingClause},
     CypherQuery as RustCypherQuery, ExecutionStrategy as RustExecutionStrategy,
     GraphConfig as RustGraphConfig, GraphError as RustGraphError, InMemoryCatalog,
-    SqlQuery as RustSqlQuery, VectorSearch as RustVectorSearch,
+    SqlDialect as RustSqlDialect, SqlQuery as RustSqlQuery,
+    VectorSearch as RustVectorSearch,
 };
 use pyo3::{
     exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
@@ -59,6 +60,34 @@ impl From<ExecutionStrategy> for RustExecutionStrategy {
     }
 }
 
+/// SQL dialect for generating SQL from Cypher queries
+#[pyclass(name = "SqlDialect", module = "lance.graph")]
+#[derive(Clone, Copy)]
+pub enum SqlDialect {
+    /// Generic SQL (DataFusion default dialect)
+    Default,
+    /// Spark SQL dialect
+    Spark,
+    /// PostgreSQL dialect
+    PostgreSql,
+    /// MySQL dialect
+    MySql,
+    /// SQLite dialect
+    Sqlite,
+}
+
+impl From<SqlDialect> for RustSqlDialect {
+    fn from(dialect: SqlDialect) -> Self {
+        match dialect {
+            SqlDialect::Default => RustSqlDialect::Default,
+            SqlDialect::Spark => RustSqlDialect::Spark,
+            SqlDialect::PostgreSql => RustSqlDialect::PostgreSql,
+            SqlDialect::MySql => RustSqlDialect::MySql,
+            SqlDialect::Sqlite => RustSqlDialect::Sqlite,
+        }
+    }
+}
+
 /// Distance metric for vector similarity search
 #[pyclass(name = "DistanceMetric", module = "lance.graph")]
 #[derive(Clone, Copy)]
@@ -494,9 +523,8 @@ impl CypherQuery {
     /// ----------
     /// datasets : dict
     ///     Dictionary mapping table names to Lance datasets
-    /// dialect : str, optional
-    ///     SQL dialect to use. One of "default", "spark", "postgresql", "mysql", "sqlite".
-    ///     Defaults to "default" (generic DataFusion SQL).
+    /// dialect : SqlDialect, optional
+    ///     SQL dialect to use. Defaults to SqlDialect.Default (generic DataFusion SQL).
     ///
     /// Returns
     /// -------
@@ -507,28 +535,14 @@ impl CypherQuery {
     /// ------
     /// RuntimeError
     ///     If SQL generation fails
-    /// ValueError
-    ///     If an invalid dialect is specified
     #[pyo3(signature = (datasets, dialect=None))]
     fn to_sql(
         &self,
         py: Python,
         datasets: &Bound<'_, PyDict>,
-        dialect: Option<&str>,
+        dialect: Option<SqlDialect>,
     ) -> PyResult<String> {
-        let sql_dialect = match dialect {
-            None | Some("default") => None,
-            Some("spark") => Some(lance_graph::SqlDialect::Spark),
-            Some("postgresql") | Some("postgres") => Some(lance_graph::SqlDialect::PostgreSql),
-            Some("mysql") => Some(lance_graph::SqlDialect::MySql),
-            Some("sqlite") => Some(lance_graph::SqlDialect::Sqlite),
-            Some(other) => {
-                return Err(PyValueError::new_err(format!(
-                    "Unknown SQL dialect: '{}'. Valid options: 'default', 'spark', 'postgresql', 'mysql', 'sqlite'",
-                    other
-                )));
-            }
-        };
+        let sql_dialect = dialect.map(|d| d.into());
 
         // Convert datasets to Arrow RecordBatch map
         let arrow_datasets = python_datasets_to_batches(datasets)?;
@@ -1570,6 +1584,7 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) ->
     let graph_module = PyModule::new(py, "graph")?;
 
     graph_module.add_class::<ExecutionStrategy>()?;
+    graph_module.add_class::<SqlDialect>()?;
     graph_module.add_class::<DistanceMetric>()?;
     graph_module.add_class::<GraphConfig>()?;
     graph_module.add_class::<GraphConfigBuilder>()?;

crates/lance-graph/src/lib.rs

@@ -68,8 +68,7 @@ pub use lance_graph_catalog::{
 #[cfg(feature = "unity-catalog")]
 pub use lance_graph_catalog::{UnityCatalogConfig, UnityCatalogProvider};
 pub use lance_vector_search::VectorSearch;
-pub use query::{CypherQuery, ExecutionStrategy};
-pub use spark_dialect::{SqlDialect, SparkDialect};
+pub use query::{CypherQuery, ExecutionStrategy, SqlDialect};
 pub use sql_query::SqlQuery;
 #[cfg(feature = "delta")]
 pub use table_readers::DeltaTableReader;

crates/lance-graph/src/query.rs

@@ -16,6 +16,22 @@ use lance_namespace::models::DescribeTableRequest;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
+/// SQL dialect to use when generating SQL from Cypher queries.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum SqlDialect {
+    /// Generic SQL (DataFusion default dialect)
+    #[default]
+    Default,
+    /// Spark SQL dialect (backtick quoting, STRING type, EXTRACT, etc.)
+    Spark,
+    /// PostgreSQL dialect
+    PostgreSql,
+    /// MySQL dialect
+    MySql,
+    /// SQLite dialect
+    Sqlite,
+}
+
 /// Normalize an Arrow schema to have lowercase field names.
 ///
 /// This ensures that column names in the dataset match the normalized
@@ -302,7 +318,7 @@ impl CypherQuery {
     pub async fn to_sql(
         &self,
         datasets: HashMap<String, arrow::record_batch::RecordBatch>,
-        dialect: Option<crate::spark_dialect::SqlDialect>,
+        dialect: Option<SqlDialect>,
     ) -> Result<String> {
         use std::sync::Arc;
 

crates/lance-graph/src/spark_dialect.rs

@@ -3,9 +3,9 @@
 
 //! SQL dialect support for the DataFusion unparser.
 //!
-//! This module provides a [`SqlDialect`] enum for selecting which SQL dialect
-//! to use when unparsing DataFusion logical plans to SQL strings, and includes
-//! a [`SparkDialect`] implementation for Spark SQL.
+//! This module provides a Spark SQL dialect built using DataFusion's
+//! [`CustomDialectBuilder`], and a helper to build an [`Unparser`] for any
+//! supported [`SqlDialect`].
 //!
 //! Key Spark SQL differences from standard SQL:
 //! - Backtick (`` ` ``) identifier quoting
@@ -16,51 +16,50 @@
 //! - `LENGTH()` instead of `CHARACTER_LENGTH()`
 //! - Subqueries in FROM require aliases
 
-use std::sync::Arc;
-
-use arrow::datatypes::TimeUnit;
-use datafusion_common::Result;
-use datafusion_expr::Expr;
 use datafusion_sql::unparser::dialect::{
-    CharacterLengthStyle, DateFieldExtractStyle, DefaultDialect, Dialect, IntervalStyle,
-    MySqlDialect, PostgreSqlDialect, SqliteDialect,
+    CharacterLengthStyle, CustomDialect, CustomDialectBuilder, DateFieldExtractStyle,
+    DefaultDialect, MySqlDialect, PostgreSqlDialect, SqliteDialect,
 };
 use datafusion_sql::unparser::Unparser;
 use datafusion_sql::sqlparser::ast::{self, Ident, ObjectName, TimezoneInfo};
 
-/// SQL dialect to use when generating SQL from Cypher queries.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
-pub enum SqlDialect {
-    /// Generic SQL (DataFusion default dialect)
-    #[default]
-    Default,
-    /// Spark SQL dialect (backtick quoting, STRING type, EXTRACT, etc.)
-    Spark,
-    /// PostgreSQL dialect
-    PostgreSql,
-    /// MySQL dialect
-    MySql,
-    /// SQLite dialect
-    Sqlite,
-}
+use crate::query::SqlDialect;
 
-impl SqlDialect {
-    /// Create a DataFusion `Unparser` configured for this dialect.
-    pub fn unparser(&self) -> DialectUnparser {
-        match self {
-            SqlDialect::Default => DialectUnparser::Default(DefaultDialect {}),
-            SqlDialect::Spark => DialectUnparser::Spark(SparkDialect),
-            SqlDialect::PostgreSql => DialectUnparser::PostgreSql(PostgreSqlDialect {}),
-            SqlDialect::MySql => DialectUnparser::MySql(MySqlDialect {}),
-            SqlDialect::Sqlite => DialectUnparser::Sqlite(SqliteDialect {}),
-        }
-    }
+/// Build a Spark SQL dialect using DataFusion's `CustomDialectBuilder`.
+pub fn build_spark_dialect() -> CustomDialect {
+    CustomDialectBuilder::new()
+        .with_identifier_quote_style('`')
+        .with_supports_nulls_first_in_sort(true)
+        .with_use_timestamp_for_date64(true)
+        .with_utf8_cast_dtype(ast::DataType::Custom(
+            ObjectName::from(vec![Ident::new("STRING")]),
+            vec![],
+        ))
+        .with_large_utf8_cast_dtype(ast::DataType::Custom(
+            ObjectName::from(vec![Ident::new("STRING")]),
+            vec![],
+        ))
+        .with_date_field_extract_style(DateFieldExtractStyle::Extract)
+        .with_character_length_style(CharacterLengthStyle::Length)
+        .with_int64_cast_dtype(ast::DataType::BigInt(None))
+        .with_int32_cast_dtype(ast::DataType::Int(None))
+        .with_timestamp_cast_dtype(
+            ast::DataType::Timestamp(None, TimezoneInfo::None),
+            ast::DataType::Timestamp(None, TimezoneInfo::None),
+        )
+        .with_date32_cast_dtype(ast::DataType::Date)
+        .with_supports_column_alias_in_table_alias(true)
+        .with_requires_derived_table_alias(true)
+        .with_full_qualified_col(false)
+        .with_unnest_as_table_factor(false)
+        .with_float64_ast_dtype(ast::DataType::Double(ast::ExactNumberInfo::None))
+        .build()
 }
 
 /// Wrapper to hold the concrete dialect type and provide an `Unparser` reference.
 pub enum DialectUnparser {
     Default(DefaultDialect),
-    Spark(SparkDialect),
+    Spark(CustomDialect),
     PostgreSql(PostgreSqlDialect),
     MySql(MySqlDialect),
     Sqlite(SqliteDialect),
@@ -78,114 +77,34 @@ impl DialectUnparser {
     }
 }
 
-/// A Spark SQL dialect for unparsing DataFusion logical plans to Spark-compatible SQL.
-pub struct SparkDialect;
-
-impl Dialect for SparkDialect {
-    fn identifier_quote_style(&self, _identifier: &str) -> Option<char> {
-        Some('`')
-    }
-
-    fn supports_nulls_first_in_sort(&self) -> bool {
-        true
-    }
-
-    fn use_timestamp_for_date64(&self) -> bool {
-        true
-    }
-
-    fn interval_style(&self) -> IntervalStyle {
-        IntervalStyle::SQLStandard
-    }
-
-    fn float64_ast_dtype(&self) -> ast::DataType {
-        ast::DataType::Double(ast::ExactNumberInfo::None)
-    }
-
-    fn utf8_cast_dtype(&self) -> ast::DataType {
-        ast::DataType::Custom(
-            ObjectName::from(vec![Ident::new("STRING")]),
-            vec![],
-        )
-    }
-
-    fn large_utf8_cast_dtype(&self) -> ast::DataType {
-        ast::DataType::Custom(
-            ObjectName::from(vec![Ident::new("STRING")]),
-            vec![],
-        )
-    }
-
-    fn date_field_extract_style(&self) -> DateFieldExtractStyle {
-        DateFieldExtractStyle::Extract
-    }
-
-    fn character_length_style(&self) -> CharacterLengthStyle {
-        CharacterLengthStyle::Length
-    }
-
-    fn int64_cast_dtype(&self) -> ast::DataType {
-        ast::DataType::BigInt(None)
-    }
-
-    fn int32_cast_dtype(&self) -> ast::DataType {
-        ast::DataType::Int(None)
-    }
-
-    fn timestamp_cast_dtype(
-        &self,
-        _time_unit: &TimeUnit,
-        _tz: &Option<Arc<str>>,
-    ) -> ast::DataType {
-        ast::DataType::Timestamp(None, TimezoneInfo::None)
-    }
-
-    fn date32_cast_dtype(&self) -> ast::DataType {
-        ast::DataType::Date
-    }
-
-    fn supports_column_alias_in_table_alias(&self) -> bool {
-        true
-    }
-
-    fn requires_derived_table_alias(&self) -> bool {
-        true
-    }
-
-    fn full_qualified_col(&self) -> bool {
-        false
-    }
-
-    fn unnest_as_table_factor(&self) -> bool {
-        false
-    }
-
-    fn scalar_function_to_sql_overrides(
-        &self,
-        _unparser: &Unparser,
-        _func_name: &str,
-        _args: &[Expr],
-    ) -> Result<Option<ast::Expr>> {
-        // character_length -> length is handled by CharacterLengthStyle::Length
-        // Additional Spark-specific function mappings can be added here as needed
-        Ok(None)
+impl SqlDialect {
+    /// Create a `DialectUnparser` configured for this dialect.
+    pub fn unparser(&self) -> DialectUnparser {
+        match self {
+            SqlDialect::Default => DialectUnparser::Default(DefaultDialect {}),
+            SqlDialect::Spark => DialectUnparser::Spark(build_spark_dialect()),
+            SqlDialect::PostgreSql => DialectUnparser::PostgreSql(PostgreSqlDialect {}),
+            SqlDialect::MySql => DialectUnparser::MySql(MySqlDialect {}),
+            SqlDialect::Sqlite => DialectUnparser::Sqlite(SqliteDialect {}),
+        }
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use datafusion_sql::unparser::dialect::Dialect;
 
     #[test]
     fn test_spark_dialect_identifier_quoting() {
-        let dialect = SparkDialect;
+        let dialect = build_spark_dialect();
         assert_eq!(dialect.identifier_quote_style("table_name"), Some('`'));
         assert_eq!(dialect.identifier_quote_style("column"), Some('`'));
     }
 
     #[test]
     fn test_spark_dialect_type_mappings() {
-        let dialect = SparkDialect;
+        let dialect = build_spark_dialect();
         assert!(matches!(dialect.utf8_cast_dtype(), ast::DataType::Custom(..)));
         assert!(matches!(dialect.int64_cast_dtype(), ast::DataType::BigInt(None)));
         assert!(matches!(dialect.int32_cast_dtype(), ast::DataType::Int(None)));
@@ -194,13 +113,13 @@ mod tests {
 
     #[test]
     fn test_spark_dialect_requires_derived_table_alias() {
-        let dialect = SparkDialect;
+        let dialect = build_spark_dialect();
         assert!(dialect.requires_derived_table_alias());
     }
 
     #[test]
     fn test_spark_dialect_extract_style() {
-        let dialect = SparkDialect;
+        let dialect = build_spark_dialect();
         assert!(matches!(
             dialect.date_field_extract_style(),
             DateFieldExtractStyle::Extract
@@ -209,7 +128,7 @@ mod tests {
 
     #[test]
     fn test_spark_dialect_character_length_style() {
-        let dialect = SparkDialect;
+        let dialect = build_spark_dialect();
         assert!(matches!(
             dialect.character_length_style(),
             CharacterLengthStyle::Length