Fix unparsing OFFSET

apache · Sep 19, 2024 · 8bd089e · 8bd089e
1 parent 0243ebd
commit 8bd089e
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 2 deletions.
diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs
@@ -322,6 +322,21 @@ impl Unparser<'_> {
                     ))));
                 }
 
+                if limit.skip > 0 {
+                    let Some(query) = query.as_mut() else {
+                        return internal_err!(
+                            "Offset operator only valid in a statement context."
+                        );
+                    };
+                    query.offset(Some(ast::Offset {
+                        rows: ast::OffsetRows::None,
+                        value: ast::Expr::Value(ast::Value::Number(
+                            limit.skip.to_string(),
+                            false,
+                        )),
+                    }));
+                }
+
                 self.select_to_sql_recursively(
                     limit.input.as_ref(),
                     query,

diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -622,8 +622,11 @@ fn test_pretty_roundtrip() -> Result<()> {
     Ok(())
 }
 
-fn sql_round_trip(query: &str, expect: &str) {
-    let statement = Parser::new(&GenericDialect {})
+fn sql_round_trip<D>(dialect: D, query: &str, expect: &str)
+where
+    D: Dialect
+{
+    let statement = Parser::new(&dialect)
         .try_with_sql(query)
         .unwrap()
         .parse_statement()
@@ -788,6 +791,7 @@ fn test_table_scan_pushdown() -> Result<()> {
 #[test]
 fn test_interval_lhs_eq() {
     sql_round_trip(
+        GenericDialect{},
         "select interval '2 seconds' = interval '2 seconds'",
         "SELECT (INTERVAL '2.000000000 SECS' = INTERVAL '2.000000000 SECS')",
     );
@@ -796,7 +800,35 @@ fn test_interval_lhs_eq() {
 #[test]
 fn test_interval_lhs_lt() {
     sql_round_trip(
+        GenericDialect{},
         "select interval '2 seconds' < interval '2 seconds'",
         "SELECT (INTERVAL '2.000000000 SECS' < INTERVAL '2.000000000 SECS')",
     );
 }
+
+#[test]
+fn test_without_offset() {
+    sql_round_trip(
+        MySqlDialect{},
+        "select 1",
+        "SELECT 1",
+    );
+}
+
+#[test]
+fn test_with_offset0() {
+    sql_round_trip(
+        MySqlDialect{},
+        "select 1 offset 0",
+        "SELECT 1",
+    );
+}
+
+#[test]
+fn test_with_offset95() {
+    sql_round_trip(
+        MySqlDialect{},
+        "select 1 offset 95",
+        "SELECT 1 OFFSET 95",
+    );
+}