Repository: alteryx/featuretools

Issue:
Move query_by_values from Entity to EntitySet
- the `query_by_values` functions on Entity needs to moved to EntitySet with the Woodwork integration

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=h1) Report
> Merging [#1251](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=desc) (f71a6d2) into [main](https://codecov.io/gh/alteryx/featuretools/commit/9819e6db26e90505071a9689e5c1b713bfac13f6?el=desc) (9819e6d) will **increase** coverage by `0.00%`.
> The diff coverage is `96.92%`.

[![Impacted file tree graph](https://codecov.io/gh/alteryx/featuretools/pull/1251/graphs/tree.svg?width=650&height=150&src=pr&token=PqDLD96eXJ)](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##             main    #1251   +/-   ##
=======================================
  Coverage   98.34%   98.34%           
=======================================
  Files         134      134           
  Lines       14430    14432    +2     
=======================================
+ Hits        14191    14193    +2     
  Misses        239      239           
```


| [Impacted Files](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [featuretools/entityset/entity.py](https://codecov.io/gh/alteryx/featuretools/pull/1251/diff?src=pr&el=tree#diff-ZmVhdHVyZXRvb2xzL2VudGl0eXNldC9lbnRpdHkucHk=) | `94.26% <ø> (-0.08%)` | :arrow_down: |
| [featuretools/tests/entityset\_tests/test\_entity.py](https://codecov.io/gh/alteryx/featuretools/pull/1251/diff?src=pr&el=tree#diff-ZmVhdHVyZXRvb2xzL3Rlc3RzL2VudGl0eXNldF90ZXN0cy90ZXN0X2VudGl0eS5weQ==) | `100.00% <ø> (ø)` | |
| [featuretools/entityset/entityset.py](https://codecov.io/gh/alteryx/featuretools/pull/1251/diff?src=pr&el=tree#diff-ZmVhdHVyZXRvb2xzL2VudGl0eXNldC9lbnRpdHlzZXQucHk=) | `97.36% <95.12%> (-0.23%)` | :arrow_down: |
| [...s/computational\_backends/feature\_set\_calculator.py](https://codecov.io/gh/alteryx/featuretools/pull/1251/diff?src=pr&el=tree#diff-ZmVhdHVyZXRvb2xzL2NvbXB1dGF0aW9uYWxfYmFja2VuZHMvZmVhdHVyZV9zZXRfY2FsY3VsYXRvci5weQ==) | `98.69% <100.00%> (ø)` | |
| [featuretools/tests/entityset\_tests/test\_es.py](https://codecov.io/gh/alteryx/featuretools/pull/1251/diff?src=pr&el=tree#diff-ZmVhdHVyZXRvb2xzL3Rlc3RzL2VudGl0eXNldF90ZXN0cy90ZXN0X2VzLnB5) | `98.83% <100.00%> (+0.02%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=footer). Last update [9819e6d...c481ce0](https://codecov.io/gh/alteryx/featuretools/pull/1251?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

thehomebrewnerd: > A couple if statements in `_vals_to_series` aren't getting hit. Are there input combinations to featuretools that would result in these lines getting hit?

I can take a closer look. This code was copied directly from the code that was in Entity, so I'm assuming it was uncovered there as well. I'll confirm.

If it was uncovered previously, do you want to try to add new tests to fix the coverage as part of this PR?
thehomebrewnerd: > > A couple if statements in `_vals_to_series` aren't getting hit. Are there input combinations to featuretools that would result in these lines getting hit?
> 
> I can take a closer look. This code was copied directly from the code that was in Entity, so I'm assuming it was uncovered there as well. I'll confirm.
> 
> If it was uncovered previously, do you want to try to add new tests to fix the coverage as part of this PR?

After further review, these lines also are uncovered currently in main. Looking at the code flow, I don't think it is possible for these lines to ever get hit during normal circumstances. 

First of all, we specify that instance ids passed to `calculate_feature_matirx` should be of type `list` in the docstring. However, even if a single value were passed, this would get converted to a series here: https://github.com/alteryx/featuretools/blob/753615c89360910768811680a7a71e08ab83ad79/featuretools/computational_backends/calculate_feature_matrix.py#L191-L193

If the user does not specify instance ids, we grab a series from the dataframe here: https://github.com/alteryx/featuretools/blob/753615c89360910768811680a7a71e08ab83ad79/featuretools/computational_backends/calculate_feature_matrix.py#L178-L184

So, given those two scenarios I don't think it is possible for our normal code flow to ever call `query_by_values` and consequently `_vals_to_series` with instances that are anything but a `pd.Series`. We do have several tests that call `query_by_values` passing a list of instance values.

Unless I have missed something, I do not believe that the uncovered code is necessary. The documentation says that the user should pass a list for instance ids. Even if a single value is passed, this will get converted to a series by the time `_vals_to_series` gets called. I haven't tested what would happen if a DataFrame is passed, but since we are not telling users that is supported in the documentation, I'm not sure we need to handle that anyway.

With all this in mind we could:
1. Delete the uncovered code 
2. Add a couple new tests that use a single instance value and a DataFrame

I think the first options is better as it removes what appears to be dead code. Thoughts?
rwedge: It's dead code as far as the featuretools code flow is concerned, but those input options are described in the `query_by_values` docstring as valid inputs.  Since `query_by_values` is a public method, it could be getting used.  I know someone on the slack channel ended up using query by values for their purposes.

What if we added some tests now, and made an issue about removing those input options in the future?  See if we get any feedback.
thehomebrewnerd: > It's dead code as far as the featuretools code flow is concerned, but those input options are described in the `query_by_values` docstring as valid inputs. Since `query_by_values` is a public method, it could be getting used. I know someone on the slack channel ended up using query by values for their purposes.
> 
> What if we added some tests now, and made an issue about removing those input options in the future? See if we get any feedback.

The tests should be easy enough to add. As a side note to all this,  `query_by_values` is not currently listed in the API Reference. Do you want it added there as well?
rwedge: Let's leave it out for now

Patch:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 24d61d4a..68036800 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -6,13 +6,19 @@ Release Notes
     * Enhancements
     * Fixes
     * Changes
+        * Move ``query_by_values`` method from ``Entity`` to ``EntitySet`` (:pr:`1251`)
     * Documentation Changes
     * Testing Changes
         * Use repository-scoped token for dependency check (:pr:`1245`:, :pr:`1248`)
         * Fix install error during docs CI test (:pr:`1250`)
 
     Thanks to the following people for contributing to this release:
-    :user:`jeff-hernandez`, :user:`rwedge`
+    :user:`jeff-hernandez`, :user:`rwedge`, :user:`thehomebrewnerd`
+
+**Breaking Changes**
+
+* ``Entity.query_by_values`` has been removed and replaced by ``EntitySet.query_by_values`` with an
+    added ``entity_id`` parameter to specify which entity in the entityset should be used for the query.
 
 **v0.22.0 Nov 30, 2020**
     * Enhancements
diff --git a/featuretools/computational_backends/feature_set_calculator.py b/featuretools/computational_backends/feature_set_calculator.py
index df73a8d4..4e775bdd 100644
--- a/featuretools/computational_backends/feature_set_calculator.py
+++ b/featuretools/computational_backends/feature_set_calculator.py
@@ -209,8 +209,7 @@ class FeatureSetCalculator(object):
         need_full_entity, full_entity_features, not_full_entity_features = feature_trie.value
 
         all_features = full_entity_features | not_full_entity_features
-        entity = self.entityset[entity_id]
-        columns = self._necessary_columns(entity, all_features)
+        columns = self._necessary_columns(entity_id, all_features)
 
         # If we need the full entity then don't filter by filter_values.
         if need_full_entity:
@@ -220,12 +219,13 @@ class FeatureSetCalculator(object):
             query_variable = filter_variable
             query_values = filter_values
 
-        df = entity.query_by_values(query_values,
-                                    variable_id=query_variable,
-                                    columns=columns,
-                                    time_last=self.time_last,
-                                    training_window=self.training_window,
-                                    include_cutoff_time=include_cutoff_time)
+        df = self.entityset.query_by_values(entity_id=entity_id,
+                                            instance_vals=query_values,
+                                            variable_id=query_variable,
+                                            columns=columns,
+                                            time_last=self.time_last,
+                                            training_window=self.training_window,
+                                            include_cutoff_time=include_cutoff_time)
 
         # call to update timer
         progress_callback(0)
@@ -740,9 +740,10 @@ class FeatureSetCalculator(object):
 
         return frame
 
-    def _necessary_columns(self, entity, feature_names):
+    def _necessary_columns(self, entity_id, feature_names):
         # We have to keep all Id columns because we don't know what forward
         # relationships will come from this node.
+        entity = self.entityset[entity_id]
         index_columns = {v.id for v in entity.variables
                          if isinstance(v, (variable_types.Index,
                                            variable_types.Id,
diff --git a/featuretools/entityset/entity.py b/featuretools/entityset/entity.py
index 73723ba9..5df2f314 100644
--- a/featuretools/entityset/entity.py
+++ b/featuretools/entityset/entity.py
@@ -4,7 +4,6 @@ import warnings
 import dask.dataframe as dd
 import numpy as np
 import pandas as pd
-import pandas.api.types as pdtypes
 
 from featuretools import variable_types as vtypes
 from featuretools.utils.entity_utils import (
@@ -215,70 +214,6 @@ class Entity(object):
         new_variable = new_type.create_from(variable)
         self.variables[self.variables.index(variable)] = new_variable
 
-    def query_by_values(self, instance_vals, variable_id=None, columns=None,
-                        time_last=None, training_window=None, include_cutoff_time=True):
-        """Query instances that have variable with given value
-
-        Args:
-            instance_vals (pd.Dataframe, pd.Series, list[str] or str) :
-                Instance(s) to match.
-            variable_id (str) : Variable to query on. If None, query on index.
-            columns (list[str]) : Columns to return. Return all columns if None.
-            time_last (pd.TimeStamp) : Query data up to and including this
-                time. Only applies if entity has a time index.
-            training_window (Timedelta, optional):
-                Window defining how much time before the cutoff time data
-                can be used when calculating features. If None, all data before cutoff time is used.
-            include_cutoff_time (bool):
-                If True, data at cutoff time are included in calculating features
-
-        Returns:
-            pd.DataFrame : instances that match constraints with ids in order of underlying dataframe
-        """
-        if not variable_id:
-            variable_id = self.index
-
-        instance_vals = self._vals_to_series(instance_vals, variable_id)
-
-        training_window = _check_timedelta(training_window)
-
-        if training_window is not None:
-            assert training_window.has_no_observations(), "Training window cannot be in observations"
-
-        if instance_vals is None:
-            df = self.df.copy()
-
-        elif isinstance(instance_vals, pd.Series) and instance_vals.empty:
-            df = self.df.head(0)
-
-        else:
-            if is_instance(instance_vals, (dd, ks), 'Series'):
-                df = self.df.merge(instance_vals.to_frame(), how="inner", on=variable_id)
-            elif isinstance(instance_vals, pd.Series) and is_instance(self.df, ks, 'DataFrame'):
-                df = self.df.merge(ks.DataFrame({variable_id: instance_vals}), how="inner", on=variable_id)
-            else:
-                df = self.df[self.df[variable_id].isin(instance_vals)]
-
-            if isinstance(self.df, pd.DataFrame):
-                df = df.set_index(self.index, drop=False)
-
-            # ensure filtered df has same categories as original
-            # workaround for issue below
-            # github.com/pandas-dev/pandas/issues/22501#issuecomment-415982538
-            if pdtypes.is_categorical_dtype(self.df[variable_id]):
-                categories = pd.api.types.CategoricalDtype(categories=self.df[variable_id].cat.categories)
-                df[variable_id] = df[variable_id].astype(categories)
-
-        df = self._handle_time(df=df,
-                               time_last=time_last,
-                               training_window=training_window,
-                               include_cutoff_time=include_cutoff_time)
-
-        if columns is not None:
-            df = df[columns]
-
-        return df
-
     def _create_variables(self, variable_types, index, time_index, secondary_time_index):
         """Extracts the variables from a dataframe
 
@@ -509,34 +444,6 @@ class Entity(object):
 
         self.secondary_time_index = secondary_time_index
 
-    def _vals_to_series(self, instance_vals, variable_id):
-        """
-        instance_vals may be a pd.Dataframe, a pd.Series, a list, a single
-        value, or None. This function always returns a Series or None.
-        """
-        if instance_vals is None:
-            return None
-
-        # If this is a single value, make it a list
-        if not hasattr(instance_vals, '__iter__'):
-            instance_vals = [instance_vals]
-
-        # convert iterable to pd.Series
-        if isinstance(instance_vals, pd.DataFrame):
-            out_vals = instance_vals[variable_id]
-        elif is_instance(instance_vals, (pd, dd, ks), 'Series'):
-            out_vals = instance_vals.rename(variable_id)
-        else:
-            out_vals = pd.Series(instance_vals)
-
-        # no duplicates or NaN values
-        out_vals = out_vals.drop_duplicates().dropna()
-
-        # want index to have no name for the merge in query_by_values
-        out_vals.index.name = None
-
-        return out_vals
-
     def _handle_time(self, df, time_last=None, training_window=None, include_cutoff_time=True):
         """
         Filter a dataframe for all instances before time_last.
diff --git a/featuretools/entityset/entityset.py b/featuretools/entityset/entityset.py
index 0063c66a..83262b8a 100644
--- a/featuretools/entityset/entityset.py
+++ b/featuretools/entityset/entityset.py
@@ -6,6 +6,7 @@ from collections import defaultdict
 import dask.dataframe as dd
 import numpy as np
 import pandas as pd
+import pandas.api.types as pdtypes
 from pandas.api.types import is_dtype_equal, is_numeric_dtype
 
 import featuretools.variable_types.variable as vtypes
@@ -18,6 +19,7 @@ from featuretools.utils.plot_utils import (
     get_graphviz_format,
     save_graph
 )
+from featuretools.utils.wrangle import _check_timedelta
 
 ks = import_or_none('databricks.koalas')
 
@@ -968,3 +970,98 @@ class EntitySet(object):
         if to_file:
             save_graph(graph, to_file, format_)
         return graph
+
+    def query_by_values(self, entity_id, instance_vals, variable_id=None, columns=None,
+                        time_last=None, training_window=None, include_cutoff_time=True):
+        """Query instances that have variable with given value
+
+        Args:
+            entity_id (str): The id of the entity to query
+            instance_vals (pd.Dataframe, pd.Series, list[str] or str) :
+                Instance(s) to match.
+            variable_id (str) : Variable to query on. If None, query on index.
+            columns (list[str]) : Columns to return. Return all columns if None.
+            time_last (pd.TimeStamp) : Query data up to and including this
+                time. Only applies if entity has a time index.
+            training_window (Timedelta, optional):
+                Window defining how much time before the cutoff time data
+                can be used when calculating features. If None, all data before cutoff time is used.
+            include_cutoff_time (bool):
+                If True, data at cutoff time are included in calculating features
+
+        Returns:
+            pd.DataFrame : instances that match constraints with ids in order of underlying dataframe
+        """
+        entity = self[entity_id]
+        if not variable_id:
+            variable_id = entity.index
+
+        instance_vals = _vals_to_series(instance_vals, variable_id)
+
+        training_window = _check_timedelta(training_window)
+
+        if training_window is not None:
+            assert training_window.has_no_observations(), "Training window cannot be in observations"
+
+        if instance_vals is None:
+            df = entity.df.copy()
+
+        elif isinstance(instance_vals, pd.Series) and instance_vals.empty:
+            df = entity.df.head(0)
+
+        else:
+            if is_instance(instance_vals, (dd, ks), 'Series'):
+                df = entity.df.merge(instance_vals.to_frame(), how="inner", on=variable_id)
+            elif isinstance(instance_vals, pd.Series) and is_instance(entity.df, ks, 'DataFrame'):
+                df = entity.df.merge(ks.DataFrame({variable_id: instance_vals}), how="inner", on=variable_id)
+            else:
+                df = entity.df[entity.df[variable_id].isin(instance_vals)]
+
+            if isinstance(entity.df, pd.DataFrame):
+                df = df.set_index(entity.index, drop=False)
+
+            # ensure filtered df has same categories as original
+            # workaround for issue below
+            # github.com/pandas-dev/pandas/issues/22501#issuecomment-415982538
+            if pdtypes.is_categorical_dtype(entity.df[variable_id]):
+                categories = pd.api.types.CategoricalDtype(categories=entity.df[variable_id].cat.categories)
+                df[variable_id] = df[variable_id].astype(categories)
+
+        df = entity._handle_time(df=df,
+                                 time_last=time_last,
+                                 training_window=training_window,
+                                 include_cutoff_time=include_cutoff_time)
+
+        if columns is not None:
+            df = df[columns]
+
+        return df
+
+
+def _vals_to_series(instance_vals, variable_id):
+    """
+    instance_vals may be a pd.Dataframe, a pd.Series, a list, a single
+    value, or None. This function always returns a Series or None.
+    """
+    if instance_vals is None:
+        return None
+
+    # If this is a single value, make it a list
+    if not hasattr(instance_vals, '__iter__'):
+        instance_vals = [instance_vals]
+
+    # convert iterable to pd.Series
+    if isinstance(instance_vals, pd.DataFrame):
+        out_vals = instance_vals[variable_id]
+    elif is_instance(instance_vals, (pd, dd, ks), 'Series'):
+        out_vals = instance_vals.rename(variable_id)
+    else:
+        out_vals = pd.Series(instance_vals)
+
+    # no duplicates or NaN values
+    out_vals = out_vals.drop_duplicates().dropna()
+
+    # want index to have no name for the merge in query_by_values
+    out_vals.index.name = None
+
+    return out_vals


Test patch:
diff --git a/featuretools/tests/entityset_tests/test_entity.py b/featuretools/tests/entityset_tests/test_entity.py
index 4e2df723..e91350bb 100644
--- a/featuretools/tests/entityset_tests/test_entity.py
+++ b/featuretools/tests/entityset_tests/test_entity.py
@@ -138,33 +138,6 @@ def test_update_data(es):
         assert es['customers'].df["id"].iloc[0] == 0
 
 
-def test_query_by_values_returns_rows_in_given_order():
-    data = pd.DataFrame({
-        "id": [1, 2, 3, 4, 5],
-        "value": ["a", "c", "b", "a", "a"],
-        "time": [1000, 2000, 3000, 4000, 5000]
-    })
-
-    es = ft.EntitySet()
-    es = es.entity_from_dataframe(entity_id="test", dataframe=data, index="id",
-                                  time_index="time", variable_types={
-                                            "value": ft.variable_types.Categorical
-                                  })
-    query = es['test'].query_by_values(['b', 'a'], variable_id='value')
-    assert np.array_equal(query['id'], [1, 3, 4, 5])
-
-
-def test_query_by_values_secondary_time_index(es):
-    end = np.datetime64(datetime(2011, 10, 1))
-    all_instances = [0, 1, 2]
-    result = es['customers'].query_by_values(all_instances, time_last=end)
-    result = to_pandas(result, index='id')
-
-    for col in ["cancel_date", "cancel_reason"]:
-        nulls = result.loc[all_instances][col].isnull() == [False, True, True]
-        assert nulls.all(), "Some instance has data it shouldn't for column %s" % col
-
-
 def test_delete_variables(es):
     entity = es['customers']
     to_delete = ['age', 'cohort', 'email']
diff --git a/featuretools/tests/entityset_tests/test_es.py b/featuretools/tests/entityset_tests/test_es.py
index d453fdf4..6df6ce08 100644
--- a/featuretools/tests/entityset_tests/test_es.py
+++ b/featuretools/tests/entityset_tests/test_es.py
@@ -176,13 +176,53 @@ def test_add_relationship_empty_child_convert_dtype(es):
     assert es['log'].df['session_id'].dtype == 'int64'
 
 
+def test_query_by_values_returns_rows_in_given_order():
+    data = pd.DataFrame({
+        "id": [1, 2, 3, 4, 5],
+        "value": ["a", "c", "b", "a", "a"],
+        "time": [1000, 2000, 3000, 4000, 5000]
+    })
+
+    es = ft.EntitySet()
+    es = es.entity_from_dataframe(entity_id="test", dataframe=data, index="id",
+                                  time_index="time", variable_types={
+                                            "value": ft.variable_types.Categorical
+                                  })
+    query = es.query_by_values('test', ['b', 'a'], variable_id='value')
+    assert np.array_equal(query['id'], [1, 3, 4, 5])
+
+
+def test_query_by_values_secondary_time_index(es):
+    end = np.datetime64(datetime(2011, 10, 1))
+    all_instances = [0, 1, 2]
+    result = es.query_by_values('customers', all_instances, time_last=end)
+    result = to_pandas(result, index='id')
+
+    for col in ["cancel_date", "cancel_reason"]:
+        nulls = result.loc[all_instances][col].isnull() == [False, True, True]
+        assert nulls.all(), "Some instance has data it shouldn't for column %s" % col
+
+
 def test_query_by_id(es):
-    df = to_pandas(es['log'].query_by_values(instance_vals=[0]))
+    df = to_pandas(es.query_by_values('log', instance_vals=[0]))
+    assert df['id'].values[0] == 0
+
+
+def test_query_by_single_value(es):
+    df = to_pandas(es.query_by_values('log', instance_vals=0))
     assert df['id'].values[0] == 0
 
 
+def test_query_by_df(es):
+    instance_df = pd.DataFrame({'id': [1, 3], 'vals': [0, 1]})
+    df = to_pandas(es.query_by_values('log', instance_vals=instance_df))
+
+    assert np.array_equal(df['id'], [1, 3])
+
+
 def test_query_by_id_with_time(es):
-    df = es['log'].query_by_values(
+    df = es.query_by_values(
+        entity_id='log',
         instance_vals=[0, 1, 2, 3, 4],
         time_last=datetime(2011, 4, 9, 10, 30, 2 * 6))
     df = to_pandas(df)
@@ -194,7 +234,8 @@ def test_query_by_id_with_time(es):
 
 
 def test_query_by_variable_with_time(es):
-    df = es['log'].query_by_values(
+    df = es.query_by_values(
+        entity_id='log',
         instance_vals=[0, 1, 2], variable_id='session_id',
         time_last=datetime(2011, 4, 9, 10, 50, 0))
     df = to_pandas(df)
@@ -210,7 +251,8 @@ def test_query_by_variable_with_time(es):
 
 
 def test_query_by_variable_with_training_window(es):
-    df = es['log'].query_by_values(
+    df = es.query_by_values(
+        entity_id='log',
         instance_vals=[0, 1, 2], variable_id='session_id',
         time_last=datetime(2011, 4, 9, 10, 50, 0),
         training_window='15m')
@@ -221,7 +263,8 @@ def test_query_by_variable_with_training_window(es):
 
 
 def test_query_by_indexed_variable(es):
-    df = es['log'].query_by_values(
+    df = es.query_by_values(
+        entity_id='log',
         instance_vals=['taco clock'],
         variable_id='product_id')
     df = to_pandas(df)
@@ -748,7 +791,7 @@ def test_concat_entitysets(es):
     assert not es.__eq__(es_2, deep=True)
 
     # make sure internal indexes work before concat
-    regions = es_1['customers'].query_by_values(['United States'], variable_id=u'région_id')
+    regions = es_1.query_by_values('customers', ['United States'], variable_id=u'région_id')
     assert regions.index.isin(es_1['customers'].df.index).all()
 
     assert es_1.__eq__(es_2)