Add type to ast::Lit::EmbeddedDocLit (#525)

partiql · Dec 5, 2024 · 27d90d1 · 27d90d1 · github-actions · Dec 6, 2024
1 parent d07a7e1
commit 27d90d1
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 8 deletions.
diff --git a/extension/partiql-extension-visualize/src/ast_to_dot.rs b/extension/partiql-extension-visualize/src/ast_to_dot.rs
@@ -185,7 +185,7 @@ fn lit_to_str(ast: &ast::Lit) -> String {
         Lit::FloatLit(l) => l.to_string(),
         Lit::DoubleLit(l) => l.to_string(),
         Lit::BoolLit(l) => (if *l { "TRUE" } else { "FALSE" }).to_string(),
-        Lit::EmbeddedDocLit(l) => format!("`{}`", l),
+        Lit::EmbeddedDocLit(l, typ) => format!("`{}`::{}", l, type_to_str(typ)),
         Lit::CharStringLit(l) => format!("'{}'", l),
         Lit::NationalCharStringLit(l) => format!("'{}'", l),
         Lit::BitStringLit(l) => format!("b'{}'", l),

diff --git a/partiql-ast/src/ast.rs b/partiql-ast/src/ast.rs
@@ -444,7 +444,7 @@ pub enum Lit {
     #[visit(skip)]
     BoolLit(bool),
     #[visit(skip)]
-    EmbeddedDocLit(String),
+    EmbeddedDocLit(String, Type),
     #[visit(skip)]
     CharStringLit(String),
     #[visit(skip)]

diff --git a/partiql-ast/src/pretty.rs b/partiql-ast/src/pretty.rs
@@ -394,7 +394,7 @@ impl PrettyDoc for Lit {
             Lit::FloatLit(inner) => arena.text(inner.to_string()),
             Lit::DoubleLit(inner) => arena.text(inner.to_string()),
             Lit::BoolLit(inner) => arena.text(inner.to_string()),
-            Lit::EmbeddedDocLit(inner) => inner.pretty_doc(arena), // TODO better pretty for embedded doc: https://github.com/partiql/partiql-lang-rust/issues/508
+            Lit::EmbeddedDocLit(inner, _typ) => inner.pretty_doc(arena), // TODO better pretty for embedded doc: https://github.com/partiql/partiql-lang-rust/issues/508
             Lit::CharStringLit(inner) => inner.pretty_doc(arena),
             Lit::NationalCharStringLit(inner) => inner.pretty_doc(arena),
             Lit::BitStringLit(inner) => inner.pretty_doc(arena),
@@ -420,9 +420,32 @@ impl PrettyDoc for Type {
     {
         match self {
             Type::CustomType(cty) => cty.pretty_doc(arena),
-            _ => {
-                todo!("Non-custom type type")
-            }
+            Type::NullType => arena.text("NULL"),
+            Type::BooleanType => arena.text("BOOL"),
+            Type::Integer2Type => arena.text("INT2"),
+            Type::Integer4Type => arena.text("INT4"),
+            Type::Integer8Type => arena.text("INT8"),
+            Type::DecimalType => arena.text("DECIMAL"),
+            Type::NumericType => arena.text("NUMERIC"),
+            Type::RealType => arena.text("REAL"),
+            Type::DoublePrecisionType => arena.text("DOUBLE PRECISION"),
+            Type::TimestampType => arena.text("TIMESTAMP"),
+            Type::CharacterType => arena.text("CHAR"),
+            Type::CharacterVaryingType => arena.text("VARCHAR"),
+            Type::MissingType => arena.text("MISSING"),
+            Type::StringType => arena.text("STRING"),
+            Type::SymbolType => arena.text("SYMBOL"),
+            Type::BlobType => arena.text("BLOB"),
+            Type::ClobType => arena.text("CLOB"),
+            Type::DateType => arena.text("DATE"),
+            Type::TimeType => arena.text("TIME"),
+            Type::ZonedTimestampType => arena.text("TIMESTAMPTZ"),
+            Type::StructType => arena.text("STRUCT"),
+            Type::TupleType => arena.text("TUPLE"),
+            Type::ListType => arena.text("LIST"),
+            Type::SexpType => arena.text("SEXP"),
+            Type::BagType => arena.text("BAG"),
+            Type::AnyType => arena.text("ANY"),
         }
     }
 }

diff --git a/partiql-logical-planner/src/lower.rs b/partiql-logical-planner/src/lower.rs
@@ -1922,7 +1922,7 @@ fn lit_to_lit(lit: &Lit) -> Result<logical::Lit, AstTransformError> {
         Lit::FloatLit(f) => logical::Lit::Double(OrderedFloat::from(*f as f64)),
         Lit::DoubleLit(f) => logical::Lit::Double(OrderedFloat::from(*f)),
         Lit::BoolLit(b) => logical::Lit::Bool(*b),
-        Lit::EmbeddedDocLit(s) => {
+        Lit::EmbeddedDocLit(s, _) => {
             logical::Lit::BoxDocument(s.clone().into_bytes(), "Ion".to_string())
         }
         Lit::CharStringLit(s) => logical::Lit::String(s.clone()),

diff --git a/partiql-parser/src/parse/partiql.lalrpop b/partiql-parser/src/parse/partiql.lalrpop
@@ -1301,7 +1301,15 @@ LiteralNumber: ast::Lit = {
 
 #[inline]
 LiteralEmbeddedDoc: ast::Lit = {
-    <ion:"EmbeddedDoc"> => ast::Lit::EmbeddedDocLit(ion.to_owned()),
+    <ion:"EmbeddedDoc"> => {
+        let ion_typ = ast::Type::CustomType(ast::CustomType {
+            parts: vec![ast::CustomTypePart::Name(ast::SymbolPrimitive {
+                value: "Ion".to_string(),
+                case: ast::CaseSensitivity::CaseInsensitive,
+            })],
+        });
+        ast::Lit::EmbeddedDocLit(ion.to_owned(), ion_typ)
+    }
 }
Benchmark suite	Current: `27d90d1`	Previous: `d07a7e1`	Ratio
`arith_agg-avg`	`762607` ns/iter (`± 13004`)	`770065` ns/iter (`± 6886`)	`0.99`
`arith_agg-avg_distinct`	`853107` ns/iter (`± 13446`)	`856721` ns/iter (`± 3282`)	`1.00`
`arith_agg-count`	`811102` ns/iter (`± 52377`)	`815223` ns/iter (`± 14799`)	`0.99`
`arith_agg-count_distinct`	`841861` ns/iter (`± 9876`)	`858669` ns/iter (`± 4273`)	`0.98`
`arith_agg-min`	`813170` ns/iter (`± 11329`)	`821056` ns/iter (`± 2841`)	`0.99`
`arith_agg-min_distinct`	`852586` ns/iter (`± 13113`)	`854399` ns/iter (`± 4373`)	`1.00`
`arith_agg-max`	`822112` ns/iter (`± 18721`)	`825900` ns/iter (`± 41288`)	`1.00`
`arith_agg-max_distinct`	`861600` ns/iter (`± 12184`)	`861984` ns/iter (`± 2669`)	`1.00`
`arith_agg-sum`	`817410` ns/iter (`± 7116`)	`821892` ns/iter (`± 6469`)	`0.99`
`arith_agg-sum_distinct`	`846422` ns/iter (`± 17409`)	`860284` ns/iter (`± 2356`)	`0.98`
`arith_agg-avg-count-min-max-sum`	`980679` ns/iter (`± 11071`)	`973599` ns/iter (`± 2729`)	`1.01`
`arith_agg-avg-count-min-max-sum-group_by`	`1236545` ns/iter (`± 41584`)	`1226677` ns/iter (`± 16934`)	`1.01`
`arith_agg-avg-count-min-max-sum-group_by-group_as`	`1851597` ns/iter (`± 26232`)	`1852614` ns/iter (`± 6514`)	`1.00`
`arith_agg-avg_distinct-count_distinct-min_distinct-max_distinct-sum_distinct`	`1204575` ns/iter (`± 13223`)	`1224367` ns/iter (`± 13913`)	`0.98`
`arith_agg-avg_distinct-count_distinct-min_distinct-max_distinct-sum_distinct-group_by`	`1505227` ns/iter (`± 11252`)	`1500889` ns/iter (`± 19101`)	`1.00`
`arith_agg-avg_distinct-count_distinct-min_distinct-max_distinct-sum_distinct-group_by-group_as`	`2123329` ns/iter (`± 20173`)	`2096925` ns/iter (`± 9815`)	`1.01`
`parse-1`	`6002` ns/iter (`± 642`)	`6086` ns/iter (`± 14`)	`0.99`
`parse-15`	`50598` ns/iter (`± 312`)	`50565` ns/iter (`± 77`)	`1.00`
`parse-30`	`99235` ns/iter (`± 404`)	`98481` ns/iter (`± 429`)	`1.01`
`compile-1`	`4106` ns/iter (`± 8`)	`4280` ns/iter (`± 12`)	`0.96`
`compile-15`	`31139` ns/iter (`± 601`)	`30186` ns/iter (`± 115`)	`1.03`
`compile-30`	`63510` ns/iter (`± 516`)	`63029` ns/iter (`± 238`)	`1.01`
`plan-1`	`70176` ns/iter (`± 215`)	`71625` ns/iter (`± 271`)	`0.98`
`plan-15`	`1093985` ns/iter (`± 9281`)	`1124229` ns/iter (`± 10520`)	`0.97`
`plan-30`	`2187788` ns/iter (`± 7953`)	`2253137` ns/iter (`± 14872`)	`0.97`
`eval-1`	`12010339` ns/iter (`± 71348`)	`11911863` ns/iter (`± 111520`)	`1.01`
`eval-15`	`76945710` ns/iter (`± 1003377`)	`77114178` ns/iter (`± 293901`)	`1.00`
`eval-30`	`147299495` ns/iter (`± 1612420`)	`147487722` ns/iter (`± 486288`)	`1.00`
`join`	`9927` ns/iter (`± 160`)	`9733` ns/iter (`± 187`)	`1.02`
`simple`	`2523` ns/iter (`± 11`)	`2542` ns/iter (`± 44`)	`0.99`
`simple-no`	`482` ns/iter (`± 2`)	`481` ns/iter (`± 1`)	`1.00`
`numbers`	`57` ns/iter (`± 0`)	`57` ns/iter (`± 0`)	`1`
`parse-simple`	`953` ns/iter (`± 5`)	`941` ns/iter (`± 4`)	`1.01`
`parse-ion`	`2703` ns/iter (`± 8`)	`2598` ns/iter (`± 15`)	`1.04`
`parse-group`	`7970` ns/iter (`± 28`)	`8131` ns/iter (`± 26`)	`0.98`
`parse-complex`	`20582` ns/iter (`± 133`)	`20806` ns/iter (`± 57`)	`0.99`
`parse-complex-fexpr`	`28121` ns/iter (`± 160`)	`28512` ns/iter (`± 2711`)	`0.99`