add columns

csst_dfs_api/common/catalog.py

@@ -9,13 +9,14 @@ class CatalogApi(object):
         self.module = Delegate().load(sub_module = "common")
         self.stub = getattr(self.module, "CatalogApi")()
     
-    def catalog_query(self, ra: float, dec: float, radius: float, catalog_name: str, min_mag: float,  max_mag: float,  obstime: int, limit: int):
+    def catalog_query(self, ra: float, dec: float, radius: float, catalog_name: str, columns: tuple, min_mag: float,  max_mag: float,  obstime: int, limit: int):
         ''' retrieval catalog
 
         :param ra: in deg
         :param dec:  in deg
         :param radius:  in deg
         :param catalog_name: one of ['gaia3','','']
+        :param columns: tuple of str, like ('ra','dec','phot_g_mean_mag')
         :param min_mag: minimal magnitude
         :param max_mag: maximal magnitude
         :param obstime: seconds  
@@ -32,12 +33,13 @@ class CatalogApi(object):
     def to_table(self, query_result):
         return to_fits_table(query_result)
 
-    def gaia3_query(self, ra: float, dec: float, radius: float, min_mag: float,  max_mag: float,  obstime: int, limit: int):
+    def gaia3_query(self, ra: float, dec: float, radius: float, columns: tuple,min_mag: float,  max_mag: float,  obstime: int, limit: int):
         """retrieval GAIA EDR 3
         
         :param ra: in deg
         :param dec:  in deg
         :param radius:  in deg
+        :param columns: tuple of str, like ('ra','dec','phot_g_mean_mag')
         :param min_mag: minimal magnitude
         :param max_mag: maximal magnitude
         :param obstime: seconds  
@@ -46,6 +48,8 @@ class CatalogApi(object):
         :returns: csst_dfs_common.models.Result
         """ 
         try:
-            return self.stub.gaia3_query(ra, dec, radius, min_mag, max_mag, obstime, limit)
+            if not columns:
+                raise Exception("columns is empty")
+            return self.stub.gaia3_query(ra, dec, radius, columns, min_mag, max_mag, obstime, limit)
         except Exception as e:
             return Result.error(message=repr(e))

csst_dfs_api/common/utils.py

@@ -80,7 +80,36 @@ def fields_dtypes(rec):
     dtypes = tuple(dtypes)
     return fields, dtypes
 
+def tuple_fields_dtypes(rec: tuple):
+    dtypes = []
+    for _, f in rec:
+        if f.type == int:
+            dtypes.append('i8')
+        elif f.type == float:
+            dtypes.append('f8')        
+        elif f.type == str:
+            dtypes.append('S2')
+        elif f.type == list:
+            dtypes.append('(12,)f8')       
+        else:
+            dtypes.append('S2')                
+    dtypes = tuple(dtypes)
+    return dtypes
+
 def to_table(query_result):
+    if not query_result.success or not query_result.data:
+        return Table()
+    fields = query_result['columns']
+    dtypes = tuple_fields_dtypes(query_result.data[0])
+    t = Table(names = fields, dtype = dtypes)
+    t.meta['comments'] = [str(query_result.data[0].__class__)]
+    t.meta['total'] = query_result['totalCount']
+
+    for rec in query_result.data:
+        t.add_row(rec)
+    return t
+
+def object_list_to_table(query_result):
     if not query_result.success or not query_result.data:
         return Table()
     fields, dtypes = fields_dtypes(query_result.data[0])
@@ -90,4 +119,4 @@ def to_table(query_result):
 
     for rec in query_result.data:
         t.add_row(tuple([rec.__getattribute__(k) for k in fields]))
-    return t
+    return t    
\ No newline at end of file

tests/common/test_common_catalog.py

@@ -13,18 +13,18 @@ class CommonCatalogTestCase(unittest.TestCase):
     def test_catalog_query(self):
         t= time.time()
         result = self.api.catalog_query(
-            ra=193.299, 
-            dec=19.08851, 
-            radius=2, 
+            ra=90, 
+            dec=24.5, 
+            radius=1, 
             catalog_name='gaia3', 
             min_mag=-1, 
             max_mag=-1, 
             obstime = -1, 
             limit = 0
         )
-        print(result)
+        # print(result)
         # dt = self.api.to_table(result)
         # df = dt.to_pandas()
         # print(df.head())
         print('used:', time.time()-t)
-        # print('return:', result['totalCount'\])
+        print('return:', len(result.data))