Move file options to the main options class

python/src/space/core/datasets.py

@@ -20,8 +20,7 @@
 import pyarrow as pa
 from substrait.algebra_pb2 import ReadRel, Rel
 
-from space.core.ops.utils import FileOptions
-from space.core.options import JoinOptions, ReadOptions
+from space.core.options import FileOptions, JoinOptions, ReadOptions
 from space.core.runners import LocalRunner
 from space.core.storage import Storage
 from space.core.transform.plans import LogicalPlanBuilder

python/src/space/core/fs/array_record.py

@@ -14,26 +14,11 @@
 #
 """ArrayRecord file utilities."""
 
-from dataclasses import dataclass
 from typing import List, Optional
 
 from space.core.utils.lazy_imports_utils import array_record_module as ar
 
 
-# pylint: disable=line-too-long
-@dataclass
-class ArrayRecordOptions:
-  """Options of ArrayRecord file writer."""
-  # Max uncompressed bytes per file.
-  max_uncompressed_file_bytes = 100 * 1024 * 1024
-  # ArrayRecord lib options.
-  # See https://github.com/google/array_record/blob/2ac1d904f6be31e5aa2f09549774af65d84bff5a/cpp/array_record_writer.h#L83
-  # Group size 1 maximizes random read performance.
-  # Match the options of TFDS:
-  # https://github.com/tensorflow/datasets/blob/92ebd18102b62cf85557ba4b905c970203d8914d/tensorflow_datasets/core/sequential_writer.py#L108
-  options: str = "group_size:1"
-
-
 def read_record_file(file_path: str,
                      positions: Optional[List[int]] = None) -> List[bytes]:
   """Read records of an ArrayRecord file.

python/src/space/core/fs/parquet.py

@@ -14,22 +14,12 @@
 #
 """Parquet file utilities."""
 
-from dataclasses import dataclass
 from typing import List
 
 import pyarrow as pa
 import pyarrow.parquet as pq
 
 
-@dataclass
-class ParquetWriterOptions:
-  """Options of Parquet file writer."""
-  # Max uncompressed bytes per row group.
-  max_uncompressed_row_group_bytes = 100 * 1024
-  # Max uncompressed bytes per file.
-  max_uncompressed_file_bytes: int = 1 * 1024 * 1024
-
-
 def write_parquet_file(file_path: str, schema: pa.Schema,
                        data: List[pa.Table]) -> pq.FileMetaData:
   """Materialize a single Parquet file."""

python/src/space/core/loaders/array_record.py

@@ -24,8 +24,8 @@
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.append import LocalAppendOp
+from space.core.options import FileOptions
 from space.core.schema import arrow
 from space.core.serializers import DictSerializer
 from space.core.utils.paths import StoragePathsMixin

python/src/space/core/ops/append.py

@@ -25,8 +25,8 @@
 from space.core.manifests import IndexManifestWriter
 from space.core.manifests import RecordManifestWriter
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.base import BaseOp, InputData
+from space.core.options import FileOptions
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.core.schema import arrow

python/src/space/core/ops/delete.py

@@ -24,9 +24,9 @@
 from pyroaring import BitMap  # type: ignore[import-not-found]
 
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.append import LocalAppendOp
 from space.core.ops.base import BaseOp
+from space.core.options import FileOptions
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.core.utils import errors

python/src/space/core/ops/insert.py

@@ -27,8 +27,7 @@
 from space.core.ops.base import BaseOp, InputData
 from space.core.ops.delete import FileSetDeleteOp
 from space.core.ops.read import FileSetReadOp
-from space.core.ops.utils import FileOptions
-from space.core.options import ReadOptions
+from space.core.options import FileOptions, ReadOptions
 import space.core.proto.metadata_pb2 as meta
 import space.core.proto.runtime_pb2 as rt
 from space.core.storage import Storage

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

@@ -14,30 +14,18 @@
 #
 """Utilities for operation classes."""
 
-from dataclasses import dataclass, field as dataclass_field
 from typing import List, Optional, Set
 
 import numpy as np
 import pyarrow as pa
 import pyarrow.compute as pc
 
-from space.core.fs.array_record import ArrayRecordOptions
-from space.core.fs.parquet import ParquetWriterOptions
 from space.core.schema import arrow
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.core.utils import errors
 
 
-@dataclass
-class FileOptions:
-  """Options of file IO."""
-  parquet_options: ParquetWriterOptions = dataclass_field(
-      default_factory=ParquetWriterOptions)
-  array_record_options: ArrayRecordOptions = dataclass_field(
-      default_factory=ArrayRecordOptions)
-
-
 def update_index_storage_stats(
     base: meta.StorageStatistics,
     update: meta.StorageStatistics,

python/src/space/core/options.py

@@ -14,7 +14,8 @@
 #
 """Options of Space core lib."""
 
-from dataclasses import dataclass
+from __future__ import annotations
+from dataclasses import dataclass, field as dataclass_field
 from typing import Any, Callable, List, Optional
 
 import pyarrow.compute as pc
@@ -57,16 +58,45 @@ def __post_init__(self):
 
 
 @dataclass
-class Range:
-  """A range of a field."""
-  # Always inclusive.
-  min_: Any
+class FileOptions:
+  """Options of file IO."""
+  # Parquet file options.
+  parquet_options: ParquetWriterOptions = dataclass_field(
+      default_factory=ParquetWriterOptions)
 
-  # Default exclusive.
-  max_: Any
+  # ArrayRecord file options.
+  array_record_options: ArrayRecordOptions = dataclass_field(
+      default_factory=ArrayRecordOptions)
 
-  # Max is inclusive when true.
-  include_max: bool = False
+
+@dataclass
+class ParquetWriterOptions:
+  """Options of Parquet file writer."""
+  # Max uncompressed bytes per row group.
+  max_uncompressed_row_group_bytes: int = 100 * 1024
+
+  # Max uncompressed bytes per file.
+  max_uncompressed_file_bytes: int = 1 * 1024 * 1024
+
+
+# pylint: disable=line-too-long
+@dataclass
+class ArrayRecordOptions:
+  """Options of ArrayRecord file writer."""
+  # Max uncompressed bytes per file.
+  max_uncompressed_file_bytes: int = 100 * 1024 * 1024
+
+  # ArrayRecord lib options.
+  #
+  # See https://github.com/google/array_record/blob/2ac1d904f6be31e5aa2f09549774af65d84bff5a/cpp/array_record_writer.h#L83
+  # Default group size 1 maximizes random read performance.
+  # It matches the options of TFDS:
+  # https://github.com/tensorflow/datasets/blob/92ebd18102b62cf85557ba4b905c970203d8914d/tensorflow_datasets/core/sequential_writer.py#L108
+  #
+  # A larger group size improves read throughput from Cloud Storage, because
+  # each RPC reads a larger chunk of data, which performs better on Cloud
+  # Storage.
+  options: str = "group_size:1"
 
 
 @dataclass
@@ -78,3 +108,16 @@ class JoinOptions:
   # TODO: to support ReadOptions for left and right views, e.g., filter_,
   # snapshot_id
   # TODO: to support join type in PyArrow, only `inner` is supported now.
+
+
+@dataclass
+class Range:
+  """A range of a field."""
+  # Always inclusive.
+  min_: Any
+
+  # Default exclusive.
+  max_: Any
+
+  # Max is inclusive when true.
+  include_max: bool = False

python/src/space/core/runners.py

@@ -27,13 +27,12 @@
 from space.core.loaders.array_record import LocalArrayRecordLoadOp
 from space.core.loaders.parquet import LocalParquetLoadOp
 from space.core.ops.append import LocalAppendOp
-from space.core.ops.append import FileOptions
 from space.core.ops.base import InputData, InputIteratorFn
 from space.core.ops.change_data import ChangeData, read_change_data
 from space.core.ops.delete import FileSetDeleteOp
 from space.core.ops.insert import InsertOptions, LocalInsertOp
 from space.core.ops.read import FileSetReadOp
-from space.core.options import JoinOptions, ReadOptions
+from space.core.options import FileOptions, JoinOptions, ReadOptions
 import space.core.proto.runtime_pb2 as rt
 from space.core.storage import Storage, Version
 from space.core.utils import errors

python/src/space/core/views.py

@@ -22,10 +22,8 @@
 import pyarrow as pa
 from substrait.algebra_pb2 import Rel
 
-from space.core.options import JoinOptions
 from space.core.fs.factory import create_fs
-from space.core.ops.utils import FileOptions
-from space.core.options import ReadOptions
+from space.core.options import FileOptions, JoinOptions, ReadOptions
 import space.core.proto.metadata_pb2 as meta
 from space.core.schema import FieldIdManager
 from space.core.storage import Storage

python/src/space/ray/ops/append.py

@@ -21,9 +21,9 @@
 import ray
 
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.append import BaseAppendOp, LocalAppendOp
 from space.core.ops.base import InputData, InputIteratorFn
+from space.core.options import FileOptions
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.ray.options import RayOptions

python/src/space/ray/ops/delete.py

@@ -21,8 +21,8 @@
 import ray
 
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.delete import BaseDeleteOp, FileSetDeleteOp
+from space.core.options import FileOptions
 from space.core.proto import metadata_pb2 as meta
 from space.core.proto import runtime_pb2 as rt
 from space.core.storage import Storage

python/src/space/ray/ops/insert.py

@@ -23,7 +23,7 @@
 from space.ray.ops.append import RayAppendOp
 from space.core.ops.insert import InsertOptions, LocalInsertOp
 from space.core.ops.insert import filter_matched
-from space.core.ops.utils import FileOptions
+from space.core.options import FileOptions
 import space.core.proto.metadata_pb2 as meta
 import space.core.proto.runtime_pb2 as rt
 from space.core.storage import Storage

python/src/space/ray/runners.py

@@ -29,12 +29,11 @@
 from space.core.runners import BaseReadOnlyRunner, BaseReadWriteRunner
 from space.core.runners import StorageMixin
 from space.core.ops import utils
-from space.core.ops.utils import FileOptions
 from space.core.ops.base import InputData, InputIteratorFn
 from space.core.ops.change_data import ChangeData, ChangeType
 from space.core.ops.delete import FileSetDeleteOp
 from space.core.ops.insert import InsertOptions
-from space.core.options import JoinOptions, ReadOptions
+from space.core.options import FileOptions, JoinOptions, ReadOptions
 import space.core.proto.runtime_pb2 as rt
 from space.core.utils import errors
 from space.ray.ops.append import RayAppendOp

python/tests/core/ops/test_append.py

@@ -17,7 +17,7 @@
 import pyarrow.parquet as pq
 
 from space.core.ops.append import LocalAppendOp
-from space.core.ops.utils import FileOptions
+from space.core.options import FileOptions
 import space.core.proto.metadata_pb2 as meta
 from space.core.storage import Storage
 
@@ -64,8 +64,7 @@ def test_write_pydict_all_types(self, tmp_path, all_types_schema,
     assert patch.storage_statistics_update == meta.StorageStatistics(
         num_rows=5, index_compressed_bytes=114, index_uncompressed_bytes=126)
 
-  def test_write_pydict_with_record_fields(self, tmp_path,
-                                           record_fields_schema,
+  def test_write_pydict_with_record_fields(self, tmp_path, record_fields_schema,
                                            record_fields_input_data):
     location = tmp_path / "dataset"
     storage = Storage.create(location=str(location),

python/tests/core/ops/test_delete.py

@@ -16,9 +16,9 @@
 import pyarrow.compute as pc
 
 from space.core.ops.append import LocalAppendOp
-from space.core.ops.utils import FileOptions
 from space.core.ops.delete import FileSetDeleteOp
 from space.core.ops.read import FileSetReadOp
+from space.core.options import FileOptions
 from space.core.storage import Storage
 
 _default_file_options = FileOptions()

python/tests/core/ops/test_read.py

@@ -16,9 +16,8 @@
 import pyarrow.compute as pc
 
 from space.core.ops.append import LocalAppendOp
-from space.core.ops.utils import FileOptions
 from space.core.ops.read import FileSetReadOp
-from space.core.options import ReadOptions
+from space.core.options import FileOptions, ReadOptions
 from space.core.storage import Storage
 
 _default_file_options = FileOptions()