PYTHON-2922 Add support for custom projections for find_arrow_all (#40)

Author: blink1073

bindings/python/pymongoarrow/api.py

@@ -58,13 +58,13 @@ def find_arrow_all(collection, query, *, schema, **kwargs):
     context = PyMongoArrowContext.from_schema(
         schema, codec_options=collection.codec_options)
 
-    for opt in ('session', 'cursor_type', 'projection'):
+    for opt in ('session', 'cursor_type'):
         if kwargs.pop(opt, None):
             warnings.warn(
                 f'Ignoring option {opt!r} as it is not supported by '
                 'PyMongoArrow', UserWarning, stacklevel=2)
 
-    kwargs['projection'] = schema._get_projection()
+    kwargs.setdefault('projection', schema._get_projection())
     raw_batch_cursor = collection.find_raw_batches(
         query, **kwargs)
     for batch in raw_batch_cursor:

bindings/python/test/test_arrow.py

@@ -72,6 +72,21 @@ def test_find_simple(self):
         self.assertEqual(find_cmd.command['projection'],
                          {'_id': True, 'data': True})
 
+    def test_find_projection(self):
+        expected = Table.from_pydict(
+            {'_id': [4, 3], 'data': [None, 60]},
+            ArrowSchema([('_id', int32()), ('data', int64())]))
+        projection = {'_id': True, 'data': {'$multiply': [2, '$data']}}
+        table = self.run_find({'_id': {'$gt': 2}},
+                              schema=self.schema,
+                              sort=[('_id', DESCENDING)],
+                              projection=projection)
+        self.assertEqual(table, expected)
+
+        find_cmd = self.cmd_listener.results['started'][-1]
+        self.assertEqual(find_cmd.command_name, 'find')
+        self.assertEqual(find_cmd.command['projection'], projection)
+
     def test_find_multiple_batches(self):
         orig_method = self.coll.find_raw_batches
 
@@ -107,13 +122,17 @@ def test_aggregate_simple(self):
         expected = Table.from_pydict(
             {'_id': [1, 2, 3, 4], 'data': [20, 40, 60, None]},
             ArrowSchema([('_id', int32()), ('data', int64())]))
+        projection = {'_id': True, 'data': {'$multiply': [2, '$data']}}
         table = self.run_aggregate(
-            [{'$project': {'_id': True, 'data': {'$multiply': [2, '$data']}}}],
+            [{'$project': projection }],
             schema=self.schema)
         self.assertEqual(table, expected)
 
         agg_cmd = self.cmd_listener.results['started'][-1]
         self.assertEqual(agg_cmd.command_name, 'aggregate')
+        assert len(agg_cmd.command['pipeline']) == 2
+        self.assertEqual(agg_cmd.command['pipeline'][0]['$project'],
+                         projection)
         self.assertEqual(agg_cmd.command['pipeline'][1]['$project'],
                          {'_id': True, 'data': True})
 

bindings/python/test/test_numpy.py

@@ -81,14 +81,16 @@ def test_aggregate_simple(self):
         expected = {
             '_id': np.array([1, 2, 3, 4], dtype=np.int32),
             'data': np.array([20, 40, 60, np.nan], dtype=np.float64)}
-
+        projection = {'_id': True, 'data': {'$multiply': [2, '$data']}}
         actual = aggregate_numpy_all(
-            self.coll, [{'$project': {
-                '_id': True, 'data': {'$multiply': [2, '$data']}}}],
+            self.coll, [{'$project': projection}],
             schema=self.schema)
         self.assert_numpy_equal(actual, expected)
 
         agg_cmd = self.cmd_listener.results['started'][-1]
         self.assertEqual(agg_cmd.command_name, 'aggregate')
+        assert len(agg_cmd.command['pipeline']) == 2
+        self.assertEqual(agg_cmd.command['pipeline'][0]['$project'],
+                         projection)
         self.assertEqual(agg_cmd.command['pipeline'][1]['$project'],
                          {'_id': True, 'data': True})

bindings/python/test/test_pandas.py

@@ -71,13 +71,16 @@ def test_aggregate_simple(self):
         expected = pd.DataFrame(
             data={'_id': [1, 2, 3, 4], 'data': [20, 40, 60, None]}).astype(
             {'_id': 'int32'})
+        projection = {'_id': True, 'data': {'$multiply': [2, '$data']}}
         table = aggregate_pandas_all(
-            self.coll, [{'$project': {
-                '_id': True, 'data': {'$multiply': [2, '$data']}}}],
+            self.coll, [{'$project': projection}],
             schema=self.schema)
         self.assertTrue(table.equals(expected))
 
         agg_cmd = self.cmd_listener.results['started'][-1]
         self.assertEqual(agg_cmd.command_name, 'aggregate')
+        assert len(agg_cmd.command['pipeline']) == 2
+        self.assertEqual(agg_cmd.command['pipeline'][0]['$project'],
+                         projection)
         self.assertEqual(agg_cmd.command['pipeline'][1]['$project'],
                          {'_id': True, 'data': True})