mirror of
https://github.com/NVIDIA-NeMo/DataDesigner
synced 2026-05-24 09:48:29 +00:00
65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import pandas as pd
|
|
from pydantic import ValidationError
|
|
import pytest
|
|
|
|
from data_designer.config.columns import DataDesignerColumnType, SamplerColumnConfig, SamplerType
|
|
from data_designer.engine.analysis.column_profilers.base import (
|
|
ColumnConfigWithDataFrame,
|
|
ColumnProfilerMetadata,
|
|
)
|
|
|
|
|
|
def test_column_config_with_dataframe_valid_column_config_with_dataframe():
|
|
df = pd.DataFrame({"test_column": [1, 2, 3]})
|
|
column_config = SamplerColumnConfig(name="test_column", sampler_type=SamplerType.CATEGORY, params={})
|
|
|
|
config_with_df = ColumnConfigWithDataFrame(column_config=column_config, df=df)
|
|
|
|
assert config_with_df.column_config.name == "test_column"
|
|
assert "test_column" in config_with_df.df.columns
|
|
assert config_with_df.df["test_column"].tolist() == [1, 2, 3]
|
|
|
|
|
|
def test_column_config_with_dataframe_column_not_found_validation_error():
|
|
df = pd.DataFrame({"other_column": [1, 2, 3]})
|
|
column_config = SamplerColumnConfig(name="test_column", sampler_type=SamplerType.CATEGORY, params={})
|
|
|
|
with pytest.raises(ValidationError, match="Column 'test_column' not found in DataFrame"):
|
|
ColumnConfigWithDataFrame(column_config=column_config, df=df)
|
|
|
|
|
|
def test_column_config_with_dataframe_pyarrow_backend_conversion():
|
|
df = pd.DataFrame({"test_column": [1, 2, 3]})
|
|
column_config = SamplerColumnConfig(name="test_column", sampler_type=SamplerType.CATEGORY, params={})
|
|
|
|
config_with_df = ColumnConfigWithDataFrame(column_config=column_config, df=df)
|
|
|
|
assert all(isinstance(dtype, pd.ArrowDtype) for dtype in config_with_df.df.dtypes)
|
|
|
|
|
|
def test_column_config_with_dataframe_as_tuple_method():
|
|
df = pd.DataFrame({"test_column": [1, 2, 3]})
|
|
column_config = SamplerColumnConfig(name="test_column", sampler_type=SamplerType.CATEGORY, params={})
|
|
|
|
config_with_df = ColumnConfigWithDataFrame(column_config=column_config, df=df)
|
|
column_config_result, df_result = config_with_df.as_tuple()
|
|
|
|
assert column_config_result == column_config
|
|
|
|
assert df_result["test_column"].tolist() == df["test_column"].tolist()
|
|
|
|
|
|
def test_column_profiler_metadata_creation():
|
|
metadata = ColumnProfilerMetadata(
|
|
name="test_profiler",
|
|
description="Test profiler",
|
|
applicable_column_types=[DataDesignerColumnType.SAMPLER, DataDesignerColumnType.LLM_TEXT],
|
|
required_resources=None,
|
|
)
|
|
|
|
assert metadata.name == "test_profiler"
|
|
assert metadata.description == "Test profiler"
|
|
assert metadata.applicable_column_types == [DataDesignerColumnType.SAMPLER, DataDesignerColumnType.LLM_TEXT]
|