Add writing record fields in Append op.

python/pyproject.toml

@@ -1,34 +1,29 @@
 [project]
 name = "space"
 version = "0.0.1"
-authors = [
-  { name="Space team", email="[email protected]" },
-]
+authors = [{ name = "Space team", email = "[email protected]" }]
 description = "A storage framework for machine learning datasets"
-license = {text = "Apache-2.0"}
+license = { text = "Apache-2.0" }
 classifiers = [
   "License :: OSI Approved :: Apache Software License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3.8",
   "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
-  "Programming Language :: Python :: 3.11"
+  "Programming Language :: Python :: 3.11",
 ]
 requires-python = ">=3.8"
 dependencies = [
+  "array-record",
   "numpy",
   "protobuf",
   "pyarrow >= 14.0.0",
   "tensorflow_datasets",
-  "typing_extensions"
+  "typing_extensions",
 ]
 
 [project.optional-dependencies]
-dev = [
-  "pyarrow-stubs",
-  "tensorflow",
-  "types-protobuf"
-]
+dev = ["pyarrow-stubs", "tensorflow", "types-protobuf"]
 
 [project.urls]
 Homepage = "https://github.com/google/space"
@@ -49,4 +44,9 @@ disable = ['fixme']
 
 [tool.pylint.MAIN]
 ignore = 'space/core/proto'
-ignored-modules = ['space.core.proto', 'google.protobuf', 'substrait']
+ignored-modules = [
+  'space.core.proto',
+  'google.protobuf',
+  'substrait',
+  'array_record',
+]

python/src/space/core/manifests/__init__.py

@@ -12,3 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+"""Manifest files writer and reader implementation."""
+
+from space.core.manifests.index import IndexManifestWriter
+from space.core.manifests.record import RecordManifestWriter

python/src/space/core/manifests/index.py

@@ -19,13 +19,13 @@
 import pyarrow as pa
 import pyarrow.parquet as pq
 
+from space.core.manifests.utils import write_parquet_file
 import space.core.proto.metadata_pb2 as meta
+from space.core.schema import constants
 from space.core.schema.arrow import field_id, field_id_to_column_id_dict
 from space.core.utils import paths
 
 # Manifest file fields.
-_FILE_PATH_FIELD = '_FILE'
-_NUM_ROWS_FIELD = '_NUM_ROWS'
 _INDEX_COMPRESSED_BYTES_FIELD = '_INDEX_COMPRESSED_BYTES'
 _INDEX_UNCOMPRESSED_BYTES_FIELD = '_INDEX_UNCOMPRESSED_BYTES'
 
@@ -57,7 +57,8 @@ def _manifest_schema(
   """Build the index manifest file schema, based on storage schema."""
   primary_keys_ = set(primary_keys)
 
-  fields = [(_FILE_PATH_FIELD, pa.utf8()), (_NUM_ROWS_FIELD, pa.int64()),
+  fields = [(constants.FILE_PATH_FIELD, pa.utf8()),
+            (constants.NUM_ROWS_FIELD, pa.int64()),
             (_INDEX_COMPRESSED_BYTES_FIELD, pa.int64()),
             (_INDEX_UNCOMPRESSED_BYTES_FIELD, pa.int64())]
 
@@ -209,16 +210,6 @@ def finish(self) -> Optional[str]:
     if manifest_data.num_rows == 0:
       return None
 
-    return _write_index_manifest(self._metadata_dir, self._manifest_schema,
-                                 manifest_data)
-
-
-def _write_index_manifest(metadata_dir: str, schema: pa.Schema,
-                          data: pa.Table) -> str:
-  # TODO: currently assume this file is small, so always write a single file.
-  file_path = paths.new_index_manifest_path(metadata_dir)
-  writer = pq.ParquetWriter(file_path, schema)
-  writer.write_table(data)
-
-  writer.close()
-  return file_path
+    file_path = paths.new_index_manifest_path(self._metadata_dir)
+    write_parquet_file(file_path, self._manifest_schema, manifest_data)
+    return file_path

python/src/space/core/manifests/record.py

@@ -0,0 +1,77 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Record manifest files writer and reader implementation."""
+
+from typing import List, Optional
+
+import pyarrow as pa
+
+from space.core.manifests.utils import write_parquet_file
+import space.core.proto.metadata_pb2 as meta
+from space.core.utils import paths
+from space.core.schema import constants
+
+
+def _manifest_schema() -> pa.Schema:
+  fields = [(constants.FILE_PATH_FIELD, pa.utf8()),
+            (constants.FIELD_ID_FIELD, pa.int32()),
+            (constants.NUM_ROWS_FIELD, pa.int64()),
+            (constants.UNCOMPRESSED_BYTES_FIELD, pa.int64())]
+  return pa.schema(fields)  # type: ignore[arg-type]
+
+
+class RecordManifestWriter:
+  """Writer of record manifest files."""
+
+  def __init__(self, metadata_dir: str):
+    self._metadata_dir = metadata_dir
+    self._manifest_schema = _manifest_schema()
+
+    self._file_paths: List[str] = []
+    self._field_ids: List[int] = []
+    self._num_rows: List[int] = []
+    self._uncompressed_bytes: List[int] = []
+
+  def write(self, file_path: str, field_id: int,
+            storage_statistics: meta.StorageStatistics) -> None:
+    """Write a new manifest row.
+
+    Args:
+      file_path: a relative file path of the index file.
+      field_id: the field ID of the associated field for this ArrayRecord file.
+      storage_statistics: storage statistics of the file.
+    """
+    self._file_paths.append(file_path)
+    self._field_ids.append(field_id)
+    self._num_rows.append(storage_statistics.num_rows)
+    self._uncompressed_bytes.append(
+        storage_statistics.record_uncompressed_bytes)
+
+  def finish(self) -> Optional[str]:
+    """Materialize the manifest file and return the file path."""
+    if not self._file_paths:
+      return None
+
+    arrays = [
+        self._file_paths, self._field_ids, self._num_rows,
+        self._uncompressed_bytes
+    ]
+    manifest_data = pa.Table.from_arrays(
+        arrays,  # type: ignore[arg-type]
+        schema=self._manifest_schema)  # type: ignore[call-arg]
+
+    file_path = paths.new_record_manifest_path(self._metadata_dir)
+    write_parquet_file(file_path, self._manifest_schema, manifest_data)
+    return file_path

python/src/space/core/manifests/utils.py

@@ -0,0 +1,28 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Manifest utilities."""
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+
+def write_parquet_file(file_path: str, schema: pa.Schema,
+                       data: pa.Table) -> None:
+  """Materialize a single Parquet file."""
+  # TODO: currently assume this file is small, so always write a single file.
+  writer = pq.ParquetWriter(file_path, schema)
+  writer.write_table(data)
+
+  writer.close()