spec: Add fileCount to dataset input & output stats

Signed-off-by: Martynov Maxim <martinov_m_s_@mail.ru>
OpenLineage · Apr 3, 2024 · 35f9b65 · 35f9b65
1 parent 9cb16d0
commit 35f9b65
Show file tree

Hide file tree

Showing 13 changed files with 55 additions and 22 deletions.
diff --git a/client/java/src/test/java/io/openlineage/client/OpenLineageTest.java b/client/java/src/test/java/io/openlineage/client/OpenLineageTest.java
@@ -227,6 +227,7 @@ void factory() throws JsonProcessingException {
                             ol.newDataQualityMetricsInputDatasetFacetBuilder()
                                 .rowCount(10L)
                                 .bytes(20L)
+                                .fileCount(5L)
                                 .columnMetrics(
                                     ol.newDataQualityMetricsInputDatasetFacetColumnMetricsBuilder()
                                         .put(
@@ -258,7 +259,12 @@ void factory() throws JsonProcessingException {
                         .build())
                 .outputFacets(
                     ol.newOutputDatasetOutputFacetsBuilder()
-                        .outputStatistics(ol.newOutputStatisticsOutputDatasetFacet(10L, 20L))
+                        .outputStatistics(
+                            ol.newOutputStatisticsOutputDatasetFacetBuilder()
+                                .rowCount(10L)
+                                .size(20L)
+                                .fileCount(5L)
+                                .build())
                         .build())
                 .build());
 
@@ -296,6 +302,7 @@ void factory() throws JsonProcessingException {
           inputDataset.getInputFacets().getDataQualityMetrics();
       assertEquals((Long) 10L, dq.getRowCount());
       assertEquals((Long) 20L, dq.getBytes());
+      assertEquals((Long) 5L, dq.getFileCount());
       DataQualityMetricsInputDatasetFacetColumnMetricsAdditional colMetrics =
           dq.getColumnMetrics().getAdditionalProperties().get("mycol");
       assertEquals((Double) 10D, colMetrics.getCount());
@@ -315,6 +322,7 @@ void factory() throws JsonProcessingException {
       assertEquals(roundTrip(json), roundTrip(mapper.writeValueAsString(read)));
       assertEquals((Long) 10L, outputDataset.getOutputFacets().getOutputStatistics().getRowCount());
       assertEquals((Long) 20L, outputDataset.getOutputFacets().getOutputStatistics().getSize());
+      assertEquals((Long) 5L, outputDataset.getOutputFacets().getOutputStatistics().getFileCount());
 
       assertEquals(json, mapper.writeValueAsString(read));
     }

diff --git a/client/python/openlineage/client/facet.py b/client/python/openlineage/client/facet.py
@@ -145,10 +145,11 @@ def _get_schema() -> str:
 
 @attr.s
 class OutputStatisticsOutputDatasetFacet(BaseFacet):
-    rowCount: int = attr.ib()  # noqa: N815
+    rowCount: Optional[int] = attr.ib(default=None)  # noqa: N815
     size: Optional[int] = attr.ib(default=None)
+    fileCount: Optional[int] = attr.ib(default=None)  # noqa: N815
 
-    _additional_skip_redact: ClassVar[List[str]] = ["rowCount", "size"]
+    _additional_skip_redact: ClassVar[List[str]] = ["rowCount", "size", "fileCount"]
 
     @staticmethod
     def _get_schema() -> str:
@@ -170,6 +171,7 @@ class ColumnMetric:
 class DataQualityMetricsInputDatasetFacet(BaseFacet):
     rowCount: Optional[int] = attr.ib(default=None)  # noqa: N815
     bytes: Optional[int] = attr.ib(default=None)  # noqa: A003
+    fileCount: Optional[int] = attr.ib(default=None)  # noqa: N815
     columnMetrics: Dict[str, ColumnMetric] = attr.ib(factory=dict)  # noqa: N815
 
     @staticmethod

diff --git a/.../io/openlineage/spark/agent/facets/builder/OutputStatisticsOutputDatasetFacetBuilder.java b/.../io/openlineage/spark/agent/facets/builder/OutputStatisticsOutputDatasetFacetBuilder.java
@@ -39,13 +39,16 @@ protected void build(
           "outputStatistics",
           context
               .getOpenLineage()
-              .newOutputStatisticsOutputDatasetFacet(
+              .newOutputStatisticsOutputDatasetFacetBuilder()
+              .rowCount(
                   Optional.of(metrics.get(JobMetricsHolder.Metric.WRITE_RECORDS))
                       .map(Number::longValue)
-                      .orElse(null),
+                      .orElse(null))
+              .size(
                   Optional.of(metrics.get(JobMetricsHolder.Metric.WRITE_BYTES))
                       .map(Number::longValue)
-                      .orElse(null)));
+                      .orElse(null))
+              .build());
     }
   }
 }
diff --git a/...ration/spark/shared/src/test/java/io/openlineage/spark/agent/OpenLineageRunEventTest.java b/...ration/spark/shared/src/test/java/io/openlineage/spark/agent/OpenLineageRunEventTest.java
@@ -81,6 +81,7 @@ public void testSerializeRunEvent() throws IOException, URISyntaxException {
                         ol.newDataQualityMetricsInputDatasetFacetBuilder()
                             .rowCount(10L)
                             .bytes(20L)
+                            .fileCount(5L)
                             .columnMetrics(
                                 ol.newDataQualityMetricsInputDatasetFacetColumnMetricsBuilder()
                                     .put(
@@ -107,7 +108,12 @@ public void testSerializeRunEvent() throws IOException, URISyntaxException {
                 "output",
                 null,
                 ol.newOutputDatasetOutputFacetsBuilder()
-                    .outputStatistics(ol.newOutputStatisticsOutputDatasetFacet(10L, 20L))
+                    .outputStatistics(
+                        ol.newOutputStatisticsOutputDatasetFacetBuilder()
+                            .rowCount(10L)
+                            .size(20L)
+                            .fileCount(5L)
+                            .build())
                     .build()));
     OpenLineage.RunEvent runStateUpdate =
         ol.newRunEventBuilder()

diff --git a/integration/spark/shared/src/test/resources/test_data/serde/openlineage-event.json b/integration/spark/shared/src/test/resources/test_data/serde/openlineage-event.json
@@ -46,6 +46,7 @@
         "dataQualityMetrics": {
           "rowCount": 10,
           "bytes": 20,
+          "fileCount": 5,
           "columnMetrics": {
             "mycol": {
               "nullCount": 1,
@@ -70,7 +71,8 @@
       "outputFacets": {
         "outputStatistics": {
           "rowCount": 10,
-          "size": 20
+          "size": 20,
+          "fileCount": 5
         }
       }
     }

diff --git a/proxy/fluentd/events/event_full.json b/proxy/fluentd/events/event_full.json
@@ -87,7 +87,8 @@
           "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
           "_schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/DataQualityMetricsInputDatasetFacet",
           "rowCount": 1000,
-          "bytes": 1048576
+          "bytes": 1048576,
+          "fileCount": 5
         },
         "dataQualityAssertions": {
           "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
@@ -155,7 +156,8 @@
           "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
           "_schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/OutputStatisticsOutputDatasetFacet",
           "rowCount": 2000,
-          "size": 2097152
+          "size": 2097152,
+          "fileCount": 5
         }
       },
       "facets": {

diff --git a/proxy/fluentd/events/event_invalid_input_dataset_facet.json b/proxy/fluentd/events/event_invalid_input_dataset_facet.json
@@ -17,12 +17,13 @@
           "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
           "_schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/DataQualityMetricsInputDatasetFacet",
           "noRowCount": 1000,
-          "bytes": 1048576
+          "bytes": 1048576,
+          "fileCount": 5
         }
       }
     }
   ],
-  "outputs": [ ],
+  "outputs": [],
   "producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
   "schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/RunEvent"
 }
diff --git a/proxy/fluentd/events/event_invalid_output_dataset_facet.json b/proxy/fluentd/events/event_invalid_output_dataset_facet.json
@@ -8,7 +8,7 @@
     "namespace": "my-scheduler-namespace",
     "name": "myjob"
   },
-  "inputs": [ ],
+  "inputs": [],
   "outputs": [
     {
       "namespace": "my-datasource-namespace",
@@ -17,8 +17,9 @@
         "outputStatistics": {
           "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
           "_schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/OutputStatisticsOutputDatasetFacet",
-          "notARowCount": 2000,
-          "size": 2097152
+          "rowCount": false,
+          "size": 2097152,
+          "fileCount": 5
         }
       }
     }

diff --git a/proxy/fluentd/test/plugin/test_parser_openlineage.rb b/proxy/fluentd/test/plugin/test_parser_openlineage.rb
@@ -130,7 +130,7 @@ class OpenlineageParserTest < Test::Unit::TestCase
     err = assert_raise Fluent::ParserError do
       @parser.instance.parse(ol_event)
     end
-    assert_match(/Openlineage validation failed: (.+) "rowCount" is a required property/, err.message)
+    assert_match(/Openlineage validation failed: (.+) path \"\/outputs\/0\/outputFacets\/outputStatistics\/rowCount/, err.message)
   end
 
   private

diff --git a/spec/facets/DataQualityMetricsInputDatasetFacet.json b/spec/facets/DataQualityMetricsInputDatasetFacet.json
@@ -1,6 +1,6 @@
 {
   "$schema": "https://json-schema.org/draft/2020-12/schema",
-  "$id": "https://openlineage.io/spec/facets/1-0-0/DataQualityMetricsInputDatasetFacet.json",
+  "$id": "https://openlineage.io/spec/facets/1-0-1/DataQualityMetricsInputDatasetFacet.json",
   "$defs": {
     "DataQualityMetricsInputDatasetFacet": {
       "allOf": [
@@ -19,6 +19,10 @@
               "description": "The size in bytes",
               "type": "integer"
             },
+            "fileCount": {
+              "description": "The number of files evaluated",
+              "type": "integer"
+            },
             "columnMetrics": {
               "description": "The property key is the column name",
               "type": "object",

diff --git a/spec/facets/OutputStatisticsOutputDatasetFacet.json b/spec/facets/OutputStatisticsOutputDatasetFacet.json
@@ -1,6 +1,6 @@
 {
   "$schema": "https://json-schema.org/draft/2020-12/schema",
-  "$id": "https://openlineage.io/spec/facets/1-0-0/OutputStatisticsOutputDatasetFacet.json",
+  "$id": "https://openlineage.io/spec/facets/1-0-1/OutputStatisticsOutputDatasetFacet.json",
   "$defs": {
     "OutputStatisticsOutputDatasetFacet": {
       "allOf": [
@@ -17,9 +17,12 @@
             "size": {
               "description": "The size in bytes written to the dataset",
               "type": "integer"
+            },
+            "fileCount": {
+              "description": "The number of files written to the dataset",
+              "type": "integer"
             }
-          },
-          "required": ["rowCount"]
+          }
         }
       ]
     }

diff --git a/spec/tests/DataQualityMetricsInputDatasetFacet/1.json b/spec/tests/DataQualityMetricsInputDatasetFacet/1.json
@@ -16,7 +16,7 @@
       }
     },
     "rowCount": 10,
-
+    "fileCount": 5,
     "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
     "_schemaURL": "https://openlineage.io/spec/facets/1-0-0/DataQualityMetricsInputDatasetFacet.json"
   }

diff --git a/spec/tests/OutputStatisticsOutputDatasetFacet/1.json b/spec/tests/OutputStatisticsOutputDatasetFacet/1.json
@@ -3,6 +3,7 @@
     "_producer": "https://github.com/OpenLineage/OpenLineage/blob/v1-0-0/client",
     "_schemaURL": "https://openlineage.io/spec/1-0-1/OpenLineage.json#/definitions/OutputStatisticsOutputDatasetFacet",
     "rowCount": 2000,
-    "size": 2097152
+    "size": 2097152,
+    "fileCount": 5
   }
 }