diff --git a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py index 55db094..c25387f 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py +++ b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py @@ -110,9 +110,13 @@ def get_all_existing_ddb_tables(connection: DuckDBPyConnection) -> tuple[str]: def relation_is_empty(relation: DuckDBPyRelation) -> bool: """Check if a duckdb relation is empty""" - if relation.limit(1).shape[0] > 0: - return False - return True + if relation.limit(1).shape[0] == 0: + return True + + cl_all_null_check_exprs = [ + f'bool_and("{cl}" is NULL)' for cl in relation.columns if cl != RECORD_INDEX_COLUMN_NAME + ] + return relation.aggregate(" AND ".join(cl_all_null_check_exprs)).fetchone()[0] # type: ignore def get_duckdb_type_from_annotation(type_annotation: Any) -> DuckDBPyType: @@ -329,7 +333,7 @@ def duckdb_get_entity_count(cls): def _duckdb_check_entity_empty(self, entity: DuckDBPyRelation) -> bool: # pylint: disable=W0613 - return entity.shape[0] == 0 + return relation_is_empty(entity) def duckdb_check_entity_empty(cls): diff --git a/src/dve/core_engine/backends/implementations/duckdb/rules.py b/src/dve/core_engine/backends/implementations/duckdb/rules.py index 2dd6a04..187a880 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/rules.py +++ b/src/dve/core_engine/backends/implementations/duckdb/rules.py @@ -372,18 +372,12 @@ def join_header(self, entities: DuckDBEntities, *, config: HeaderJoin) -> Messag ), ) - joined_rel = ( - source_rel - .cross(target_rel) - .select( - StarExpression( - exclude=[f"{config.target_name}.{c}" for c in target_rel.columns] - ), - FunctionExpression( - "struct_pack", - *[ColumnExpression(f"{config.target_name}.{c}") for c in target_rel.columns], - ).alias(config.header_column_name) - ) + joined_rel = source_rel.cross(target_rel).select( + StarExpression(exclude=[f"{config.target_name}.{c}" for c in target_rel.columns]), + FunctionExpression( + "struct_pack", + *[ColumnExpression(f"{config.target_name}.{c}") for c in target_rel.columns], + ).alias(config.header_column_name), ) entities[config.new_entity_name or config.entity_name] = joined_rel diff --git a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py index 1257306..f3ce03f 100644 --- a/src/dve/core_engine/backends/implementations/spark/spark_helpers.py +++ b/src/dve/core_engine/backends/implementations/spark/spark_helpers.py @@ -359,7 +359,17 @@ def object_to_spark_literal(obj: SparkLiteralType) -> Column: def df_is_empty(dataframe: DataFrame) -> bool: """Return whether a DataFrame is empty.""" - return not dataframe.select(lit(1)).head(1) + + if not dataframe.select(lit(1)).head(1): + return True + + return ( + not dataframe.dropna( + how="all", subset=[cl for cl in dataframe.columns if cl != RECORD_INDEX_COLUMN_NAME] + ) + .select(lit(1)) + .head(1) + ) def _spark_read_parquet(self, path: URI, **kwargs) -> DataFrame: @@ -457,7 +467,7 @@ def spark_get_entity_count(cls): def _spark_check_entity_empty(self, entity: DataFrame) -> bool: # pylint: disable=W0613 - return entity.count() == 0 + return df_is_empty(entity) def spark_check_entity_empty(cls): diff --git a/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_duckdb_helpers.py b/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_duckdb_helpers.py index 0e81cf4..74b7b3a 100644 --- a/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_duckdb_helpers.py +++ b/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_duckdb_helpers.py @@ -25,6 +25,7 @@ relation_is_empty ) from dve.core_engine.backends.utilities import datetime_format_to_regex +from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME from dve.metadata_parser.domain_types import formatteddatetime @pytest.fixture @@ -266,6 +267,18 @@ def test_relation_is_empty(temp_ddb_conn: DuckDBPyConnection): _, con = temp_ddb_conn rel = con.sql("SELECT 'abc' AS test").filter("test IS NULL") assert relation_is_empty(rel) + + rel2 = con.sql("SELECT NULL as fld1, NULL as fld2 UNION ALL SELECT NULL as fld1, NULL as fld2") + assert relation_is_empty(rel2) + + rel3 = con.sql("SELECT 1 as fld1, NULL as fld2 UNION ALL SELECT NULL as fld1, NULL as fld2") + assert not relation_is_empty(rel3) + + rel4 = con.sql(f"SELECT 1 as {RECORD_INDEX_COLUMN_NAME}, NULL as fld1, NULL as fld2 UNION ALL SELECT 2, NULL, NULL") + assert relation_is_empty(rel4) + + rel5 = con.sql(f"SELECT 1 as {RECORD_INDEX_COLUMN_NAME}, NULL as fld1, NULL as fld2 UNION ALL SELECT 2, NULL, 'hello'") + assert not relation_is_empty(rel5) @pytest.mark.parametrize("date_format, expected_regex", [("%Y-%m-%d", r"^[0-9]{4}\-[0-9]{2}\-[0-9]{2}$"), ("%Y-%m-%dT%H:%M:%S", r"^[0-9]{4}\-[0-9]{2}\-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}$"), diff --git a/tests/test_core_engine/test_backends/test_implementations/test_spark/test_spark_helpers.py b/tests/test_core_engine/test_backends/test_implementations/test_spark/test_spark_helpers.py index 41c9e93..34ecb6c 100644 --- a/tests/test_core_engine/test_backends/test_implementations/test_spark/test_spark_helpers.py +++ b/tests/test_core_engine/test_backends/test_implementations/test_spark/test_spark_helpers.py @@ -30,7 +30,9 @@ get_type_from_annotation, object_to_spark_literal, python_to_java_datetime_format, + df_is_empty ) +from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME from dve.metadata_parser.domain_types import conformatteddate from .....fixtures import spark # pylint: disable=unused-import @@ -354,3 +356,31 @@ def test_spark_filter_contract_errors(spark: SparkSession, example_data_contract ("%Y-%m-%dT%H:%M:%S", "yyyy-MM-dd'T'HH:mm:ss")]) def test_python_to_java_datetime_format(python_dt, expected_java_dt): assert python_to_java_datetime_format(python_dt) == expected_java_dt + + +def test_df_is_empty(spark): + df = spark.createDataFrame([], schema=StructType([StructField("col1", StringType()), StructField("col2", LongType())])) + assert df_is_empty(df) + df2 = spark.createDataFrame([(None, None), (None, None), (None, None)], + schema=StructType([StructField("col1", StringType()), StructField("col2", LongType())])) + assert df_is_empty(df2) + df3 = spark.createDataFrame([(None, None), (None, None), (None, 3)], + schema=StructType([StructField("col1", StringType()), StructField("col2", LongType())])) + assert not df_is_empty(df3) + + _schema_w_ri = schema=StructType([ + StructField(RECORD_INDEX_COLUMN_NAME, LongType()), + StructField("col1", StringType()), + StructField("col2", LongType()) + ]) + + df4 = spark.createDataFrame([(1, None, None), (2, None, None), (3, None, None)], + schema=_schema_w_ri) + assert df_is_empty(df4) + + df5 = df4.union(spark.createDataFrame([(4, None, 8)], schema=_schema_w_ri)) + assert not df_is_empty(df5) + + + +