Fix #1326 (#1328)

levkk · web-flow · commit 98f114891db5 · 2024-02-23T10:08:27.000-08:00
diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -955,7 +955,8 @@ impl Model {
                                             .map_or(snapshot::NULL_CATEGORY_KEY.to_string(), |k| k.to_string())
                                     }
                                     pgrx_pg_sys::NUMERICOID => {
-                                        let element: Result<Option<AnyNumeric>, TryFromDatumError> = tuple.get_by_index(index);
+                                        let element: Result<Option<AnyNumeric>, TryFromDatumError> =
+                                            tuple.get_by_index(index);
                                         element
                                             .unwrap()
                                             .map_or(snapshot::NULL_CATEGORY_KEY.to_string(), |k| k.to_string())
@@ -999,7 +1000,8 @@ impl Model {
                                         features.push(element.unwrap().map_or(f32::NAN, |v| v as f32));
                                     }
                                     pgrx_pg_sys::NUMERICOID => {
-                                        let element: Result<Option<AnyNumeric>, TryFromDatumError> = tuple.get_by_index(index);
+                                        let element: Result<Option<AnyNumeric>, TryFromDatumError> =
+                                            tuple.get_by_index(index);
                                         features.push(element.unwrap().map_or(f32::NAN, |v| v.try_into().unwrap()));
                                     }
                                     // TODO handle NULL to NaN for arrays
diff --git a/pgml-extension/src/orm/snapshot.rs b/pgml-extension/src/orm/snapshot.rs
@@ -749,7 +749,7 @@ impl Snapshot {
                 .map(|c| c.quoted_name())
                 .collect::<Vec<String>>()
                 .join(", "),
-            self.relation_name(),
+            self.relation_name_quoted(),
             match self.materialized {
                 // If the snapshot is materialized, we already randomized it.
                 true => "",
@@ -990,7 +990,10 @@ impl Snapshot {
                                 "int8" => row[column.position].value::<i64>().unwrap().map(|v| v.to_string()),
                                 "float4" => row[column.position].value::<f32>().unwrap().map(|v| v.to_string()),
                                 "float8" => row[column.position].value::<f64>().unwrap().map(|v| v.to_string()),
-                                "numeric" => row[column.position].value::<AnyNumeric>().unwrap().map(|v| v.to_string()),
+                                "numeric" => row[column.position]
+                                    .value::<AnyNumeric>()
+                                    .unwrap()
+                                    .map(|v| v.to_string()),
                                 "bpchar" | "text" | "varchar" => {
                                     row[column.position].value::<String>().unwrap().map(|v| v.to_string())
                                 }
@@ -1084,7 +1087,7 @@ impl Snapshot {
                                         check_column_size(column, vec.len());
 
                                         for j in vec {
-                                            vector.push(j.rescale::<6,0>().unwrap().try_into().unwrap())
+                                            vector.push(j.rescale::<6, 0>().unwrap().try_into().unwrap())
                                         }
                                     }
                                     _ => error!(
@@ -1101,7 +1104,10 @@ impl Snapshot {
                                     "int8" => row[column.position].value::<i64>().unwrap().map(|v| v as f32),
                                     "float4" => row[column.position].value::<f32>().unwrap(),
                                     "float8" => row[column.position].value::<f64>().unwrap().map(|v| v as f32),
-                                    "numeric" => row[column.position].value::<AnyNumeric>().unwrap().map(|v| v.rescale::<6,0>().unwrap().try_into().unwrap()),
+                                    "numeric" => row[column.position]
+                                        .value::<AnyNumeric>()
+                                        .unwrap()
+                                        .map(|v| v.rescale::<6, 0>().unwrap().try_into().unwrap()),
                                     _ => error!(
                                         "Unhandled type for quantitative scalar column: {} {:?}",
                                         column.name, column.pg_type
@@ -1156,6 +1162,16 @@ impl Snapshot {
             false => self.relation_name.clone(),
         }
     }
+
+    fn relation_name_quoted(&self) -> String {
+        match self.materialized {
+            true => self.snapshot_name(), // Snapshot name is already safe.
+            false => {
+                let (schema_name, table_name) = Self::fully_qualified_table(&self.relation_name);
+                format!("\"{}\".\"{}\"", schema_name, table_name)
+            }
+        }
+    }
 }
 
 #[inline]

Original file line number	Diff line number	Diff line change
`@@ -955,7 +955,8 @@ impl Model {`
`955`	`955`	`.map_or(snapshot::NULL_CATEGORY_KEY.to_string(), \|k\| k.to_string())`
`956`	`956`	`}`
`957`	`957`	`pgrx_pg_sys::NUMERICOID => {`
`958`		`- let element: Result<Option<AnyNumeric>, TryFromDatumError> = tuple.get_by_index(index);`
	`958`	`+ let element: Result<Option<AnyNumeric>, TryFromDatumError> =`
	`959`	`+ tuple.get_by_index(index);`
`959`	`960`	`element`
`960`	`961`	`.unwrap()`
`961`	`962`	`.map_or(snapshot::NULL_CATEGORY_KEY.to_string(), \|k\| k.to_string())`
`@@ -999,7 +1000,8 @@ impl Model {`
`999`	`1000`	`features.push(element.unwrap().map_or(f32::NAN, \|v\| v as f32));`
`1000`	`1001`	`}`
`1001`	`1002`	`pgrx_pg_sys::NUMERICOID => {`
`1002`		`- let element: Result<Option<AnyNumeric>, TryFromDatumError> = tuple.get_by_index(index);`
	`1003`	`+ let element: Result<Option<AnyNumeric>, TryFromDatumError> =`
	`1004`	`+ tuple.get_by_index(index);`
`1003`	`1005`	`features.push(element.unwrap().map_or(f32::NAN, \|v\| v.try_into().unwrap()));`
`1004`	`1006`	`}`
`1005`	`1007`	`// TODO handle NULL to NaN for arrays`