Skip to content

Commit

Permalink
fix(ingest/delta-lake): skip file count if require_files is false (da…
Browse files Browse the repository at this point in the history
  • Loading branch information
mayurinehate authored Oct 14, 2024
1 parent 3387110 commit b74ba11
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -223,15 +223,14 @@ def ingest_table(
)

customProperties = {
"number_of_files": str(get_file_count(delta_table)),
"partition_columns": str(delta_table.metadata().partition_columns),
"table_creation_time": str(delta_table.metadata().created_time),
"id": str(delta_table.metadata().id),
"version": str(delta_table.version()),
"location": self.source_config.complete_path,
}
if not self.source_config.require_files:
del customProperties["number_of_files"] # always 0
if self.source_config.require_files:
customProperties["number_of_files"] = str(get_file_count(delta_table))

dataset_properties = DatasetPropertiesClass(
description=delta_table.metadata().description,
Expand Down
6 changes: 2 additions & 4 deletions metadata-ingestion/tests/unit/test_mlflow_source.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import datetime
from pathlib import Path
from typing import Any, TypeVar, Union
from typing import Any, Union

import pytest
from mlflow import MlflowClient
Expand All @@ -11,8 +11,6 @@
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.source.mlflow import MLflowConfig, MLflowSource

T = TypeVar("T")


@pytest.fixture
def tracking_uri(tmp_path: Path) -> str:
Expand Down Expand Up @@ -46,7 +44,7 @@ def model_version(
)


def dummy_search_func(page_token: Union[None, str], **kwargs: Any) -> PagedList[T]:
def dummy_search_func(page_token: Union[None, str], **kwargs: Any) -> PagedList[str]:
dummy_pages = dict(
page_1=PagedList(items=["a", "b"], token="page_2"),
page_2=PagedList(items=["c", "d"], token="page_3"),
Expand Down

0 comments on commit b74ba11

Please sign in to comment.