polyglot_sql/dialects/
databricks.rs

1//! Databricks Dialect
2//!
3//! Databricks-specific transformations based on sqlglot patterns.
4//! Databricks extends Spark SQL with additional features:
5//! - Colon operator for JSON extraction (col:path)
6//! - DATEADD/DATEDIFF with specific syntax
7//! - NULL type mapped to VOID
8//! - Native REGEXP_LIKE and TRY_CAST support
9
10use super::{DialectImpl, DialectType};
11use crate::error::Result;
12use crate::expressions::{
13    AggregateFunction, Cast, DataType, DateAddFunc, Expression, Function, IntervalUnit, Literal,
14    UnaryFunc, VarArgFunc,
15};
16#[cfg(feature = "generate")]
17use crate::generator::GeneratorConfig;
18use crate::tokens::TokenizerConfig;
19
20/// Databricks dialect
21pub struct DatabricksDialect;
22
23impl DialectImpl for DatabricksDialect {
24    fn dialect_type(&self) -> DialectType {
25        DialectType::Databricks
26    }
27
28    fn tokenizer_config(&self) -> TokenizerConfig {
29        let mut config = TokenizerConfig::default();
30        // Databricks uses backticks for identifiers (NOT double quotes)
31        config.identifiers.clear();
32        config.identifiers.insert('`', '`');
33        // Databricks (like Hive/Spark) uses double quotes as string delimiters
34        config.quotes.insert("\"".to_string(), "\"".to_string());
35        // Databricks uses backslash escapes in strings (inherited from Hive/Spark)
36        config.string_escapes.push('\\');
37        // Databricks supports DIV keyword for integer division
38        config
39            .keywords
40            .insert("DIV".to_string(), crate::tokens::TokenType::Div);
41        config
42            .keywords
43            .insert("REPAIR".to_string(), crate::tokens::TokenType::Command);
44        config
45            .keywords
46            .insert("MSCK".to_string(), crate::tokens::TokenType::Command);
47        // Databricks numeric literal suffixes (same as Hive/Spark)
48        config
49            .numeric_literals
50            .insert("L".to_string(), "BIGINT".to_string());
51        config
52            .numeric_literals
53            .insert("S".to_string(), "SMALLINT".to_string());
54        config
55            .numeric_literals
56            .insert("Y".to_string(), "TINYINT".to_string());
57        config
58            .numeric_literals
59            .insert("D".to_string(), "DOUBLE".to_string());
60        config
61            .numeric_literals
62            .insert("F".to_string(), "FLOAT".to_string());
63        config
64            .numeric_literals
65            .insert("BD".to_string(), "DECIMAL".to_string());
66        // Databricks allows identifiers to start with digits (like Hive/Spark)
67        config.identifiers_can_start_with_digit = true;
68        // Databricks (like Spark): STRING_ESCAPES_ALLOWED_IN_RAW_STRINGS = False
69        // Backslashes in raw strings are always literal (no escape processing)
70        config.string_escapes_allowed_in_raw_strings = false;
71        config
72    }
73
74    #[cfg(feature = "generate")]
75
76    fn generator_config(&self) -> GeneratorConfig {
77        use crate::generator::IdentifierQuoteStyle;
78        GeneratorConfig {
79            identifier_quote: '`',
80            identifier_quote_style: IdentifierQuoteStyle::BACKTICK,
81            dialect: Some(DialectType::Databricks),
82            struct_field_sep: ": ",
83            create_function_return_as: false,
84            tablesample_seed_keyword: "REPEATABLE",
85            identifiers_can_start_with_digit: true,
86            // Databricks uses COMMENT 'value' without = sign
87            schema_comment_with_eq: false,
88            ..Default::default()
89        }
90    }
91
92    #[cfg(feature = "transpile")]
93
94    fn transform_expr(&self, expr: Expression) -> Result<Expression> {
95        match expr {
96            // IFNULL -> COALESCE in Databricks
97            Expression::IfNull(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
98                original_name: None,
99                expressions: vec![f.this, f.expression],
100                inferred_type: None,
101            }))),
102
103            // NVL -> COALESCE in Databricks
104            Expression::Nvl(f) => Ok(Expression::Coalesce(Box::new(VarArgFunc {
105                original_name: None,
106                expressions: vec![f.this, f.expression],
107                inferred_type: None,
108            }))),
109
110            // TryCast is native in Databricks
111            Expression::TryCast(c) => Ok(Expression::TryCast(c)),
112
113            // SafeCast -> TRY_CAST in Databricks
114            Expression::SafeCast(c) => Ok(Expression::TryCast(c)),
115
116            // ILIKE is native in Databricks (Spark 3+)
117            Expression::ILike(op) => Ok(Expression::ILike(op)),
118
119            // UNNEST -> EXPLODE in Databricks
120            Expression::Unnest(f) => Ok(Expression::Explode(Box::new(UnaryFunc::new(f.this)))),
121
122            // EXPLODE is native to Databricks
123            Expression::Explode(f) => Ok(Expression::Explode(f)),
124
125            // ExplodeOuter is supported
126            Expression::ExplodeOuter(f) => Ok(Expression::ExplodeOuter(f)),
127
128            // RANDOM -> RAND in Databricks
129            Expression::Random(_) => Ok(Expression::Rand(Box::new(crate::expressions::Rand {
130                seed: None,
131                lower: None,
132                upper: None,
133            }))),
134
135            // Rand is native
136            Expression::Rand(r) => Ok(Expression::Rand(r)),
137
138            // || (Concat) -> CONCAT in Databricks
139            Expression::Concat(op) => Ok(Expression::Function(Box::new(Function::new(
140                "CONCAT".to_string(),
141                vec![op.left, op.right],
142            )))),
143
144            // RegexpLike is native in Databricks
145            Expression::RegexpLike(op) => Ok(Expression::RegexpLike(op)),
146
147            // Cast with typed literal: TIMESTAMP 'x'::TYPE -> CAST(CAST('x' AS TYPE) AS TIMESTAMP)
148            // This is a complex sqlglot transformation where:
149            // 1. The inner typed literal (e.g., TIMESTAMP 'x') becomes CAST('x' AS <target_type>)
150            // 2. The outer result is wrapped in CAST(... AS <original_literal_type>)
151            Expression::Cast(c) => self.transform_cast(*c),
152
153            // Generic function transformations
154            Expression::Function(f) => self.transform_function(*f),
155
156            // Generic aggregate function transformations
157            Expression::AggregateFunction(f) => self.transform_aggregate_function(f),
158
159            // DateSub -> DATE_ADD(date, -val) in Databricks
160            Expression::DateSub(f) => {
161                // Convert string literals to numbers (interval values are often stored as strings)
162                let val = match f.interval {
163                    Expression::Literal(lit) if matches!(lit.as_ref(), crate::expressions::Literal::String(s) if s.parse::<i64>().is_ok()) =>
164                    {
165                        let crate::expressions::Literal::String(s) = lit.as_ref() else {
166                            unreachable!()
167                        };
168                        Expression::Literal(Box::new(crate::expressions::Literal::Number(
169                            s.clone(),
170                        )))
171                    }
172                    other => other,
173                };
174                let neg_val = Expression::Neg(Box::new(crate::expressions::UnaryOp {
175                    this: val,
176                    inferred_type: None,
177                }));
178                Ok(Expression::Function(Box::new(Function::new(
179                    "DATE_ADD".to_string(),
180                    vec![f.this, neg_val],
181                ))))
182            }
183
184            // DateAdd -> native Databricks DATE_ADD forms.
185            Expression::DateAdd(f) => Ok(Self::transform_date_add(*f)),
186
187            // Pass through everything else
188            _ => Ok(expr),
189        }
190    }
191}
192
193#[cfg(feature = "transpile")]
194impl DatabricksDialect {
195    fn transform_function(&self, f: Function) -> Result<Expression> {
196        let name_upper = f.name.to_uppercase();
197        match name_upper.as_str() {
198            // IFNULL -> COALESCE
199            "IFNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
200                original_name: None,
201                expressions: f.args,
202                inferred_type: None,
203            }))),
204
205            // NVL -> COALESCE
206            "NVL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
207                original_name: None,
208                expressions: f.args,
209                inferred_type: None,
210            }))),
211
212            // ISNULL -> COALESCE
213            "ISNULL" if f.args.len() == 2 => Ok(Expression::Coalesce(Box::new(VarArgFunc {
214                original_name: None,
215                expressions: f.args,
216                inferred_type: None,
217            }))),
218
219            // ROW -> STRUCT (no auto-naming for cross-dialect conversion)
220            "ROW" => Ok(Expression::Function(Box::new(Function::new(
221                "STRUCT".to_string(),
222                f.args,
223            )))),
224
225            // NAMED_STRUCT('a', 1) -> STRUCT(1 AS a) for SQLGlot Databricks outputs
226            "NAMED_STRUCT" if f.args.len() % 2 == 0 => {
227                let original_args = f.args.clone();
228                let mut struct_args = Vec::new();
229                for pair in f.args.chunks(2) {
230                    if let Expression::Literal(lit) = &pair[0] {
231                        if let Literal::String(field_name) = lit.as_ref() {
232                            struct_args.push(Expression::Alias(Box::new(
233                                crate::expressions::Alias {
234                                    this: pair[1].clone(),
235                                    alias: crate::expressions::Identifier::new(field_name),
236                                    column_aliases: Vec::new(),
237                                    alias_explicit_as: false,
238                                    alias_keyword: None,
239                                    pre_alias_comments: Vec::new(),
240                                    trailing_comments: Vec::new(),
241                                    inferred_type: None,
242                                },
243                            )));
244                            continue;
245                        }
246                    }
247                    return Ok(Expression::Function(Box::new(Function::new(
248                        "NAMED_STRUCT".to_string(),
249                        original_args,
250                    ))));
251                }
252                Ok(Expression::Function(Box::new(Function::new(
253                    "STRUCT".to_string(),
254                    struct_args,
255                ))))
256            }
257
258            // GETDATE -> CURRENT_TIMESTAMP
259            "GETDATE" => Ok(Expression::CurrentTimestamp(
260                crate::expressions::CurrentTimestamp {
261                    precision: None,
262                    sysdate: false,
263                },
264            )),
265
266            // NOW -> CURRENT_TIMESTAMP
267            "NOW" => Ok(Expression::CurrentTimestamp(
268                crate::expressions::CurrentTimestamp {
269                    precision: None,
270                    sysdate: false,
271                },
272            )),
273
274            // CURDATE -> CURRENT_DATE
275            "CURDATE" => Ok(Expression::CurrentDate(crate::expressions::CurrentDate)),
276
277            // CURRENT_DATE() with parens -> CURRENT_DATE (no parens)
278            "CURRENT_DATE" if f.args.is_empty() => {
279                Ok(Expression::CurrentDate(crate::expressions::CurrentDate))
280            }
281
282            // RANDOM -> RAND
283            "RANDOM" => Ok(Expression::Rand(Box::new(crate::expressions::Rand {
284                seed: None,
285                lower: None,
286                upper: None,
287            }))),
288
289            // GROUP_CONCAT -> COLLECT_LIST + ARRAY_JOIN
290            "GROUP_CONCAT" if !f.args.is_empty() => {
291                let mut args = f.args;
292                let first = args.remove(0);
293                let separator = args.pop();
294                let collect_list = Expression::Function(Box::new(Function::new(
295                    "COLLECT_LIST".to_string(),
296                    vec![first],
297                )));
298                if let Some(sep) = separator {
299                    Ok(Expression::Function(Box::new(Function::new(
300                        "ARRAY_JOIN".to_string(),
301                        vec![collect_list, sep],
302                    ))))
303                } else {
304                    Ok(Expression::Function(Box::new(Function::new(
305                        "ARRAY_JOIN".to_string(),
306                        vec![collect_list],
307                    ))))
308                }
309            }
310
311            // STRING_AGG -> COLLECT_LIST + ARRAY_JOIN in Databricks
312            "STRING_AGG" if !f.args.is_empty() => {
313                let mut args = f.args;
314                let first = args.remove(0);
315                let separator = args.pop();
316                let collect_list = Expression::Function(Box::new(Function::new(
317                    "COLLECT_LIST".to_string(),
318                    vec![first],
319                )));
320                if let Some(sep) = separator {
321                    Ok(Expression::Function(Box::new(Function::new(
322                        "ARRAY_JOIN".to_string(),
323                        vec![collect_list, sep],
324                    ))))
325                } else {
326                    Ok(Expression::Function(Box::new(Function::new(
327                        "ARRAY_JOIN".to_string(),
328                        vec![collect_list],
329                    ))))
330                }
331            }
332
333            // LISTAGG -> COLLECT_LIST + ARRAY_JOIN
334            "LISTAGG" if !f.args.is_empty() => {
335                let mut args = f.args;
336                let first = args.remove(0);
337                let separator = args.pop();
338                let collect_list = Expression::Function(Box::new(Function::new(
339                    "COLLECT_LIST".to_string(),
340                    vec![first],
341                )));
342                if let Some(sep) = separator {
343                    Ok(Expression::Function(Box::new(Function::new(
344                        "ARRAY_JOIN".to_string(),
345                        vec![collect_list, sep],
346                    ))))
347                } else {
348                    Ok(Expression::Function(Box::new(Function::new(
349                        "ARRAY_JOIN".to_string(),
350                        vec![collect_list],
351                    ))))
352                }
353            }
354
355            // ARRAY_AGG -> COLLECT_LIST in Databricks
356            "ARRAY_AGG" if !f.args.is_empty() => Ok(Expression::Function(Box::new(Function::new(
357                "COLLECT_LIST".to_string(),
358                f.args,
359            )))),
360
361            // SUBSTR -> SUBSTRING
362            "SUBSTR" => Ok(Expression::Function(Box::new(Function::new(
363                "SUBSTRING".to_string(),
364                f.args,
365            )))),
366
367            // LEN -> LENGTH
368            "LEN" if f.args.len() == 1 => Ok(Expression::Length(Box::new(UnaryFunc::new(
369                f.args.into_iter().next().unwrap(),
370            )))),
371
372            // CHARINDEX -> LOCATE (with swapped args, like Spark)
373            "CHARINDEX" if f.args.len() >= 2 => {
374                let mut args = f.args;
375                let substring = args.remove(0);
376                let string = args.remove(0);
377                // LOCATE(substring, string)
378                Ok(Expression::Function(Box::new(Function::new(
379                    "LOCATE".to_string(),
380                    vec![substring, string],
381                ))))
382            }
383
384            // POSITION -> LOCATE
385            "POSITION" if f.args.len() == 2 => {
386                let args = f.args;
387                Ok(Expression::Function(Box::new(Function::new(
388                    "LOCATE".to_string(),
389                    args,
390                ))))
391            }
392
393            // STRPOS -> LOCATE (with same arg order)
394            "STRPOS" if f.args.len() == 2 => {
395                let args = f.args;
396                let string = args[0].clone();
397                let substring = args[1].clone();
398                // STRPOS(string, substring) -> LOCATE(substring, string)
399                Ok(Expression::Function(Box::new(Function::new(
400                    "LOCATE".to_string(),
401                    vec![substring, string],
402                ))))
403            }
404
405            // INSTR is native in Databricks
406            "INSTR" => Ok(Expression::Function(Box::new(f))),
407
408            // LOCATE is native in Databricks
409            "LOCATE" => Ok(Expression::Function(Box::new(f))),
410
411            // ARRAY_LENGTH -> SIZE
412            "ARRAY_LENGTH" if f.args.len() == 1 => Ok(Expression::Function(Box::new(
413                Function::new("SIZE".to_string(), f.args),
414            ))),
415
416            // CARDINALITY -> SIZE
417            "CARDINALITY" if f.args.len() == 1 => Ok(Expression::Function(Box::new(
418                Function::new("SIZE".to_string(), f.args),
419            ))),
420
421            // SIZE is native
422            "SIZE" => Ok(Expression::Function(Box::new(f))),
423
424            // ARRAY_CONTAINS is native in Databricks
425            "ARRAY_CONTAINS" => Ok(Expression::Function(Box::new(f))),
426
427            // CONTAINS -> ARRAY_CONTAINS in Databricks (for array operations)
428            // But keep CONTAINS for string contains (from CONTAINS_SUBSTR transpilation)
429            "CONTAINS" if f.args.len() == 2 => {
430                // Check if this is a string CONTAINS (LOWER() args pattern from CONTAINS_SUBSTR)
431                let is_string_contains = matches!(&f.args[0], Expression::Lower(_))
432                    && matches!(&f.args[1], Expression::Lower(_));
433                if is_string_contains {
434                    Ok(Expression::Function(Box::new(f)))
435                } else {
436                    Ok(Expression::Function(Box::new(Function::new(
437                        "ARRAY_CONTAINS".to_string(),
438                        f.args,
439                    ))))
440                }
441            }
442
443            // TO_DATE is native in Databricks
444            "TO_DATE" => Ok(Expression::Function(Box::new(f))),
445
446            // TO_TIMESTAMP is native in Databricks
447            "TO_TIMESTAMP" => Ok(Expression::Function(Box::new(f))),
448
449            // DATE_FORMAT is native in Databricks
450            "DATE_FORMAT" => Ok(Expression::Function(Box::new(f))),
451
452            // strftime -> DATE_FORMAT in Databricks
453            "STRFTIME" if f.args.len() >= 2 => {
454                let mut args = f.args;
455                let format = args.remove(0);
456                let date = args.remove(0);
457                Ok(Expression::Function(Box::new(Function::new(
458                    "DATE_FORMAT".to_string(),
459                    vec![date, format],
460                ))))
461            }
462
463            // TO_CHAR is supported natively in Databricks (unlike Spark)
464            "TO_CHAR" => Ok(Expression::Function(Box::new(f))),
465
466            // DATE_TRUNC is native in Databricks
467            "DATE_TRUNC" => Ok(Expression::Function(Box::new(f))),
468
469            // DATEADD normalizes to Databricks DATE_ADD with an uppercased unit.
470            "DATEADD" => {
471                if f.args.len() == 2 {
472                    Ok(Expression::Function(Box::new(Function::new(
473                        "DATE_ADD".to_string(),
474                        f.args,
475                    ))))
476                } else {
477                    let transformed_args = self.uppercase_first_arg_if_identifier(f.args);
478                    let function_name = if matches!(
479                        transformed_args.first(),
480                        Some(Expression::Identifier(unit))
481                            if unit.name.eq_ignore_ascii_case("WEEK")
482                    ) {
483                        "DATEADD"
484                    } else {
485                        "DATE_ADD"
486                    };
487                    Ok(Expression::Function(Box::new(Function::new(
488                        function_name.to_string(),
489                        transformed_args,
490                    ))))
491                }
492            }
493
494            // DATE_ADD -> DATEADD in Databricks (2-arg form only)
495            // 2-arg with interval: DATE_ADD(date, interval) -> DATEADD(DAY, interval, date)
496            // 2-arg with number: DATE_ADD(date, -2) -> keep as DATE_ADD(date, -2)
497            // 3-arg: DATE_ADD(unit, amount, date) -> keep as DATE_ADD(UNIT, amount, date)
498            "DATE_ADD" => {
499                if f.args.len() == 2 {
500                    let is_simple_number = matches!(
501                        &f.args[1],
502                        Expression::Literal(lit) if matches!(lit.as_ref(), crate::expressions::Literal::Number(_))
503                    ) || matches!(&f.args[1], Expression::Neg(_));
504                    if is_simple_number {
505                        // Keep as DATE_ADD(date, num_days)
506                        Ok(Expression::Function(Box::new(Function::new(
507                            "DATE_ADD".to_string(),
508                            f.args,
509                        ))))
510                    } else {
511                        let mut args = f.args;
512                        let date = args.remove(0);
513                        let interval = args.remove(0);
514                        let unit = Expression::Identifier(crate::expressions::Identifier {
515                            name: "DAY".to_string(),
516                            quoted: false,
517                            trailing_comments: Vec::new(),
518                            span: None,
519                        });
520                        Ok(Expression::Function(Box::new(Function::new(
521                            "DATEADD".to_string(),
522                            vec![unit, interval, date],
523                        ))))
524                    }
525                } else {
526                    let transformed_args = self.uppercase_first_arg_if_identifier(f.args);
527                    Ok(Expression::Function(Box::new(Function::new(
528                        "DATE_ADD".to_string(),
529                        transformed_args,
530                    ))))
531                }
532            }
533
534            // DATEDIFF is native in Databricks - uppercase the unit if present
535            // 2-arg: DATEDIFF(end, start) -> DATEDIFF(DAY, start, end)
536            // 3-arg: DATEDIFF(unit, start, end) -> DATEDIFF(UNIT, start, end)
537            "DATEDIFF" => {
538                if f.args.len() == 2 {
539                    let mut args = f.args;
540                    let end_date = args.remove(0);
541                    let start_date = args.remove(0);
542                    let unit = Expression::Identifier(crate::expressions::Identifier {
543                        name: "DAY".to_string(),
544                        quoted: false,
545                        trailing_comments: Vec::new(),
546                        span: None,
547                    });
548                    Ok(Expression::Function(Box::new(Function::new(
549                        "DATEDIFF".to_string(),
550                        vec![unit, start_date, end_date],
551                    ))))
552                } else {
553                    let transformed_args = self.uppercase_first_arg_if_identifier(f.args);
554                    Ok(Expression::Function(Box::new(Function::new(
555                        "DATEDIFF".to_string(),
556                        transformed_args,
557                    ))))
558                }
559            }
560
561            // DATE_DIFF -> DATEDIFF with uppercased unit
562            "DATE_DIFF" => {
563                let transformed_args = self.uppercase_first_arg_if_identifier(f.args);
564                Ok(Expression::Function(Box::new(Function::new(
565                    "DATEDIFF".to_string(),
566                    transformed_args,
567                ))))
568            }
569
570            // JSON_EXTRACT -> Use colon operator in generation, but keep as function for now
571            "JSON_EXTRACT" => Ok(Expression::Function(Box::new(f))),
572
573            // JSON_EXTRACT_SCALAR -> same handling
574            "JSON_EXTRACT_SCALAR" => Ok(Expression::Function(Box::new(f))),
575
576            // GET_JSON_OBJECT is native in Databricks.
577            "GET_JSON_OBJECT" if f.args.len() == 2 => {
578                let mut args = f.args;
579                let col = args.remove(0);
580                let path_arg = args.remove(0);
581
582                Ok(Expression::Function(Box::new(Function::new(
583                    "GET_JSON_OBJECT".to_string(),
584                    vec![col, self.normalize_get_json_object_path(path_arg)],
585                ))))
586            }
587
588            // FROM_JSON is native in Databricks
589            "FROM_JSON" => Ok(Expression::Function(Box::new(f))),
590
591            // PARSE_JSON is native in Databricks
592            "PARSE_JSON" => Ok(Expression::Function(Box::new(f))),
593
594            // COLLECT_LIST is native in Databricks
595            "COLLECT_LIST" => Ok(Expression::Function(Box::new(f))),
596
597            // COLLECT_SET is native in Databricks
598            "COLLECT_SET" => Ok(Expression::Function(Box::new(f))),
599
600            // RLIKE is native in Databricks
601            "RLIKE" => Ok(Expression::Function(Box::new(f))),
602
603            // REGEXP -> RLIKE in Databricks
604            "REGEXP" if f.args.len() == 2 => Ok(Expression::Function(Box::new(Function::new(
605                "RLIKE".to_string(),
606                f.args,
607            )))),
608
609            // REGEXP_LIKE is native in Databricks
610            "REGEXP_LIKE" => Ok(Expression::Function(Box::new(f))),
611
612            // LEVENSHTEIN is native in Databricks
613            "LEVENSHTEIN" => Ok(Expression::Function(Box::new(f))),
614
615            // SEQUENCE is native (for GENERATE_SERIES)
616            "GENERATE_SERIES" if !f.args.is_empty() => Ok(Expression::Function(Box::new(
617                Function::new("SEQUENCE".to_string(), f.args),
618            ))),
619
620            // SEQUENCE is native
621            "SEQUENCE" => Ok(Expression::Function(Box::new(f))),
622
623            // FLATTEN is native in Databricks
624            "FLATTEN" => Ok(Expression::Function(Box::new(f))),
625
626            // ARRAY_SORT is native
627            "ARRAY_SORT" => Ok(Expression::Function(Box::new(f))),
628
629            // ARRAY_DISTINCT is native
630            "ARRAY_DISTINCT" => Ok(Expression::Function(Box::new(f))),
631
632            // TRANSFORM is native (for array transformation)
633            "TRANSFORM" => Ok(Expression::Function(Box::new(f))),
634
635            // FILTER is native (for array filtering)
636            "FILTER" => Ok(Expression::Function(Box::new(f))),
637
638            // FROM_UTC_TIMESTAMP - wrap first argument in CAST(... AS TIMESTAMP) if not already
639            "FROM_UTC_TIMESTAMP" if f.args.len() >= 2 => {
640                let mut args = f.args;
641                let first_arg = args.remove(0);
642
643                // Check if first arg is already a Cast to TIMESTAMP
644                let wrapped_arg = if self.is_cast_to_timestamp(&first_arg) {
645                    first_arg
646                } else {
647                    // Wrap in CAST(... AS TIMESTAMP)
648                    Expression::Cast(Box::new(Cast {
649                        this: first_arg,
650                        to: DataType::Timestamp {
651                            precision: None,
652                            timezone: false,
653                        },
654                        trailing_comments: Vec::new(),
655                        double_colon_syntax: false,
656                        format: None,
657                        default: None,
658                        inferred_type: None,
659                    }))
660                };
661
662                let mut new_args = vec![wrapped_arg];
663                new_args.extend(args);
664
665                Ok(Expression::Function(Box::new(Function::new(
666                    "FROM_UTC_TIMESTAMP".to_string(),
667                    new_args,
668                ))))
669            }
670
671            // UNIFORM(low, high, RANDOM(seed)) -> UNIFORM(low, high, seed) or UNIFORM(low, high)
672            "UNIFORM" if f.args.len() == 3 => {
673                let mut args = f.args;
674                let low = args.remove(0);
675                let high = args.remove(0);
676                let gen = args.remove(0);
677                match gen {
678                    Expression::Function(func) if func.name.to_uppercase() == "RANDOM" => {
679                        if func.args.len() == 1 {
680                            // RANDOM(seed) -> extract seed
681                            let seed = func.args.into_iter().next().unwrap();
682                            Ok(Expression::Function(Box::new(Function::new(
683                                "UNIFORM".to_string(),
684                                vec![low, high, seed],
685                            ))))
686                        } else {
687                            // RANDOM() -> drop gen arg
688                            Ok(Expression::Function(Box::new(Function::new(
689                                "UNIFORM".to_string(),
690                                vec![low, high],
691                            ))))
692                        }
693                    }
694                    Expression::Rand(r) => {
695                        if let Some(seed) = r.seed {
696                            Ok(Expression::Function(Box::new(Function::new(
697                                "UNIFORM".to_string(),
698                                vec![low, high, *seed],
699                            ))))
700                        } else {
701                            Ok(Expression::Function(Box::new(Function::new(
702                                "UNIFORM".to_string(),
703                                vec![low, high],
704                            ))))
705                        }
706                    }
707                    _ => Ok(Expression::Function(Box::new(Function::new(
708                        "UNIFORM".to_string(),
709                        vec![low, high, gen],
710                    )))),
711                }
712            }
713
714            // REGEXP_SUBSTR(subject, pattern, ...) -> REGEXP_EXTRACT(subject, pattern)
715            "REGEXP_SUBSTR" if f.args.len() >= 2 => {
716                let subject = f.args[0].clone();
717                let pattern = f.args[1].clone();
718                Ok(Expression::Function(Box::new(Function::new(
719                    "REGEXP_EXTRACT".to_string(),
720                    vec![subject, pattern],
721                ))))
722            }
723
724            // BIT_GET -> GETBIT
725            "BIT_GET" => Ok(Expression::Function(Box::new(Function::new(
726                "GETBIT".to_string(),
727                f.args,
728            )))),
729
730            // Pass through everything else
731            _ => Ok(Expression::Function(Box::new(f))),
732        }
733    }
734
735    fn transform_aggregate_function(
736        &self,
737        f: Box<crate::expressions::AggregateFunction>,
738    ) -> Result<Expression> {
739        let name_upper = f.name.to_uppercase();
740        match name_upper.as_str() {
741            // COUNT_IF is native in Databricks (Spark 3+)
742            "COUNT_IF" => Ok(Expression::AggregateFunction(f)),
743
744            // ANY_VALUE is native in Databricks (Spark 3+)
745            "ANY_VALUE" => Ok(Expression::AggregateFunction(f)),
746
747            // GROUP_CONCAT -> COLLECT_LIST + ARRAY_JOIN
748            "GROUP_CONCAT" if !f.args.is_empty() => {
749                let mut args = f.args;
750                let first = args.remove(0);
751                let separator = args.pop();
752                let collect_list = Expression::Function(Box::new(Function::new(
753                    "COLLECT_LIST".to_string(),
754                    vec![first],
755                )));
756                if let Some(sep) = separator {
757                    Ok(Expression::Function(Box::new(Function::new(
758                        "ARRAY_JOIN".to_string(),
759                        vec![collect_list, sep],
760                    ))))
761                } else {
762                    Ok(Expression::Function(Box::new(Function::new(
763                        "ARRAY_JOIN".to_string(),
764                        vec![collect_list],
765                    ))))
766                }
767            }
768
769            // STRING_AGG -> COLLECT_LIST + ARRAY_JOIN
770            "STRING_AGG" if !f.args.is_empty() => {
771                let mut args = f.args;
772                let first = args.remove(0);
773                let separator = args.pop();
774                let collect_list = Expression::Function(Box::new(Function::new(
775                    "COLLECT_LIST".to_string(),
776                    vec![first],
777                )));
778                if let Some(sep) = separator {
779                    Ok(Expression::Function(Box::new(Function::new(
780                        "ARRAY_JOIN".to_string(),
781                        vec![collect_list, sep],
782                    ))))
783                } else {
784                    Ok(Expression::Function(Box::new(Function::new(
785                        "ARRAY_JOIN".to_string(),
786                        vec![collect_list],
787                    ))))
788                }
789            }
790
791            // LISTAGG -> COLLECT_LIST + ARRAY_JOIN
792            "LISTAGG" if !f.args.is_empty() => {
793                let mut args = f.args;
794                let first = args.remove(0);
795                let separator = args.pop();
796                let collect_list = Expression::Function(Box::new(Function::new(
797                    "COLLECT_LIST".to_string(),
798                    vec![first],
799                )));
800                if let Some(sep) = separator {
801                    Ok(Expression::Function(Box::new(Function::new(
802                        "ARRAY_JOIN".to_string(),
803                        vec![collect_list, sep],
804                    ))))
805                } else {
806                    Ok(Expression::Function(Box::new(Function::new(
807                        "ARRAY_JOIN".to_string(),
808                        vec![collect_list],
809                    ))))
810                }
811            }
812
813            // ARRAY_AGG -> COLLECT_LIST
814            "ARRAY_AGG" if !f.args.is_empty() => Ok(Expression::Function(Box::new(Function::new(
815                "COLLECT_LIST".to_string(),
816                f.args,
817            )))),
818
819            // STDDEV is native in Databricks
820            "STDDEV" => Ok(Expression::AggregateFunction(f)),
821
822            // VARIANCE is native in Databricks
823            "VARIANCE" => Ok(Expression::AggregateFunction(f)),
824
825            // APPROX_COUNT_DISTINCT is native in Databricks
826            "APPROX_COUNT_DISTINCT" => Ok(Expression::AggregateFunction(f)),
827
828            // APPROX_DISTINCT -> APPROX_COUNT_DISTINCT
829            "APPROX_DISTINCT" if !f.args.is_empty() => {
830                Ok(Expression::AggregateFunction(Box::new(AggregateFunction {
831                    source_dialect: None,
832
833                    name: "APPROX_COUNT_DISTINCT".to_string(),
834                    args: f.args,
835                    distinct: f.distinct,
836                    filter: f.filter,
837                    order_by: Vec::new(),
838                    limit: None,
839                    ignore_nulls: None,
840                    inferred_type: None,
841                })))
842            }
843
844            // Pass through everything else
845            _ => Ok(Expression::AggregateFunction(f)),
846        }
847    }
848
849    /// Transform Cast expressions - handles typed literals being cast
850    ///
851    /// When we have a typed literal (TIMESTAMP 'x', DATE 'x', TIME 'x') being cast to another type,
852    /// Databricks/Spark transforms it as follows:
853    ///
854    /// `TIMESTAMP 'x'::TYPE` -> `CAST(CAST('x' AS TYPE) AS TIMESTAMP)`
855    ///
856    /// This reverses the types - the inner cast is to the target type,
857    /// the outer cast is to the original literal type.
858    fn transform_cast(&self, c: Cast) -> Result<Expression> {
859        // Check if the inner expression is a typed literal
860        match &c.this {
861            // TIMESTAMP 'value'::TYPE -> CAST(CAST('value' AS TYPE) AS TIMESTAMP)
862            Expression::Literal(lit) if matches!(lit.as_ref(), Literal::Timestamp(_)) => {
863                let Literal::Timestamp(value) = lit.as_ref() else {
864                    unreachable!()
865                };
866                // Create inner cast: CAST('value' AS target_type)
867                let inner_cast = Expression::Cast(Box::new(Cast {
868                    this: Expression::Literal(Box::new(Literal::String(value.clone()))),
869                    to: c.to,
870                    trailing_comments: Vec::new(),
871                    double_colon_syntax: false,
872                    format: None,
873                    default: None,
874                    inferred_type: None,
875                }));
876                // Create outer cast: CAST(inner_cast AS TIMESTAMP)
877                Ok(Expression::Cast(Box::new(Cast {
878                    this: inner_cast,
879                    to: DataType::Timestamp {
880                        precision: None,
881                        timezone: false,
882                    },
883                    trailing_comments: c.trailing_comments,
884                    double_colon_syntax: false,
885                    format: None,
886                    default: None,
887                    inferred_type: None,
888                })))
889            }
890            // DATE 'value'::TYPE -> CAST(CAST('value' AS TYPE) AS DATE)
891            Expression::Literal(lit) if matches!(lit.as_ref(), Literal::Date(_)) => {
892                let Literal::Date(value) = lit.as_ref() else {
893                    unreachable!()
894                };
895                let inner_cast = Expression::Cast(Box::new(Cast {
896                    this: Expression::Literal(Box::new(Literal::String(value.clone()))),
897                    to: c.to,
898                    trailing_comments: Vec::new(),
899                    double_colon_syntax: false,
900                    format: None,
901                    default: None,
902                    inferred_type: None,
903                }));
904                Ok(Expression::Cast(Box::new(Cast {
905                    this: inner_cast,
906                    to: DataType::Date,
907                    trailing_comments: c.trailing_comments,
908                    double_colon_syntax: false,
909                    format: None,
910                    default: None,
911                    inferred_type: None,
912                })))
913            }
914            // TIME 'value'::TYPE -> CAST(CAST('value' AS TYPE) AS TIME)
915            Expression::Literal(lit) if matches!(lit.as_ref(), Literal::Time(_)) => {
916                let Literal::Time(value) = lit.as_ref() else {
917                    unreachable!()
918                };
919                let inner_cast = Expression::Cast(Box::new(Cast {
920                    this: Expression::Literal(Box::new(Literal::String(value.clone()))),
921                    to: c.to,
922                    trailing_comments: Vec::new(),
923                    double_colon_syntax: false,
924                    format: None,
925                    default: None,
926                    inferred_type: None,
927                }));
928                Ok(Expression::Cast(Box::new(Cast {
929                    this: inner_cast,
930                    to: DataType::Time {
931                        precision: None,
932                        timezone: false,
933                    },
934                    trailing_comments: c.trailing_comments,
935                    double_colon_syntax: false,
936                    format: None,
937                    default: None,
938                    inferred_type: None,
939                })))
940            }
941            // For all other cases, pass through the Cast unchanged
942            _ => Ok(Expression::Cast(Box::new(c))),
943        }
944    }
945
946    /// Check if an expression is a CAST to TIMESTAMP
947    fn is_cast_to_timestamp(&self, expr: &Expression) -> bool {
948        if let Expression::Cast(cast) = expr {
949            matches!(cast.to, DataType::Timestamp { .. })
950        } else {
951            false
952        }
953    }
954
955    /// Helper to uppercase the first argument if it's an identifier or column (for DATEDIFF, DATEADD units)
956    fn uppercase_first_arg_if_identifier(&self, mut args: Vec<Expression>) -> Vec<Expression> {
957        use crate::expressions::Identifier;
958        if !args.is_empty() {
959            match &args[0] {
960                Expression::Identifier(id) => {
961                    args[0] = Expression::Identifier(Identifier {
962                        name: id.name.to_uppercase(),
963                        quoted: id.quoted,
964                        trailing_comments: id.trailing_comments.clone(),
965                        span: None,
966                    });
967                }
968                Expression::Var(v) => {
969                    args[0] = Expression::Identifier(Identifier {
970                        name: v.this.to_uppercase(),
971                        quoted: false,
972                        trailing_comments: Vec::new(),
973                        span: None,
974                    });
975                }
976                Expression::Column(col) if col.table.is_none() => {
977                    // Unqualified column name like "day" should be treated as a unit
978                    args[0] = Expression::Identifier(Identifier {
979                        name: col.name.name.to_uppercase(),
980                        quoted: col.name.quoted,
981                        trailing_comments: col.name.trailing_comments.clone(),
982                        span: None,
983                    });
984                }
985                _ => {}
986            }
987        }
988        args
989    }
990
991    fn normalize_get_json_object_path(&self, path_arg: Expression) -> Expression {
992        let Expression::Literal(lit) = &path_arg else {
993            return path_arg;
994        };
995        let crate::expressions::Literal::String(path) = lit.as_ref() else {
996            return path_arg;
997        };
998
999        let Some(segment) = path.strip_prefix("$.") else {
1000            return path_arg;
1001        };
1002
1003        if segment.is_empty()
1004            || segment.contains('.')
1005            || segment.contains('[')
1006            || segment
1007                .chars()
1008                .all(|c| c.is_ascii_alphanumeric() || c == '_')
1009        {
1010            return path_arg;
1011        }
1012
1013        Expression::Literal(Box::new(crate::expressions::Literal::String(format!(
1014            "$[\"{}\"]",
1015            segment.replace('"', "\\\"")
1016        ))))
1017    }
1018
1019    fn transform_date_add(f: DateAddFunc) -> Expression {
1020        if f.unit == IntervalUnit::Day {
1021            Expression::Function(Box::new(Function::new(
1022                "DATE_ADD".to_string(),
1023                vec![f.this, f.interval],
1024            )))
1025        } else {
1026            Expression::Function(Box::new(Function::new(
1027                "DATE_ADD".to_string(),
1028                vec![
1029                    Expression::Identifier(crate::expressions::Identifier {
1030                        name: Self::interval_unit_name(f.unit).to_string(),
1031                        quoted: false,
1032                        trailing_comments: Vec::new(),
1033                        span: None,
1034                    }),
1035                    f.interval,
1036                    f.this,
1037                ],
1038            )))
1039        }
1040    }
1041
1042    fn interval_unit_name(unit: IntervalUnit) -> &'static str {
1043        match unit {
1044            IntervalUnit::Year => "YEAR",
1045            IntervalUnit::Quarter => "QUARTER",
1046            IntervalUnit::Month => "MONTH",
1047            IntervalUnit::Week => "WEEK",
1048            IntervalUnit::Day => "DAY",
1049            IntervalUnit::Hour => "HOUR",
1050            IntervalUnit::Minute => "MINUTE",
1051            IntervalUnit::Second => "SECOND",
1052            IntervalUnit::Millisecond => "MILLISECOND",
1053            IntervalUnit::Microsecond => "MICROSECOND",
1054            IntervalUnit::Nanosecond => "NANOSECOND",
1055        }
1056    }
1057}
1058
1059#[cfg(test)]
1060mod tests {
1061    use super::*;
1062    use crate::Dialect;
1063
1064    #[test]
1065    fn test_timestamp_literal_cast() {
1066        // TIMESTAMP 'value'::DATE -> CAST(CAST('value' AS DATE) AS TIMESTAMP)
1067        // This is test [47] in the Databricks dialect identity fixtures
1068        let sql = "SELECT TIMESTAMP '2025-04-29 18.47.18'::DATE";
1069        let expected = "SELECT CAST(CAST('2025-04-29 18.47.18' AS DATE) AS TIMESTAMP)";
1070
1071        let d = Dialect::get(DialectType::Databricks);
1072        let ast = d.parse(sql).expect("Parse failed");
1073        let transformed = d.transform(ast[0].clone()).expect("Transform failed");
1074        let output = d.generate(&transformed).expect("Generate failed");
1075
1076        assert_eq!(
1077            output, expected,
1078            "Timestamp literal cast transformation failed"
1079        );
1080    }
1081
1082    #[test]
1083    fn test_from_utc_timestamp_wraps_column() {
1084        // Test [48]: FROM_UTC_TIMESTAMP(foo, 'timezone') -> FROM_UTC_TIMESTAMP(CAST(foo AS TIMESTAMP), 'timezone')
1085        let sql = "SELECT DATE_FORMAT(CAST(FROM_UTC_TIMESTAMP(foo, 'America/Los_Angeles') AS TIMESTAMP), 'yyyy-MM-dd HH:mm:ss') AS foo FROM t";
1086        let expected = "SELECT DATE_FORMAT(CAST(FROM_UTC_TIMESTAMP(CAST(foo AS TIMESTAMP), 'America/Los_Angeles') AS TIMESTAMP), 'yyyy-MM-dd HH:mm:ss') AS foo FROM t";
1087
1088        let d = Dialect::get(DialectType::Databricks);
1089        let ast = d.parse(sql).expect("Parse failed");
1090        let transformed = d.transform(ast[0].clone()).expect("Transform failed");
1091        let output = d.generate(&transformed).expect("Generate failed");
1092
1093        assert_eq!(output, expected, "FROM_UTC_TIMESTAMP transformation failed");
1094    }
1095
1096    #[test]
1097    fn test_from_utc_timestamp_keeps_existing_cast() {
1098        // Test [50]: FROM_UTC_TIMESTAMP(x::TIMESTAMP, tz) -> FROM_UTC_TIMESTAMP(CAST(x AS TIMESTAMP), tz)
1099        // When already cast to TIMESTAMP, keep it but convert :: syntax to CAST()
1100        let sql = "FROM_UTC_TIMESTAMP(x::TIMESTAMP, tz)";
1101        let expected = "FROM_UTC_TIMESTAMP(CAST(x AS TIMESTAMP), tz)";
1102
1103        let d = Dialect::get(DialectType::Databricks);
1104        let ast = d.parse(sql).expect("Parse failed");
1105        let transformed = d.transform(ast[0].clone()).expect("Transform failed");
1106        let output = d.generate(&transformed).expect("Generate failed");
1107
1108        assert_eq!(
1109            output, expected,
1110            "FROM_UTC_TIMESTAMP with existing CAST failed"
1111        );
1112    }
1113
1114    #[test]
1115    fn test_deep_clone_version_as_of() {
1116        let sql = "CREATE TABLE events_clone DEEP CLONE events VERSION AS OF 5";
1117        let d = Dialect::get(DialectType::Databricks);
1118        let ast = d.parse(sql).expect("Parse failed");
1119        let output = d.generate(&ast[0]).expect("Generate failed");
1120
1121        assert_eq!(output, sql);
1122    }
1123
1124    #[test]
1125    fn test_deep_clone_timestamp_as_of() {
1126        let sql = "CREATE TABLE events_clone DEEP CLONE events TIMESTAMP AS OF '2024-01-01'";
1127        let d = Dialect::get(DialectType::Databricks);
1128        let ast = d.parse(sql).expect("Parse failed");
1129        let output = d.generate(&ast[0]).expect("Generate failed");
1130
1131        assert_eq!(output, sql);
1132    }
1133
1134    #[test]
1135    fn test_shallow_clone_still_roundtrips() {
1136        let sql = "CREATE TABLE events_clone SHALLOW CLONE events";
1137        let d = Dialect::get(DialectType::Databricks);
1138        let ast = d.parse(sql).expect("Parse failed");
1139        let output = d.generate(&ast[0]).expect("Generate failed");
1140
1141        assert_eq!(output, sql);
1142    }
1143
1144    #[test]
1145    fn test_repair_table_commands_roundtrip() {
1146        let d = Dialect::get(DialectType::Databricks);
1147        let cases = [
1148            "REPAIR TABLE events",
1149            "MSCK REPAIR TABLE events",
1150            "REPAIR TABLE events ADD PARTITIONS",
1151            "REPAIR TABLE events DROP PARTITIONS",
1152            "REPAIR TABLE events SYNC PARTITIONS",
1153            "REPAIR TABLE events SYNC METADATA",
1154        ];
1155
1156        for sql in cases {
1157            let ast = d.parse(sql).expect("Parse failed");
1158            let output = d.generate(&ast[0]).expect("Generate failed");
1159            assert_eq!(output, sql);
1160        }
1161    }
1162
1163    #[test]
1164    fn test_apply_changes_commands_roundtrip() {
1165        let d = Dialect::get(DialectType::Databricks);
1166        let cases = [
1167            "APPLY CHANGES INTO silver.orders FROM STREAM(bronze.orders) KEYS (id) SEQUENCE BY ts",
1168            "APPLY CHANGES INTO LIVE.silver_orders FROM STREAM(LIVE.bronze_orders) KEYS (id) IGNORE NULL UPDATES SEQUENCE BY ts",
1169            "APPLY CHANGES INTO LIVE.silver_orders FROM STREAM(LIVE.bronze_orders) KEYS (id) APPLY AS DELETE WHEN operation = 'DELETE' SEQUENCE BY ts COLUMNS * EXCEPT (operation) STORED AS SCD TYPE 1",
1170            "APPLY CHANGES INTO LIVE.silver_orders FROM STREAM(LIVE.bronze_orders) KEYS (id) SEQUENCE BY ts STORED AS SCD TYPE 2 TRACK HISTORY ON * EXCEPT (updated_at)",
1171            "AUTO CDC INTO silver.orders FROM STREAM(bronze.orders) KEYS (id) SEQUENCE BY ts",
1172            "CREATE FLOW apply_cdc AS AUTO CDC INTO silver.orders FROM STREAM(bronze.orders) KEYS (id) SEQUENCE BY ts",
1173        ];
1174
1175        for sql in cases {
1176            let ast = d.parse(sql).expect("Parse failed");
1177            let output = d.generate(&ast[0]).expect("Generate failed");
1178            assert_eq!(output, sql);
1179        }
1180    }
1181
1182    #[test]
1183    fn test_generate_symlink_format_manifest_roundtrip() {
1184        let d = Dialect::get(DialectType::Databricks);
1185        let cases = [
1186            "GENERATE symlink_format_manifest FOR TABLE events",
1187            "GENERATE symlink_format_manifest FOR TABLE catalog.schema.events",
1188        ];
1189
1190        for sql in cases {
1191            let ast = d.parse(sql).expect("Parse failed");
1192            let output = d.generate(&ast[0]).expect("Generate failed");
1193            assert_eq!(output, sql);
1194        }
1195    }
1196
1197    #[test]
1198    fn test_convert_to_delta_roundtrip() {
1199        let d = Dialect::get(DialectType::Databricks);
1200        let cases = [
1201            "CONVERT TO DELTA parquet.`/mnt/data/events`",
1202            "CONVERT TO DELTA database_name.table_name",
1203            "CONVERT TO DELTA parquet.`s3://my-bucket/path/to/table` PARTITIONED BY (date DATE)",
1204            "CONVERT TO DELTA database_name.table_name NO STATISTICS",
1205        ];
1206
1207        for sql in cases {
1208            let ast = d.parse(sql).expect("Parse failed");
1209            let output = d.generate(&ast[0]).expect("Generate failed");
1210            assert_eq!(output, sql);
1211        }
1212    }
1213}
polyglot_sql/dialects/databricks.rs

polyglot_sql/dialects/
databricks.rs
