Fix bug on existing data types and reloptions

greenplum-db · Jul 26, 2023 · 4584d25 · 4584d25
1 parent 56d146f
commit 4584d25
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 14 deletions.
diff --git a/doc/source/notebooks/embedding.ipynb b/doc/source/notebooks/embedding.ipynb
@@ -28,8 +28,8 @@
       "Requirement already satisfied, skipping upgrade: dill==0.3.6 in /home/gpadmin/.local/lib/python3.9/site-packages (from greenplum-python==1.0.1) (0.3.6)\n",
       "Building wheels for collected packages: greenplum-python\n",
       "  Building wheel for greenplum-python (PEP 517) ... \u001b[?25ldone\n",
-      "\u001b[?25h  Created wheel for greenplum-python: filename=greenplum_python-1.0.1-py3-none-any.whl size=70614 sha256=4ee428916d3690ae05c591dff7e48e84e41f04d37fba06dfd6bf9543791e4d4f\n",
-      "  Stored in directory: /tmp/pip-ephem-wheel-cache-2iw1rkm7/wheels/bb/1f/99/ff8594e48ec11df99af6e0ee8611a5e560e9f44d1a3fefb351\n",
+      "\u001b[?25h  Created wheel for greenplum-python: filename=greenplum_python-1.0.1-py3-none-any.whl size=70623 sha256=6bbdfa6fb272db092d6e63def388beb58365fb6b726fc98bf40ff788ebe3143f\n",
+      "  Stored in directory: /tmp/pip-ephem-wheel-cache-oqna_kdn/wheels/bb/1f/99/ff8594e48ec11df99af6e0ee8611a5e560e9f44d1a3fefb351\n",
       "Successfully built greenplum-python\n",
       "Installing collected packages: greenplum-python\n",
       "Successfully installed greenplum-python-1.0.1\n"
@@ -75,7 +75,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -106,7 +106,7 @@
        "(2 rows)"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 4,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -127,7 +127,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -153,7 +153,7 @@
        "(1 row)"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }

diff --git a/greenplumpython/experimental/embedding.py b/greenplumpython/experimental/embedding.py
@@ -74,7 +74,7 @@ def create_index(self, column: str, model: str) -> gp.DataFrame:
                 SET LOCAL allow_system_table_mods TO ON;
 
                 WITH embedding_info AS (
-                    SELECT '{embedding_df._qualified_table_name}'::regclass::oid AS base_relid, attnum, '{model}' AS model
+                    SELECT '{embedding_df._qualified_table_name}'::regclass::oid AS embedding_relid, attnum, '{model}' AS model
                     FROM pg_attribute
                     WHERE 
                         attrelid = '{self._dataframe._qualified_table_name}'::regclass::oid AND
@@ -85,10 +85,11 @@ def create_index(self, column: str, model: str) -> gp.DataFrame:
                     reloptions, 
                     format('_pygp_emb_%s=%s', attnum::text, to_json(embedding_info))
                 )
-                FROM embedding_info;
+                FROM embedding_info
+                WHERE oid = '{self._dataframe._qualified_table_name}'::regclass::oid;
 
                 WITH embedding_info AS (
-                    SELECT '{embedding_df._qualified_table_name}'::regclass::oid AS base_relid, attnum, '{model}' AS model
+                    SELECT '{embedding_df._qualified_table_name}'::regclass::oid AS embedding_relid, attnum, '{model}' AS model
                     FROM pg_attribute
                     WHERE
                         attrelid = '{self._dataframe._qualified_table_name}'::regclass::oid AND
@@ -100,7 +101,7 @@ def create_index(self, column: str, model: str) -> gp.DataFrame:
                     '{embedding_df._qualified_table_name}'::regclass::oid AS objid,
                     0::oid AS objsubid,
                     'pg_class'::regclass::oid AS refclassid,
-                    embedding_info.base_relid AS refobjid,
+                    embedding_info.embedding_relid AS refobjid,
                     embedding_info.attnum AS refobjsubid,
                     'n' AS deptype
                 FROM embedding_info;
@@ -131,14 +132,14 @@ def search(self, column: str, query: Any, top_k: int) -> gp.DataFrame:
                 WHERE option LIKE format('_pygp_emb_%s=%%', attnum)
             ), embedding_info AS (
                 SELECT * 
-                FROM embedding_info_json, json_to_record(val) AS (attnum int4, base_relid oid, model text)
+                FROM embedding_info_json, json_to_record(val) AS (attnum int4, embedding_relid oid, model text)
             )
             SELECT nspname, relname, attname, model
             FROM embedding_info, pg_class, pg_namespace, pg_attribute
             WHERE 
-                pg_class.oid = base_relid AND
+                pg_class.oid = embedding_relid AND
                 relnamespace = pg_namespace.oid AND
-                base_relid = attrelid AND
+                embedding_relid = attrelid AND
                 pg_attribute.attnum = 2;
             """
         )

diff --git a/greenplumpython/type.py b/greenplumpython/type.py
@@ -230,7 +230,7 @@ def to_pg_type(
             args: Tuple[type, ...] = annotation.__args__
             if for_return:
                 return f"SETOF {to_pg_type(args[0], db)}"  # type: ignore
-            if args[0] in _defined_types:
+            else:
                 return f"{to_pg_type(args[0], db)}[]"  # type: ignore
         raise NotImplementedError()
     else: