[Data] Add override_num_blocks parameter to `get_parquet_metadata_p…

…rovider` (ray-project#45767) Signed-off-by: Balaji Veeramani <[email protected]>
harborn · Jun 5, 2024 · e233555 · e233555
1 parent 6031b4a
commit e233555
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/ray/data/datasource/_default_metadata_providers.py b/python/ray/data/datasource/_default_metadata_providers.py
@@ -13,7 +13,7 @@ def get_generic_metadata_provider(file_extensions: Optional[List[str]]):
  return DefaultFileMetadataProvider()
 
 
-def get_parquet_metadata_provider():
+def get_parquet_metadata_provider(override_num_blocks: Optional[int] = None):
  # Used by `read_parquet`
  return ParquetMetadataProvider()
 

diff --git a/python/ray/data/read_api.py b/python/ray/data/read_api.py
@@ -769,7 +769,7 @@ def read_parquet(
  files.
  """
  if meta_provider is None:
- meta_provider = get_parquet_metadata_provider()
+ meta_provider = get_parquet_metadata_provider(override_num_blocks)
  arrow_parquet_args = _resolve_parquet_args(
  tensor_column_schema,
  **arrow_parquet_args,