perf: improve query_region for large number of centers

Author: ManonMarchand

CHANGES.rst

@@ -56,6 +56,9 @@ simbad
   empty for an object, there will now be a line with masked values instead of no line in
   the result [#3199]
 
+- perf: use a TAP upload in ``query_region`` when there are more than 300 coordinates.
+  This prevents timeouts. [#3235]
+
 xmatch
 ^^^^^^
 

astroquery/simbad/core.py

@@ -18,7 +18,7 @@
 
 from astroquery.query import BaseVOQuery
 from astroquery.utils import commons
-from astroquery.exceptions import LargeQueryWarning, NoResultsWarning
+from astroquery.exceptions import NoResultsWarning
 from astroquery.simbad.utils import (_catch_deprecated_fields_with_arguments,
                                      _wildcard_to_regexp, CriteriaTranslator,
                                      query_criteria_fields)
@@ -698,7 +698,7 @@ def query_objects(self, object_names, *, wildcard=False, criteria=None,
 
     @deprecated_renamed_argument(["equinox", "epoch", "cache"],
                                  new_name=[None]*3,
-                                 alternative=["astropy.coordinates.SkyCoord"]*3,
+                                 alternative=["astropy.coordinates.SkyCoord"]*2 + [None],
                                  since=['0.4.8']*3, relax=True)
     def query_region(self, coordinates, radius=2*u.arcmin, *,
                      criteria=None, get_query_payload=False,
@@ -768,42 +768,54 @@ def query_region(self, coordinates, radius=2*u.arcmin, *,
         center = center.transform_to("icrs")
 
         top, columns, joins, instance_criteria = self._get_query_parameters()
+        if criteria:
+            instance_criteria.append(f"({criteria})")
 
         if center.isscalar:
             radius = coord.Angle(radius)
             instance_criteria.append(f"CONTAINS(POINT('ICRS', basic.ra, basic.dec), "
                                      f"CIRCLE('ICRS', {center.ra.deg}, {center.dec.deg}, "
                                      f"{radius.to(u.deg).value})) = 1")
+            return self._query(top, columns, joins, instance_criteria,
+                               get_query_payload=get_query_payload)
 
+        # from uploadLimit in SIMBAD's capabilities
+        # http://simbad.cds.unistra.fr/simbad/sim-tap/capabilities
+        if len(center) > 200000:
+            raise ValueError("'query_region' can process up to 200000 centers. For "
+                             "larger queries, split your centers list.")
+
+        # `radius` as `str` is iterable, but contains only one value.
+        if isiterable(radius) and not isinstance(radius, str):
+            if len(radius) != len(center):
+                raise ValueError(f"Mismatch between radii of length {len(radius)}"
+                                 f" and center coordinates of length {len(center)}.")
+            radius = [coord.Angle(item).to(u.deg).value for item in radius]
         else:
-            if len(center) > 10000:
-                warnings.warn(
-                    "For very large queries, you may receive a timeout error.  SIMBAD suggests"
-                    " splitting queries with >10000 entries into multiple threads",
-                    LargeQueryWarning, stacklevel=2
-                )
-
-            # `radius` as `str` is iterable, but contains only one value.
-            if isiterable(radius) and not isinstance(radius, str):
-                if len(radius) != len(center):
-                    raise ValueError(f"Mismatch between radii of length {len(radius)}"
-                                     f" and center coordinates of length {len(center)}.")
-                radius = [coord.Angle(item) for item in radius]
-            else:
-                radius = [coord.Angle(radius)] * len(center)
+            radius = [coord.Angle(radius).to(u.deg).value] * len(center)
 
+        # for small number of centers, not using the upload method is faster
+        if len(center) <= 300:
             cone_criteria = [(f"CONTAINS(POINT('ICRS', basic.ra, basic.dec), CIRCLE('ICRS', "
-                             f"{center.ra.deg}, {center.dec.deg}, {radius.to(u.deg).value})) = 1 ")
+                             f"{center.ra.deg}, {center.dec.deg}, {radius})) = 1 ")
                              for center, radius in zip(center, radius)]
 
             cone_criteria = f" ({'OR '.join(cone_criteria)})"
             instance_criteria.append(cone_criteria)
 
-        if criteria:
-            instance_criteria.append(f"({criteria})")
+            return self._query(top, columns, joins, instance_criteria,
+                               get_query_payload=get_query_payload)
+
+        # for longer centers list, we use a TAP upload
+        upload_centers = Table({"ra": center.ra.deg, "dec": center.dec.deg,
+                                "radius": radius})
+        sub_query = "(SELECT ra, dec, radius FROM TAP_UPLOAD.centers) AS centers"
+        instance_criteria.append("CONTAINS(POINT('ICRS', basic.ra, basic.dec), CIRCLE"
+                                 "('ICRS', centers.ra, centers.dec, centers.radius)) = 1 ")
 
         return self._query(top, columns, joins, instance_criteria,
-                           get_query_payload=get_query_payload)
+                           from_table=f"{sub_query}, basic",
+                           get_query_payload=get_query_payload, centers=upload_centers)
 
     @deprecated_renamed_argument(["verbose", "cache"], new_name=[None, None],
                                  since=['0.4.8', '0.4.8'], relax=True)

astroquery/simbad/tests/test_simbad.py

@@ -13,7 +13,7 @@
 
 from ... import simbad
 from .test_simbad_remote import multicoords
-from astroquery.exceptions import LargeQueryWarning, NoResultsWarning
+from astroquery.exceptions import NoResultsWarning
 
 
 GALACTIC_COORDS = SkyCoord(l=-67.02084 * u.deg, b=-29.75447 * u.deg, frame="galactic")
@@ -448,11 +448,9 @@ def test_query_region_with_criteria():
     adql = simbad.core.Simbad.query_region(ICRS_COORDS, radius="0.1s",
                                            criteria="galdim_majaxis>0.2",
                                            get_query_payload=True)["QUERY"]
-    assert adql.endswith("AND (galdim_majaxis>0.2)")
+    assert "(galdim_majaxis>0.2)" in adql
 
 
-# transform large query warning into error to save execution time
-@pytest.mark.filterwarnings("error:For very large queries")
 @pytest.mark.usefixtures("_mock_simbad_class")
 def test_query_region_errors():
     with pytest.raises(u.UnitsError):
@@ -462,9 +460,24 @@ def test_query_region_errors():
     with pytest.raises(ValueError, match="Mismatch between radii of length 3 "
                        "and center coordinates of length 2."):
         simbad.SimbadClass().query_region(multicoords, radius=[1, 2, 3] * u.deg)
-    centers = SkyCoord([0] * 10001, [0] * 10001, unit="deg", frame="icrs")
-    with pytest.raises(LargeQueryWarning, match="For very large queries, you may receive a timeout error.*"):
-        simbad.core.Simbad.query_region(centers, radius="2m", get_query_payload=True)["QUERY"]
+
+
+@pytest.mark.usefixtures("_mock_simbad_class")
+def test_query_region_error_on_long_list_of_centers(monkeypatch):
+    # initiating a SkyCoord longer than 200000 takes a few seconds
+    monkeypatch.setattr(SkyCoord, "__len__", lambda self: 200001)
+    centers = SkyCoord([0, 0], [0, 0], unit="deg", frame="icrs")
+    with pytest.raises(ValueError, match="'query_region' can process up to 200000 centers.*"):
+        simbad.core.Simbad.query_region(centers, radius="2m")
+
+
+@pytest.mark.usefixtures("_mock_simbad_class")
+def test_query_region_upload():
+    centers = SkyCoord([0] * 301, [0] * 301, unit="deg", frame="icrs")
+    adql = simbad.core.Simbad.query_region(centers, radius=["2m"] * 301,
+                                           get_query_payload=True)["QUERY"]
+    assert adql.endswith("WHERE CONTAINS(POINT('ICRS', basic.ra, basic.dec), CIRCLE"
+                         "('ICRS', centers.ra, centers.dec, centers.radius)) = 1 ")
 
 
 @pytest.mark.usefixtures("_mock_simbad_class")

astroquery/simbad/tests/test_simbad_remote.py

@@ -1,4 +1,5 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
+import numpy as np
 import pytest
 
 from astropy.coordinates import SkyCoord
@@ -79,6 +80,13 @@ def test_query_regions(self):
         # filtering on main_id to retrieve the two cone centers
         assert {"M  81", "M  10"} == set(result["main_id"].data.data)
 
+    def test_query_regions_long_list(self):
+        self.simbad.ROW_LIMIT = -1
+        # we create a list of centers longer than 300 to trigger the TAP upload case
+        centers = SkyCoord(np.arange(0, 360, 1), np.arange(0, 180, 0.5) - 90, unit="deg")
+        result = self.simbad.query_region(centers, radius="1m")
+        assert len(result) > 90
+
     def test_query_object_ids(self):
         self.simbad.ROW_LIMIT = -1
         result = self.simbad.query_objectids("Polaris")