[Data] Added support for projection pushdown into Parquet reads #56500

alexeykudinkin · 2025-09-13T02:38:26Z

Why are these changes needed?

Adding support for projection pushdown into Read ops reading from parquet tables.

Changes

Adding ProjectionPushdown optimization rule
Abstracting LogicalOperatorSupportsProjectionPushdown, etc interfaces to implement projection pushdown for various Read ops as well as Datasources
Added tests for projection push-down

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

gemini-code-assist

Code Review

This pull request introduces an important optimization, ProjectionPushdown, which serves two purposes: fusing consecutive Project operators and pushing column projections into ParquetDatasource during reads. The implementation for fusing Project operators is well-covered with extensive tests. However, I've identified a few critical issues. The primary feature of pushing projections into Parquet reads is not tested at all. Additionally, the implementation mutates the datasource object in-place, which is a significant issue for an optimizer rule and can lead to incorrect behavior. I've left detailed comments on these points.

python/ray/data/_internal/logical/rules/projection_pushdown.py

python/ray/data/tests/test_projection_pushdown.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed typo Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin · 2025-09-17T22:42:14Z

python/ray/data/_internal/datasource/parquet_datasource.py

-        # `read_schema` is the schema object that will be used to perform
-        # read operations.
-        # It should be None, unless user has specified the schema or columns.
-        # We don't use the inferred schema for read, because we infer the schema based
-        # on the first file. Thus, files with different schemas will end up producing
-        # blocks with wrong schema.
-        # See https://github.com/ray-project/ray/issues/47960 for more context.
-        read_schema = schema
-        inferred_schema = _infer_schema(
-            pq_ds, schema, columns, partitioning, _block_udf
-        )


This is to defer derivation of the target schema until the actual reading starts to make sure that pushed down projections are picked up appropriately

srinathk10

Minor comments. LGTM otherwise.

srinathk10 · 2025-09-18T17:12:44Z

python/ray/data/_internal/logical/interfaces/logical_operator.py

        return True
+
+
+class LogicalOperatorSupportsProjectionPushdown(LogicalOperator):


nit: May be move these into LogicalOperator, bc projection pushdown can be supported by all LogicalOperator.

Actually, purposefully abstracted it to avoid polluting shared interface.

We actually don't need to do pushdown explicitly in other ops (like Map) simply b/c it will be done by Operator fusion already.

srinathk10 · 2025-09-18T17:15:48Z

python/ray/data/tests/test_parquet.py

    assert ds.count() == partitioned_ds.count()


+def test_projection_pushdown_on_count(ray_start_regular_shared, temp_dir):


nit: Missing project pushdown

It's done internally inside count itself

srinathk10 · 2025-09-18T17:19:56Z

python/ray/data/tests/test_projection_fusion.py

+
+    name: str
+    expressions_list: List[Dict[str, str]]  # List of {name: expression_desc}
+    expected_levels: int


nit: Level terminology is a new. May be Sequence is a better termilogy.

Just cloning tests from RT as is

goutamvenkat-anyscale · 2025-09-18T18:12:01Z

python/ray/data/_internal/logical/interfaces/logical_operator.py

+    def get_current_projection(self) -> Optional[List[str]]:
+        return None
+
+    def apply_projection(self, columns: Optional[List[str]]) -> LogicalOperator:


TODO for later: This should include exclude columns as well

goutamvenkat-anyscale · 2025-09-18T18:14:58Z

python/ray/data/_internal/logical/interfaces/logical_operator.py

+    def get_current_projection(self) -> Optional[List[str]]:
+        return None
+
+    def apply_projection(self, columns: Optional[List[str]]) -> LogicalOperator:


Can you please add a TODO here and tag me to use expressions here?

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Zhiqiang Ma <zhiqiang.ma@intel.com>

## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes Signed-off-by: Alexey Kudinkin <ak@anyscale.com> --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: zac <zac@anyscale.com>

## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Marco Stephan <marco@magic.dev>

## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Douglas Strodtman <douglas@anyscale.com>

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Original PR #56500 by alexeykudinkin Original: ray-project/ray#56500

…nto Parquet reads Merged from original PR #56500 Original: ray-project/ray#56500

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…project#56500)   ## Why are these changes needed? Adding support for projection pushdown into `Read` ops reading from parquet tables. Changes --- - Adding `ProjectionPushdown` optimization rule - Abstracting `LogicalOperatorSupportsProjectionPushdown`, etc interfaces to implement projection pushdown for various Read ops as well as `Datasource`s - Added tests for projection push-down ## Related issue number  ## Checks - [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR. - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've included any doc changes needed for https://docs.ray.io/en/master/. - [ ] I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in `doc/source/tune/api/` under the corresponding `.rst` file. - [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/ - Testing Strategy - [ ] Unit tests - [ ] Release tests - [ ] This PR is not tested :( --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

alexeykudinkin requested a review from a team as a code owner September 13, 2025 02:38

gemini-code-assist bot reviewed Sep 13, 2025

View reviewed changes

python/ray/data/_internal/logical/rules/projection_pushdown.py Outdated Show resolved Hide resolved

python/ray/data/tests/test_projection_pushdown.py Show resolved Hide resolved

ray-gardener bot added the data Ray Data-related issues label Sep 13, 2025

alexeykudinkin added go add ONLY when ready to merge, run all tests and removed data Ray Data-related issues labels Sep 14, 2025

ray-gardener bot added the data Ray Data-related issues label Sep 14, 2025

ray-project deleted a comment from gemini-code-assist bot Sep 17, 2025

alexeykudinkin added 5 commits September 16, 2025 22:51

Scaffolded ProjectionPushdown rule

44693e0

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

5f1c208

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing refs

1cdcd33

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed in-place modification

aa36885

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

f80507c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin force-pushed the ak/prj-pdwn-fix branch from 5638317 to f80507c Compare September 17, 2025 06:08

alexeykudinkin requested a review from goutamvenkat-anyscale September 17, 2025 06:09

alexeykudinkin added 16 commits September 17, 2025 10:23

Added typing;

3c36e27

Fixed typo Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

18af794

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added test for projection pushdown

e82e932

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed logger

ef9a477

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Abstracted projection push-down methods for Datasources

86aff22

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Abstracted projection push-down mixing for Read op

61c2995

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Relocated to interfaces

b52e125

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

f2ae3ca

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Revisited ProjectionPushdown rule to utilize new interfaces

2ba27c4

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Missing imports

7545151

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

8346e69

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed Operator._apply_transform to avoid modifying in place

b9a1c40

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed Read missing to update legacy field

fc32668

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaned up tests

974eb2e

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added tests for projection upon counting

c2d6917

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

9adfeb0

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin added 10 commits September 17, 2025 15:32

Updated tests to assert that empty projection reads 0 bytes

64a1158

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

3c9d875

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added test bazel target

b7fc8ef

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

c1195d6

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed invalid ref

e059500

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Skip overlapping fields

efd08f4

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed tests

93ea4f6

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed memory estimation for empty projections

c54bf6c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed test

e5e93fd

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

fff51d6

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin requested a review from srinathk10 September 18, 2025 01:14

Clean up invalid tests

3af7980

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin commented Sep 18, 2025

View reviewed changes

srinathk10 approved these changes Sep 18, 2025

View reviewed changes

goutamvenkat-anyscale reviewed Sep 18, 2025

View reviewed changes

alexeykudinkin merged commit 41fdf01 into ray-project:master Sep 18, 2025
5 checks passed

snorkelopstesting1-a11y mentioned this pull request Oct 22, 2025

[Data] Added support for projection pushdown into Parquet reads snorkel-marlin-repos/ray-project_ray_pr_56500_b464507f-8c79-4ef9-9a1b-79e67ef6d558#1

Merged

		return True


		class LogicalOperatorSupportsProjectionPushdown(LogicalOperator):

		assert ds.count() == partitioned_ds.count()


		def test_projection_pushdown_on_count(ray_start_regular_shared, temp_dir):

[Data] Added support for projection pushdown into Parquet reads #56500

[Data] Added support for projection pushdown into Parquet reads #56500

Uh oh!

Conversation

alexeykudinkin commented Sep 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Changes

Related issue number

Checks

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

srinathk10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

alexeykudinkin commented Sep 13, 2025 •

edited

Loading