work in progress

Author: jfeser

bench/tpcds/q38.txt

@@ -104,7 +104,7 @@ type ('p, 'r) query =
   | Dedup of 'r
   | Relation of Relation.t
   | Range of ('p * 'p)
-  | Limit of int
+  | Limit of int * 'r
   | AEmpty
   | AScalar of 'p scalar
   | AList of 'r list_

bench/tpcds/q8.txt

@@ -45,6 +45,7 @@ module Query = struct
 
   let list' = alist
   let list a c = list' { l_keys = a; l_values = c }
+  let limit = limit
 end
 
 (** Construct annotated queries. Discards any existing metadata. *)
@@ -110,6 +111,7 @@ module Annot = struct
   let list' x = wrap @@ Query.list' (strip_meta_alist x)
   let tuple a b = wrap @@ Query.tuple (strip_meta_list a) b
   let call a = wrap @@ Query.call a
+  let limit l r = wrap @@ Query.limit l r
 
   let hash_idx ?key_layout a c d =
     wrap

lib/ast.ml

@@ -52,6 +52,7 @@ module Annot : sig
   val list' : _ annot Ast.list_ -> < > annot
   val tuple : _ annot list -> Ast.tuple -> < > annot
   val call : string -> < > annot
+  val limit : int -> < > annot -> < > annot
 
   val hash_idx :
     ?key_layout:_ annot -> _ annot -> _ annot -> _ pred list -> < > annot

lib/constructors.ml

@@ -28,27 +28,38 @@ let conv_alias = sprintf "%s_%s"
 let conv_sql db_schema =
   let open Sqlgg in
   let open Constructors.Annot in
-  let aliases = ref (Set.empty (module String)) in
-  let alias_of_name = ref (Map.empty (module String)) in
+  let fresh = Fresh.create () in
 
-  let conv_order _ q =
-    Log.warn (fun m -> m "Dropping orderby clause.");
-    q
-  in
   let conv_limit l q =
-    if Option.is_some l then Log.warn (fun m -> m "Dropping limit clause.");
-    q
+    match l with
+    | None -> q
+    | Some (_, true) -> limit 1 q
+    | Some (_, false) ->
+        Log.warn (fun m -> m "Dropping limit clause.");
+        q
   in
+
   let conv_distinct d q = if d then dedup q else q in
 
-  let rec conv_filter f q =
-    match f with Some e -> filter (conv_expr e) q | None -> q
+  let rec conv_order order q subst =
+    let key =
+      List.map order ~f:(fun (sql_expr, sql_dir) ->
+          let dir =
+            match sql_dir with Some `Asc -> Asc | Some `Desc | None -> Desc
+          in
+          (Subst.subst_pred subst (conv_expr sql_expr), dir))
+    in
+    order_by key q
+  and conv_filter f q subst =
+    match f with
+    | Some e -> filter (Subst.subst_pred subst (conv_expr e)) q
+    | None -> q
   and conv_source (s, alias) =
     let q =
       match s with
       | `Subquery s -> conv_query s
       | `Table t -> relation (Db.Schema.relation db_schema t)
-      | `Nested n -> conv_nested n
+      | `Nested n -> conv_nested Map.empty n
     in
     match alias with
     | Some a ->
@@ -62,29 +73,29 @@ let conv_sql db_schema =
         in
         select select_list q
     | None -> q
-  and conv_nested (q, qs) =
+  and conv_nested subst (q, qs) =
     match qs with
     | [] -> conv_source q
     | (q', j) :: qs' -> (
         match j with
         | `Cross | `Default ->
-            join (bool true) (conv_source q) (conv_nested (q', qs'))
+            join (bool true) (conv_source q) (conv_nested subst (q', qs'))
         | `Search e ->
-            join (conv_expr e) (conv_source q) (conv_nested (q', qs'))
+            join (conv_expr subst e) (conv_source q)
+              (conv_nested subst (q', qs'))
         | `Using _ | `Natural -> failwith "Join type not supported")
   and conv_subquery q = conv_query q
-  and conv_column (c : Sql.Col_name.t) =
-    let n =
+  and conv_column subst (c : Sql.Col_name.t) =
+    let name =
       match c.tname with
-      | Some a -> if Set.mem !aliases a then conv_alias a c.cname else c.cname
-      | None -> (
-          match Map.find !alias_of_name c.cname with
-          | Some a -> conv_alias a c.cname
-          | None -> c.cname)
+      | Some table_name ->
+          `Name { Name.name = Name.Attr (table_name, c.cname); type_ = None }
+      | None -> `Name (Name.create c.cname)
     in
-    Name.create n
-  and conv_expr e =
-    let open Pred in
+    Subst.subst_pred subst name
+  and conv_expr subst e =
+    let conv_expr = conv_expr subst in
+    let module Infix = Pred.Infix in
     match e with
     | Sql.Value v -> (
         match v with
@@ -94,8 +105,9 @@ let conv_sql db_schema =
         | Bool x -> `Bool x
         | Float x -> `Fixed (Fixed_point.of_float x)
         | Null -> `Null None)
-    | Param _ | Choices (_, _) | Inserted _ | Sequence _ ->
-        failwith "unsupported"
+    | Param ((Some name, _), _) -> `Name (Name.create name)
+    | Param ((None, _), _) -> `Name (Name.create (Fresh.name fresh "param_%d"))
+    | Choices (_, _) | Inserted _ | Sequence _ -> failwith "unsupported"
     | Case (branches, else_) ->
         let else_ =
           Option.map else_ ~f:conv_expr |> Option.value ~default:(`Null None)
@@ -145,7 +157,7 @@ let conv_sql db_schema =
             Error.create "Unsupported op" op [%sexp_of: Sql.op] |> Error.raise)
     | Subquery (s, `Exists) -> `Exists (conv_subquery s)
     | Subquery (s, `AsValue) -> `First (conv_subquery s)
-    | Column c -> `Name (conv_column c)
+    | Column c -> conv_column subst c
   and conv_select s =
     (* Build query in order. First, FROM *)
     let query =

lib/constructors.mli

@@ -272,6 +272,9 @@ let resolve_open resolve stage outer_ctx =
       let r, value_ctx = rsame outer_ctx r in
       let pred = resolve_pred stage (Ctx.merge outer_ctx value_ctx) pred in
       (Filter (pred, r), value_ctx)
+  | Limit (l, r) ->
+      let r, value_ctx = rsame outer_ctx r in
+      (Limit (l, r), value_ctx)
   | DepJoin { d_lhs; d_rhs } ->
       let d_lhs, lctx = rsame outer_ctx d_lhs in
       let lctx = Ctx.zero lctx in

lib/of_sql.ml

@@ -48,7 +48,8 @@ let schema_query_open schema r =
     | DepJoin { d_rhs = r; _ }
     | Filter (_, r)
     | Dedup r
-    | OrderBy { rel = r; _ } ->
+    | OrderBy { rel = r; _ }
+    | Limit (_, r) ->
         schema r
     | AOrderedIdx { oi_keys = lhs; oi_values = rhs; _ }
     | AHashIdx { hi_keys = lhs; hi_values = rhs; _ } ->

lib/resolve.ml

@@ -281,9 +281,10 @@ class virtual ['self] base_endo =
       let r0 = self#visit_string env c0 in
       if c0 == r0 then this else Call r0
 
-    method visit_Limit env this c0 =
-      let r0 = self#visit_int env c0 in
-      if c0 == r0 then this else Limit r0
+    method visit_Limit env this l r =
+      let l' = self#visit_int env l in
+      let r' = self#visit_t env r in
+      if l == l' && r == r' then this else Limit (l', r')
 
     method visit_query env this =
       match this with
@@ -303,7 +304,7 @@ class virtual ['self] base_endo =
       | AHashIdx c0 as this -> self#visit_AHashIdx env this c0
       | AOrderedIdx c0 as this -> self#visit_AOrderedIdx env this c0
       | Call c0 as this -> self#visit_Call env this c0
-      | Limit c0 as this -> self#visit_Limit env this c0
+      | Limit (c0, r) as this -> self#visit_Limit env this c0 r
 
     method visit_annot env this =
       let r0 = self#visit_query env this.node in
@@ -579,9 +580,10 @@ class virtual ['self] base_map =
       let r0 = self#visit_string env c0 in
       Call r0
 
-    method visit_Limit env c0 =
-      let r0 = self#visit_int env c0 in
-      Limit r0
+    method visit_Limit env l r =
+      let l' = self#visit_int env l in
+      let r' = self#visit_t env r in
+      Limit (l', r')
 
     method visit_query env this =
       match this with
@@ -601,7 +603,7 @@ class virtual ['self] base_map =
       | AHashIdx c0 -> self#visit_AHashIdx env c0
       | AOrderedIdx c0 -> self#visit_AOrderedIdx env c0
       | Call c0 -> self#visit_Call env c0
-      | Limit c0 -> self#visit_Limit env c0
+      | Limit (l, r) -> self#visit_Limit env l r
 
     method visit_annot env this =
       let r0 = self#visit_query env this.node in
@@ -873,7 +875,10 @@ class virtual ['self] base_iter =
       ()
 
     method visit_Call env c0 = self#visit_string env c0
-    method visit_Limit env c0 = self#visit_int env c0
+
+    method visit_Limit env l r =
+      self#visit_int env l;
+      self#visit_t env r
 
     method visit_query env this =
       match this with
@@ -893,7 +898,7 @@ class virtual ['self] base_iter =
       | AHashIdx c0 -> self#visit_AHashIdx env c0
       | AOrderedIdx c0 -> self#visit_AOrderedIdx env c0
       | Call c0 -> self#visit_Call env c0
-      | Limit c0 -> self#visit_Limit env c0
+      | Limit (l, r) -> self#visit_Limit env l r
 
     method visit_annot env this =
       let r0 = self#visit_query env this.node in
@@ -1165,7 +1170,9 @@ class virtual ['self] base_reduce =
       s0
 
     method visit_Call env c0 = self#visit_string env c0
-    method visit_Limit env c0 = self#visit_int env c0
+
+    method visit_Limit env l r =
+      self#plus (self#visit_int env l) (self#visit_t env r)
 
     method visit_query env this =
       match this with
@@ -1185,7 +1192,7 @@ class virtual ['self] base_reduce =
       | AHashIdx c0 -> self#visit_AHashIdx env c0
       | AOrderedIdx c0 -> self#visit_AOrderedIdx env c0
       | Call c0 -> self#visit_Call env c0
-      | Limit c0 -> self#visit_Limit env c0
+      | Limit (l, r) -> self#visit_Limit env l r
 
     method visit_annot env this =
       let s0 = self#visit_query env this.node in
@@ -1466,9 +1473,10 @@ class virtual ['self] base_mapreduce =
       let r0, s0 = self#visit_string env c0 in
       (Call r0, s0)
 
-    method visit_Limit env c0 =
-      let r0, s0 = self#visit_int env c0 in
-      (Limit r0, s0)
+    method visit_Limit env l r =
+      let r0, s0 = self#visit_int env l in
+      let r1, s1 = self#visit_t env r in
+      (Limit (r0, r1), self#plus s0 s1)
 
     method visit_query env this =
       match this with
@@ -1488,7 +1496,7 @@ class virtual ['self] base_mapreduce =
       | AHashIdx c0 -> self#visit_AHashIdx env c0
       | AOrderedIdx c0 -> self#visit_AOrderedIdx env c0
       | Call c0 -> self#visit_Call env c0
-      | Limit c0 -> self#visit_Limit env c0
+      | Limit (c0, c1) -> self#visit_Limit env c0 c1
 
     method visit_annot env this =
       let r0, s0 = self#visit_query env this.node in

lib/schema.ml

@@ -33,21 +33,21 @@ let dataclass_name table =
 
 let singular = String.chop_suffix_if_exists ~suffix:"s"
 
-let emit_dataclass db_schema name : Python_block.t =
-  let schema = Db.Schema.attrs db_schema name in
+let emit_dataclass db_schema table_name : Python_block.t =
+  let schema = Db.Schema.attrs db_schema table_name in
   let fields =
     List.filter_map schema ~f:(fun attr ->
         match attr.constraints with
         | `Foreign_key tbl ->
             Some
               (Fmt.str "%s : '%s'"
-                 (String.chop_suffix_exn name ~suffix:"_id")
+                 (String.chop_suffix_exn attr.attr_name ~suffix:"_id")
                  (dataclass_name tbl))
         | `None | `Primary_key ->
-            Some (Fmt.str "%s : %s" name (emit_type attr.type_)))
+            Some (Fmt.str "%s : %s" attr.attr_name (emit_type attr.type_)))
   in
   let referrers =
-    List.map (Db_schema.referrers db_schema name) ~f:(fun tbl ->
+    List.map (Db_schema.referrers db_schema table_name) ~f:(fun tbl ->
         Fmt.str "%s : set['%s'] = field(default_factory=set, compare=False)" tbl
           (dataclass_name tbl))
   in
@@ -56,7 +56,7 @@ let emit_dataclass db_schema name : Python_block.t =
       Stmts [ "@dataclass(frozen=True)" ];
       Block
         {
-          header = Fmt.str "class %s:" (dataclass_name name);
+          header = Fmt.str "class %s:" (dataclass_name table_name);
           body = Stmts (fields @ referrers);
         };
     ]
@@ -385,9 +385,9 @@ let emit_expr _db_schema params ctx =
     | Column x -> lookup_column_exn ctx x
     | Param p ->
         List.find_map params ~f:(fun (name, param) ->
-            if [%compare.equal: Sql.param] p param then Some name else None)
+            if [%compare.equal: Sql.Param.t] p param then Some name else None)
         |> Option.value_exn
-             ~error:(Error.of_lazy_sexp (lazy [%message (p : Sql.param)]))
+             ~error:(Error.of_lazy_sexp (lazy [%message (p : Sql.Param.t)]))
     | (Sequence _ | Choices _ | Case _ | Subquery _ | Inserted _) as expr ->
         raise_s [%message "unsupported" (expr : Sql.op Sql.expr)]
   in
@@ -792,6 +792,7 @@ let emit_select db_schema query_str query_name select =
     }
 
 let main () =
+  Log.setup_log Warning;
   let input_str = In_channel.input_all In_channel.stdin in
   let queries = load_queries input_str in
   let db =
@@ -804,6 +805,19 @@ let main () =
   let tables =
     Db.Schema.relation_names db |> List.sort ~compare:[%compare: string]
   in
+  List.iter queries ~f:(fun (q, _) ->
+      match q with
+      | Select q ->
+          let castor_query = Of_sql.conv_sql db q in
+          print_s [%message (castor_query : Ast.t)];
+          print_s
+            [%message
+              (Resolve.resolve_exn
+                 ~params:(Set.empty (module Name))
+                 castor_query
+                : _ Ast.annot)]
+      | _ -> ());
+
   let dataclasses = List.map tables ~f:(emit_dataclass db) in
   let add_methods = List.map tables ~f:(emit_add db) in
   let remove_methods = List.map tables ~f:(emit_remove db) in

lib/visitors_gen.ml

@@ -6,5 +6,5 @@
 (library
  (public_name sqlgg)
  (libraries base stdio)
- (preprocess (pps ppx_compare ppx_sexp_conv))
+ (preprocess (pps ppx_compare ppx_hash ppx_sexp_conv))
 )